diff --git a/docs/changelog.rst b/docs/changelog.rst index 938e9db..2950dd4 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -10,7 +10,8 @@ Unreleased ---------- - New ``table.checks``, ``table.column_checks`` and ``table.table_checks`` introspection properties expose column-level and table-level ``CHECK`` constraints. (:issue:`834`) -- ``table.transform()`` now preserves ``CHECK`` constraints, including comments within their expressions. Renaming a column rewrites identifier references in them without changing string literals or function names. Dropping a column drops a check owned by that column, and raises ``TransformError`` if a remaining check depends on it. (:issue:`762`) +- ``table.transform()`` now preserves ``CHECK`` constraints, including comments within their expressions. Renaming a column rewrites identifier references in checks without changing string literals or function names. Dropping a column drops a check owned by that column, and raises ``TransformError`` if a remaining check depends on it. (:issue:`762`) +- ``table.transform()`` now preserves comments immediately before or after column definitions. These comments move with the column if it is renamed or reordered, and are removed if the column is dropped. (:issue:`762`) - ``table.transform()`` now works for tables that are referenced by views. Previously the ``ALTER TABLE ... RENAME TO`` step raised ``no such table`` if a view referenced the table being transformed. View definitions are left unchanged - see :ref:`python_api_transform_views`. This also fixes a bug where ``transform(keep_table=...)`` silently rewrote dependent views to point at the frozen backup table instead of the live one. (:issue:`831`) .. _v3_39_1: diff --git a/docs/python-api.rst b/docs/python-api.rst index 75900b6..90c5165 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -1995,6 +1995,8 @@ CHECK constraints A column-level check is removed if its owning column is dropped. Dropping a column referenced by any remaining check raises ``TransformError`` instead of creating an invalid or unexpectedly weakened schema. +Comments immediately before or after a column definition are preserved too. They move with that column if it is renamed or reordered, and are removed if the column is dropped. A comment between two column definitions is treated as belonging to the following column. + .. _python_api_transform_views: Tables referenced by views diff --git a/sqlite_utils/create_table_parser.py b/sqlite_utils/create_table_parser.py index 082c2c8..2d891ae 100644 --- a/sqlite_utils/create_table_parser.py +++ b/sqlite_utils/create_table_parser.py @@ -26,6 +26,12 @@ class Check: end: int = field(default=-1, compare=False, repr=False) +@dataclass(frozen=True) +class ColumnComments: + before: str = "" + after: str = "" + + class ParseError(ValueError): pass @@ -472,8 +478,7 @@ def _column_checks( return checks -def parse_checks(create_sql: str) -> list[Check]: - """Return CHECK constraints from a valid SQLite CREATE TABLE statement.""" +def _table_body(create_sql: str) -> tuple[str, int] | None: all_tokens = _lex(create_sql) tokens = _meaningful(all_tokens) if not tokens or not tokens[0].is_keyword("CREATE"): @@ -484,7 +489,7 @@ def parse_checks(create_sql: str) -> list[Check]: ): index += 1 if index < len(tokens) and tokens[index].is_keyword("VIRTUAL"): - return [] + return None if index >= len(tokens) or not tokens[index].is_keyword("TABLE"): raise ParseError("Expected CREATE TABLE") index += 1 @@ -501,7 +506,7 @@ def parse_checks(create_sql: str) -> list[Check]: if index + 1 < len(tokens) and tokens[index].text == ".": index += 2 if index < len(tokens) and tokens[index].is_keyword("AS"): - return [] + return None if index >= len(tokens) or tokens[index].text != "(": raise ParseError("CREATE TABLE is missing its column list") close = _matching_paren(tokens, index) @@ -512,7 +517,15 @@ def parse_checks(create_sql: str) -> list[Check]: body_start = tokens[index].end body_end = tokens[close].start - body = create_sql[body_start:body_end] + return create_sql[body_start:body_end], body_start + + +def parse_checks(create_sql: str) -> list[Check]: + """Return CHECK constraints from a valid SQLite CREATE TABLE statement.""" + body_info = _table_body(create_sql) + if body_info is None: + return [] + body, body_start = body_info body_tokens = _lex(body) checks: list[Check] = [] for item, item_start, _ in _split_spans(body, body_tokens): @@ -551,6 +564,35 @@ def parse_checks(create_sql: str) -> list[Check]: return checks +def parse_column_comments(create_sql: str) -> dict[str, ColumnComments]: + """Return comments immediately before and after each column definition.""" + body_info = _table_body(create_sql) + if body_info is None: + return {} + body, _ = body_info + comments: dict[str, ColumnComments] = {} + for item, _, _ in _split_spans(body, _lex(body)): + item_tokens = _meaningful(_lex(item)) + if not item_tokens: + continue + item_index = 0 + if item_tokens[item_index].is_keyword("CONSTRAINT"): + item_index = 2 + head = item_tokens[item_index] if item_index < len(item_tokens) else None + if ( + head + and head.kind == "word" + and head.text.upper() in _TABLE_CONSTRAINT_KEYWORDS + ): + continue + column = _unquote(item_tokens[0].text) + before = item[: item_tokens[0].start].strip() + after = item[item_tokens[-1].end :].strip() + if before or after: + comments[column] = ColumnComments(before=before, after=after) + return comments + + def _is_identifier_token(tokens: list[_Token], index: int) -> bool: token = tokens[index] if index + 1 < len(tokens) and tokens[index + 1].text in ("(", "."): @@ -576,9 +618,9 @@ def check_references_identifier(expression: str, identifier: str) -> bool: ) -def check_expression_ends_in_line_comment(expression: str) -> bool: +def sql_ends_in_line_comment(sql: str) -> bool: """Return True if appended SQL would be swallowed by a ``--`` comment.""" - tokens = _lex(expression) + tokens = _lex(sql) if not tokens: return False final = tokens[-1] diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 949f574..2e9b570 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -29,11 +29,13 @@ from sqlite_utils.plugins import ensure_plugins_loaded, pm from .create_table_parser import ( Check, + ColumnComments, ParseError, - check_expression_ends_in_line_comment, check_references_identifier, parse_checks, + parse_column_comments, rewrite_check_expression, + sql_ends_in_line_comment, ) from .utils import ( OperationalError, @@ -95,10 +97,26 @@ def quote_identifier(identifier: str) -> str: def _check_constraint_sql(check: Check) -> str: prefix = f"CONSTRAINT {quote_identifier(check.name)} " if check.name else "" - newline = "\n" if check_expression_ends_in_line_comment(check.check) else "" + newline = "\n" if sql_ends_in_line_comment(check.check) else "" return f"{prefix}CHECK ({check.check}{newline})" +def _column_definition_with_comments( + definition: str, comments: ColumnComments | None +) -> str: + if comments is None: + return definition + before = textwrap.dedent(comments.before).strip() + after = textwrap.dedent(comments.after).strip() + if before: + definition = f"{textwrap.indent(before, ' ')}\n{definition}" + if after: + definition = f"{definition} {after}" + if sql_ends_in_line_comment(after): + definition += "\n" + return definition + + _IDENTIFIER_CASEFOLD = str.maketrans( "ABCDEFGHIJKLMNOPQRSTUVWXYZ", "abcdefghijklmnopqrstuvwxyz" ) @@ -1394,6 +1412,7 @@ class Database: if_not_exists: bool = False, strict: bool = False, _checks: Iterable[Check] | None = None, + _column_comments: Mapping[str, ColumnComments] | None = None, ) -> str: """ Returns the SQL ``CREATE TABLE`` statement for creating the specified table. @@ -1439,6 +1458,10 @@ class Database: defaults = {resolve_casing(n, columns): v for n, v in (defaults or {}).items()} if column_order is not None: column_order = [resolve_casing(c, columns) for c in column_order] + column_comments = { + resolve_casing(name, columns): comments + for name, comments in (_column_comments or {}).items() + } checks = list(_checks or ()) checks_by_column: dict[str, list[Check]] = {} table_checks: list[Check] = [] @@ -1517,13 +1540,16 @@ class Database: # Refs https://github.com/simonw/sqlite-utils/issues/644 if strict and column_type_str == "FLOAT": column_type_str = "REAL" + column_definition = " {} {column_type}{column_extras}".format( + quote_identifier(column_name), + column_type=column_type_str, + column_extras=( + (" " + " ".join(column_extras)) if column_extras else "" + ), + ) column_defs.append( - " {} {column_type}{column_extras}".format( - quote_identifier(column_name), - column_type=column_type_str, - column_extras=( - (" " + " ".join(column_extras)) if column_extras else "" - ), + _column_definition_with_comments( + column_definition, column_comments.get(column_name) ) ) extra_pk = "" @@ -2713,9 +2739,10 @@ class Table(Queryable): try: existing_checks = self.checks + existing_column_comments = parse_column_comments(self.schema) except ParseError as ex: raise TransformError( - f"Could not parse CHECK constraints for table {self.name!r}: {ex}" + f"Could not parse table schema for table {self.name!r}: {ex}" ) from ex create_table_checks: list[Check] = [] for check in existing_checks: @@ -2739,6 +2766,12 @@ class Table(Queryable): ) ) + create_table_column_comments: dict[str, ColumnComments] = {} + for column, comments in existing_column_comments.items(): + owner = resolve_casing(column, existing_columns) + if owner not in drop: + create_table_column_comments[rename.get(owner) or owner] = comments + create_table_foreign_keys: list[ForeignKeyIndicator] = [] if foreign_keys is not None: @@ -2889,6 +2922,7 @@ class Table(Queryable): column_order=column_order, strict=self.strict if strict is None else strict, _checks=create_table_checks, + _column_comments=create_table_column_comments, ).strip() ) diff --git a/tests/test_create_table_parser.py b/tests/test_create_table_parser.py index 3b878f3..a7aa0c0 100644 --- a/tests/test_create_table_parser.py +++ b/tests/test_create_table_parser.py @@ -4,7 +4,13 @@ import hypothesis.strategies as st import pytest from hypothesis import given -from sqlite_utils.create_table_parser import Check, ParseError, parse_checks +from sqlite_utils.create_table_parser import ( + Check, + ColumnComments, + ParseError, + parse_checks, + parse_column_comments, +) def test_parse_column_and_table_checks(): @@ -46,6 +52,26 @@ def test_comments_are_trivia_not_constraints(): ] +def test_parse_comments_owned_by_columns(): + sql = """ + CREATE TABLE t ( + -- Before id + id /* Between name and type */ INTEGER /* After id */, + /* Between column definitions */ + value TEXT CHECK(value != '') /* After value */, + /* Before a table constraint, not a column */ + CHECK(value != 'forbidden') + ) + """ + assert parse_column_comments(sql) == { + "id": ColumnComments(before="-- Before id", after="/* After id */"), + "value": ColumnComments( + before="/* Between column definitions */", + after="/* After value */", + ), + } + + @pytest.mark.parametrize( "expression,expected", [ diff --git a/tests/test_transform.py b/tests/test_transform.py index 44fec99..980ee9d 100644 --- a/tests/test_transform.py +++ b/tests/test_transform.py @@ -1102,6 +1102,54 @@ def test_transform_preserves_check_ending_in_line_comment(fresh_db): inventory.insert({"quantity": -1}) +def test_transform_preserves_comments_owned_by_columns(fresh_db): + fresh_db.execute(""" + CREATE TABLE people ( + -- Primary identifier + id INTEGER PRIMARY KEY /* IDs are stable */, + /* Displayed to users */ + name TEXT /* May contain spaces */, + -- Age in years + age INTEGER -- May be NULL + ) + """) + people = fresh_db["people"] + people.insert({"id": 1, "name": "Cleo", "age": 5}) + people.transform( + rename={"name": "display_name"}, + types={"age": float}, + column_order=("age", "id", "name"), + ) + assert people.get(1) == {"age": 5.0, "id": 1, "display_name": "Cleo"} + schema = people.schema + assert schema.index("-- Age in years") < schema.index('"age" REAL') + assert schema.index('"age" REAL') < schema.index("-- May be NULL") + assert schema.index("-- Primary identifier") < schema.index('"id" INTEGER') + assert schema.index('"id" INTEGER') < schema.index("/* IDs are stable */") + assert schema.index("/* Displayed to users */") < schema.index( + '"display_name" TEXT' + ) + assert schema.index('"display_name" TEXT') < schema.index( + "/* May contain spaces */" + ) + + +def test_transform_drops_comments_owned_by_dropped_column(fresh_db): + fresh_db.execute(""" + CREATE TABLE t ( + /* Keep this explanation */ + id INTEGER, + /* Drop this explanation */ + obsolete TEXT /* Drop this too */ + ) + """) + fresh_db["t"].transform(drop={"obsolete"}) + schema = fresh_db["t"].schema + assert "Keep this explanation" in schema + assert "Drop this explanation" not in schema + assert "Drop this too" not in schema + + def test_transform_renames_columns_inside_check_constraints(fresh_db): fresh_db.execute(""" CREATE TABLE inventory (