.transform() preserves check constraints, refs #762

sqlite_utils/create_table_parser.py

@@ -549,3 +549,86 @@ def parse_checks(create_sql: str) -> list[Check]:
             _column_checks(item, item_tokens, column, body_start + item_start)
         )
     return checks
+
+
+def _is_identifier_token(tokens: list[_Token], index: int) -> bool:
+    token = tokens[index]
+    if index + 1 < len(tokens) and tokens[index + 1].text in ("(", "."):
+        return False
+    if index and (
+        tokens[index - 1].is_keyword("COLLATE") or tokens[index - 1].is_keyword("AS")
+    ):
+        return False
+    if token.kind == "identifier":
+        return True
+    if token.kind != "word" or token.text.upper() in _SQLITE_KEYWORDS:
+        return False
+    return True
+
+
+def check_references_identifier(expression: str, identifier: str) -> bool:
+    tokens = _meaningful(_lex(expression))
+    folded = _ascii_fold(identifier)
+    return any(
+        _is_identifier_token(tokens, index)
+        and _ascii_fold(_unquote(token.text)) == folded
+        for index, token in enumerate(tokens)
+    )
+
+
+def check_expression_ends_in_line_comment(expression: str) -> bool:
+    """Return True if appended SQL would be swallowed by a ``--`` comment."""
+    tokens = _lex(expression)
+    if not tokens:
+        return False
+    final = tokens[-1]
+    return (
+        final.kind == "comment"
+        and final.text.startswith("--")
+        and not final.text.endswith(("\n", "\r"))
+    )
+
+
+def _valid_bare_identifier(identifier: str) -> bool:
+    if not identifier or identifier.upper() in _SQLITE_KEYWORDS:
+        return False
+    first = identifier[0]
+    if not (first == "_" or first.isalpha() or ord(first) >= 0x80):
+        return False
+    return all(
+        char == "_" or char == "$" or char.isalnum() or ord(char) >= 0x80
+        for char in identifier[1:]
+    )
+
+
+def _quote_replacement(original: str, replacement: str) -> str:
+    if original.startswith('"'):
+        return '"{}"'.format(replacement.replace('"', '""'))
+    if original.startswith("`"):
+        return "`{}`".format(replacement.replace("`", "``"))
+    if original.startswith("[") and "]" not in replacement:
+        return f"[{replacement}]"
+    if _valid_bare_identifier(replacement):
+        return replacement
+    return '"{}"'.format(replacement.replace('"', '""'))
+
+
+def rewrite_check_expression(expression: str, rename: dict[str, str]) -> str:
+    """Rewrite column identifiers in a CHECK expression, preserving trivia."""
+    if not rename:
+        return expression
+    tokens = _lex(expression)
+    meaningful = _meaningful(tokens)
+    replacements = {_ascii_fold(key): value for key, value in rename.items()}
+    edits: list[tuple[int, int, str]] = []
+    for index, token in enumerate(meaningful):
+        if not _is_identifier_token(meaningful, index):
+            continue
+        replacement = replacements.get(_ascii_fold(_unquote(token.text)))
+        if replacement is not None:
+            edits.append(
+                (token.start, token.end, _quote_replacement(token.text, replacement))
+            )
+    for start, end, replacement in reversed(edits):
+        expression = expression[:start] + replacement + expression[end:]
+    return expression

sqlite_utils/db.py

@@ -27,7 +27,14 @@ from typing_extensions import Self
 
 from sqlite_utils.plugins import ensure_plugins_loaded, pm
 
-from .create_table_parser import Check, parse_checks
+from .create_table_parser import (
+    Check,
+    ParseError,
+    check_expression_ends_in_line_comment,
+    check_references_identifier,
+    parse_checks,
+    rewrite_check_expression,
+)
 from .utils import (
     OperationalError,
     chunks,
@@ -86,6 +93,12 @@ def quote_identifier(identifier: str) -> str:
     return '"{}"'.format(identifier.replace('"', '""'))
 
 
+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 ""
+    return f"{prefix}CHECK ({check.check}{newline})"
+
+
 _IDENTIFIER_CASEFOLD = str.maketrans(
     "ABCDEFGHIJKLMNOPQRSTUVWXYZ", "abcdefghijklmnopqrstuvwxyz"
 )
@@ -1380,6 +1393,7 @@ class Database:
         extracts: dict[str, str] | list[str] | None = None,
         if_not_exists: bool = False,
         strict: bool = False,
+        _checks: Iterable[Check] | None = None,
     ) -> str:
         """
         Returns the SQL ``CREATE TABLE`` statement for creating the specified table.
@@ -1425,6 +1439,19 @@ 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]
+        checks = list(_checks or ())
+        checks_by_column: dict[str, list[Check]] = {}
+        table_checks: list[Check] = []
+        for check in checks:
+            if check.column:
+                column = resolve_casing(check.column, columns)
+                if column not in columns:
+                    raise AlterError(
+                        f"No such column for CHECK constraint: {check.column}"
+                    )
+                checks_by_column.setdefault(column, []).append(check)
+            else:
+                table_checks.append(check)
         if not columns:
             raise ValueError("Tables must have at least one column")
         if not all(n in columns for n in not_null):
@@ -1481,6 +1508,10 @@ class Database:
                 column_extras.append(
                     f"REFERENCES {quote_identifier(fk.other_table)}({quote_identifier(cast(str, fk.other_column))}){_fk_actions_sql(fk)}"
                 )
+            column_extras.extend(
+                _check_constraint_sql(check)
+                for check in checks_by_column.get(column_name, ())
+            )
             column_type_str = COLUMN_TYPE_MAPPING[column_type]
             # Special case for strict tables to map FLOAT to REAL
             # Refs https://github.com/simonw/sqlite-utils/issues/644
@@ -1520,6 +1551,9 @@ class Database:
                     actions=_fk_actions_sql(fk),
                 )
             )
+        column_defs.extend(
+            f"   {_check_constraint_sql(check)}" for check in table_checks
+        )
         columns_sql = ",\n".join(column_defs)
         sql = """CREATE TABLE {if_not_exists}{table} (
 {columns_sql}{extra_pk}
@@ -2677,6 +2711,34 @@ class Table(Queryable):
         if column_order is not None:
             column_order = [resolve_casing(c, existing_columns) for c in column_order]
 
+        try:
+            existing_checks = self.checks
+        except ParseError as ex:
+            raise TransformError(
+                f"Could not parse CHECK constraints for table {self.name!r}: {ex}"
+            ) from ex
+        create_table_checks: list[Check] = []
+        for check in existing_checks:
+            owner = (
+                resolve_casing(check.column, existing_columns) if check.column else ""
+            )
+            # A column-level constraint disappears with the column that owns it.
+            if owner and owner in drop:
+                continue
+            for dropped_column in drop:
+                if check_references_identifier(check.check, dropped_column):
+                    raise TransformError(
+                        f"Cannot drop column {dropped_column!r}: it is used by "
+                        f"CHECK constraint {check.name or check.check!r}"
+                    )
+            create_table_checks.append(
+                Check(
+                    rewrite_check_expression(check.check, rename),
+                    name=check.name,
+                    column=rename.get(owner) or owner,
+                )
+            )
+
         create_table_foreign_keys: list[ForeignKeyIndicator] = []
 
         if foreign_keys is not None:
@@ -2826,6 +2888,7 @@ class Table(Queryable):
                 foreign_keys=create_table_foreign_keys,
                 column_order=column_order,
                 strict=self.strict if strict is None else strict,
+                _checks=create_table_checks,
             ).strip()
         )