remote/sqlite-utils
1
0
Fork 0
mirror of https://github.com/simonw/sqlite-utils.git synced 2026-08-13 04:44:19 +02:00
Code Issues Projects Releases Packages Wiki Activity

Preserve column before/after comments through .transform()

Browse source 
Refs #762

The before comment comes before the column definition - the after
comment is anything after it but before its trailing comma.
This commit is contained in:
Simon Willison 2026-08-11 22:37:13 -07:00
commit b37b8cf8c8
6 changed files with 171 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

56
sqlite_utils/create_table_parser.py
View file

@ -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]

52
sqlite_utils/db.py
View file

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