Preserve column before/after comments through .transform()

Refs #762

The before comment comes before the column definition - the after
comment is anything after it but before its trailing comma.

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:

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

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]

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()
         )
 

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",
     [

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 (