From 75ba58846206b2c1beb39836134e763bf34177aa Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Wed, 12 Aug 2026 18:46:18 -0700 Subject: [PATCH] Preserve composite UNIQUE constraints in transforms --- docs/changelog.rst | 1 + sqlite_utils/create_table_parser.py | 193 ++++++++++++++++++++++++++++ sqlite_utils/db.py | 119 +++++++++++++++++ tests/test_create_table_parser.py | 48 +++++++ tests/test_transform.py | 74 +++++++++-- 5 files changed, 425 insertions(+), 10 deletions(-) diff --git a/docs/changelog.rst b/docs/changelog.rst index 624808f..fe215fd 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -9,6 +9,7 @@ Unreleased ---------- +- ``table.transform()`` now preserves column-level and composite ``UNIQUE`` constraints, including constraint names, collations, sort order and ``ON CONFLICT`` behavior. Renaming columns updates those constraints, while dropping any constituent column removes the entire constraint. (:issue:`762`) - ``table.transform()`` now preserves ``AUTOINCREMENT`` primary keys and their sequence high-water marks. Previously a transform removed ``AUTOINCREMENT`` and could reuse deleted row IDs. (:issue:`602`) - New ``table.checks``, ``table.column_checks`` and ``table.table_checks`` introspection properties expose column-level and table-level ``CHECK`` constraints. (:issue:`834`) - New ``sqlite_utils.ANY`` marker type for creating and introspecting SQLite ``ANY`` columns. The Python API and CLI can create, add and transform these columns, and ``table.transform()`` and ``table.extract()`` now preserve ``ANY`` columns and their values in ``STRICT`` tables. (:issue:`790`) diff --git a/sqlite_utils/create_table_parser.py b/sqlite_utils/create_table_parser.py index d426286..7377f4f 100644 --- a/sqlite_utils/create_table_parser.py +++ b/sqlite_utils/create_table_parser.py @@ -32,6 +32,24 @@ class ColumnComments: after: str = "" +@dataclass(frozen=True) +class UniqueColumn: + name: str + collation: str = "" + order: str = "" + + +@dataclass +class Unique: + columns: tuple[UniqueColumn, ...] + name: str = "" + column: str = "" + conflict: str = "" + sql: str = field(default="", compare=False, repr=False) + start: int = field(default=-1, compare=False, repr=False) + end: int = field(default=-1, compare=False, repr=False) + + class ParseError(ValueError): pass @@ -592,6 +610,181 @@ def parse_autoincrement(create_sql: str) -> str | None: return None +_CONFLICT_ACTIONS = frozenset(("ROLLBACK", "ABORT", "FAIL", "IGNORE", "REPLACE")) + + +def _conflict_after(tokens: list[_Token], index: int) -> tuple[str, int]: + if index >= len(tokens) or not tokens[index].is_keyword("ON"): + return "", index + if index + 2 >= len(tokens) or not tokens[index + 1].is_keyword("CONFLICT"): + raise ParseError("ON after UNIQUE must be followed by CONFLICT and an action") + action = tokens[index + 2].text.upper() + if tokens[index + 2].kind != "word" or action not in _CONFLICT_ACTIONS: + raise ParseError("Invalid UNIQUE ON CONFLICT action") + return action, index + 3 + + +def _unique_columns( + item: str, tokens: list[_Token], open_index: int +) -> tuple[tuple[UniqueColumn, ...], int]: + close = _matching_paren(tokens, open_index) + inner = item[tokens[open_index].end : tokens[close].start] + columns: list[UniqueColumn] = [] + for raw_column in _split_ranges(inner, _lex(inner)): + column_tokens = _meaningful(_lex(raw_column)) + if not column_tokens or column_tokens[0].kind not in ( + "word", + "identifier", + "string", + ): + raise ParseError("UNIQUE constraint has an invalid column") + name = _unquote(column_tokens[0].text) + collation = "" + order = "" + index = 1 + if index < len(column_tokens) and column_tokens[index].is_keyword("COLLATE"): + if index + 1 >= len(column_tokens): + raise ParseError("COLLATE in UNIQUE constraint is missing its name") + collation = _unquote(column_tokens[index + 1].text) + index += 2 + if index < len(column_tokens) and ( + column_tokens[index].is_keyword("ASC") + or column_tokens[index].is_keyword("DESC") + ): + order = column_tokens[index].text.upper() + index += 1 + if index != len(column_tokens): + raise ParseError("UNIQUE constraint has an invalid indexed column") + columns.append(UniqueColumn(name, collation=collation, order=order)) + if not columns: + raise ParseError("UNIQUE constraint must include at least one column") + return tuple(columns), close + 1 + + +def _column_uniques( + item: str, tokens: list[_Token], column: str, base_offset: int +) -> list[Unique]: + uniques: list[Unique] = [] + collation = "" + collation_index = 1 + while collation_index < len(tokens): + token = tokens[collation_index] + if token.text == "(": + collation_index = _matching_paren(tokens, collation_index) + 1 + continue + if token.is_keyword("COLLATE"): + if collation_index + 1 >= len(tokens): + raise ParseError("COLLATE is missing its name") + collation = _unquote(tokens[collation_index + 1].text) + collation_index += 2 + continue + collation_index += 1 + pending_name = "" + pending_start: int | None = None + index = 1 + while index < len(tokens): + token = tokens[index] + if token.text == "(": + index = _matching_paren(tokens, index) + 1 + continue + if token.is_keyword("CONSTRAINT"): + if index + 1 >= len(tokens): + raise ParseError("CONSTRAINT is missing its name") + pending_name = _unquote(tokens[index + 1].text) + pending_start = index + index += 2 + continue + if token.is_keyword("UNIQUE"): + source_start = tokens[ + pending_start if pending_start is not None else index + ].start + conflict, next_index = _conflict_after(tokens, index + 1) + source_end = tokens[next_index - 1].end + uniques.append( + Unique( + (UniqueColumn(column, collation=collation),), + name=pending_name, + column=column, + conflict=conflict, + sql=item[source_start:source_end], + start=base_offset + source_start, + end=base_offset + source_end, + ) + ) + pending_name = "" + pending_start = None + index = next_index + continue + if ( + token.kind == "word" + and token.text.upper() in _OTHER_COLUMN_CONSTRAINT_KEYWORDS + ): + pending_name = "" + pending_start = None + index += 1 + return uniques + + +def parse_uniques(create_sql: str) -> list[Unique]: + """Return column-level and table-level UNIQUE constraints.""" + body_info = _table_body(create_sql) + if body_info is None: + return [] + body, body_start = body_info + uniques: list[Unique] = [] + for item, item_start, _ in _split_spans(body, _lex(body)): + item_tokens = _meaningful(_lex(item)) + if not item_tokens: + continue + item_index = 0 + constraint_name = "" + if item_tokens[item_index].is_keyword("CONSTRAINT"): + if len(item_tokens) < 2: + raise ParseError("CONSTRAINT is missing its name") + constraint_name = _unquote(item_tokens[1].text) + item_index = 2 + head = item_tokens[item_index] if item_index < len(item_tokens) else None + if head and head.is_keyword("UNIQUE"): + if ( + item_index + 1 >= len(item_tokens) + or item_tokens[item_index + 1].text != "(" + ): + raise ParseError("Table UNIQUE must be followed by a column list") + columns, next_index = _unique_columns(item, item_tokens, item_index + 1) + conflict, next_index = _conflict_after(item_tokens, next_index) + if next_index != len(item_tokens): + raise ParseError("Unexpected SQL after UNIQUE constraint") + source_start = item_tokens[0].start + source_end = item_tokens[next_index - 1].end + uniques.append( + Unique( + columns, + name=constraint_name, + conflict=conflict, + sql=item[source_start:source_end], + start=body_start + item_start + source_start, + end=body_start + item_start + source_end, + ) + ) + continue + if ( + head + and head.kind == "word" + and head.text.upper() in _TABLE_CONSTRAINT_KEYWORDS + ): + continue + column = _unquote(item_tokens[0].text) + uniques.extend( + _column_uniques( + item, + item_tokens, + column, + body_start + item_start, + ) + ) + return uniques + + 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) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 48c5d5d..48987a6 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -32,10 +32,13 @@ from .create_table_parser import ( Check, ColumnComments, ParseError, + Unique, + UniqueColumn, check_references_identifier, parse_autoincrement, parse_checks, parse_column_comments, + parse_uniques, rewrite_check_expression, sql_ends_in_line_comment, ) @@ -104,6 +107,24 @@ def _check_constraint_sql(check: Check) -> str: return f"{prefix}CHECK ({check.check}{newline})" +def _unique_constraint_sql(unique: Unique) -> str: + prefix = f"CONSTRAINT {quote_identifier(unique.name)} " if unique.name else "" + if unique.column: + constraint = "UNIQUE" + else: + columns = [] + for column in unique.columns: + column_sql = quote_identifier(column.name) + if column.collation: + column_sql += f" COLLATE {quote_identifier(column.collation)}" + if column.order: + column_sql += f" {column.order}" + columns.append(column_sql) + constraint = "UNIQUE ({})".format(", ".join(columns)) + conflict = f" ON CONFLICT {unique.conflict}" if unique.conflict else "" + return f"{prefix}{constraint}{conflict}" + + def _column_definition_with_comments( definition: str, comments: ColumnComments | None ) -> str: @@ -1424,6 +1445,7 @@ class Database: _checks: Iterable[Check] | None = None, _column_comments: Mapping[str, ColumnComments] | None = None, _autoincrement: str | None = None, + _uniques: Iterable[Unique] | None = None, ) -> str: """ Returns the SQL ``CREATE TABLE`` statement for creating the specified table. @@ -1486,6 +1508,60 @@ class Database: checks_by_column.setdefault(column, []).append(check) else: table_checks.append(check) + uniques_by_column: dict[str, list[Unique]] = {} + table_uniques: list[Unique] = [] + for unique in _uniques or (): + resolved_unique = Unique( + tuple( + UniqueColumn( + resolve_casing(column.name, columns), + collation=column.collation, + order=column.order, + ) + for column in unique.columns + ), + name=unique.name, + column=( + resolve_casing(unique.column, columns) if unique.column else "" + ), + conflict=unique.conflict, + ) + missing = [ + column.name + for column in resolved_unique.columns + if column.name not in columns + ] + if missing: + raise AlterError( + "No such column for UNIQUE constraint: {}".format( + ", ".join(missing) + ) + ) + if resolved_unique.column: + if ( + len(resolved_unique.columns) != 1 + or resolved_unique.columns[0].name != resolved_unique.column + ): + raise AlterError("Invalid column-level UNIQUE constraint") + if any( + column.collation or column.order + for column in resolved_unique.columns + ): + # Render this as a table constraint so the collation or sort + # order that governs uniqueness can be represented explicitly. + table_uniques.append( + Unique( + resolved_unique.columns, + name=resolved_unique.name, + conflict=resolved_unique.conflict, + ) + ) + else: + uniques_by_column.setdefault(resolved_unique.column, []).append( + resolved_unique + ) + else: + table_uniques.append(resolved_unique) if not columns: raise ValueError("Tables must have at least one column") if not all(n in columns for n in not_null): @@ -1554,6 +1630,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( + _unique_constraint_sql(unique) + for unique in uniques_by_column.get(column_name, ()) + ) column_extras.extend( _check_constraint_sql(check) for check in checks_by_column.get(column_name, ()) @@ -1600,6 +1680,9 @@ class Database: actions=_fk_actions_sql(fk), ) ) + column_defs.extend( + f" {_unique_constraint_sql(unique)}" for unique in table_uniques + ) column_defs.extend( f" {_check_constraint_sql(check)}" for check in table_checks ) @@ -2763,6 +2846,7 @@ class Table(Queryable): existing_checks = self.checks existing_column_comments = parse_column_comments(self.schema) existing_autoincrement = parse_autoincrement(self.schema) + existing_uniques = parse_uniques(self.schema) except ParseError as ex: raise TransformError( f"Could not parse table schema for table {self.name!r}: {ex}" @@ -2789,6 +2873,37 @@ class Table(Queryable): ) ) + create_table_uniques: list[Unique] = [] + for unique in existing_uniques: + columns = tuple( + UniqueColumn( + resolve_casing(column.name, existing_columns), + collation=column.collation, + order=column.order, + ) + for column in unique.columns + ) + if any(column.name in drop for column in columns): + continue + owner = ( + resolve_casing(unique.column, existing_columns) if unique.column else "" + ) + create_table_uniques.append( + Unique( + tuple( + UniqueColumn( + rename.get(column.name) or column.name, + collation=column.collation, + order=column.order, + ) + for column in columns + ), + name=unique.name, + column=rename.get(owner) or owner, + conflict=unique.conflict, + ) + ) + create_table_column_comments: dict[str, ColumnComments] = {} for column, comments in existing_column_comments.items(): owner = resolve_casing(column, existing_columns) @@ -2974,6 +3089,7 @@ class Table(Queryable): _checks=create_table_checks, _column_comments=create_table_column_comments, _autoincrement=create_table_autoincrement, + _uniques=create_table_uniques, ).strip() ) @@ -3008,6 +3124,9 @@ class Table(Queryable): {"index_name": index.name}, ).fetchall()[0][0] if index_sql is None: + if index.origin == "u": + # UNIQUE constraints are reproduced in CREATE TABLE above. + continue raise TransformError( f"Index '{index.name}' on table '{self.name}' does not have a " "CREATE INDEX statement. You must manually drop this index prior to running this " diff --git a/tests/test_create_table_parser.py b/tests/test_create_table_parser.py index 54bf221..74a089c 100644 --- a/tests/test_create_table_parser.py +++ b/tests/test_create_table_parser.py @@ -8,9 +8,12 @@ from sqlite_utils.create_table_parser import ( Check, ColumnComments, ParseError, + Unique, + UniqueColumn, parse_autoincrement, parse_checks, parse_column_comments, + parse_uniques, ) @@ -148,6 +151,51 @@ def test_parse_autoincrement(sql, expected): assert parse_autoincrement(sql) == expected +def test_parse_column_and_table_uniques(): + sql = """ + CREATE TABLE memberships ( + email TEXT COLLATE RTRIM CONSTRAINT unique_email UNIQUE ON CONFLICT IGNORE, + account_id INTEGER, + CONSTRAINT unique_membership UNIQUE ( + account_id DESC, + email COLLATE NOCASE ASC + ) ON CONFLICT REPLACE + ) + """ + sqlite3.connect(":memory:").execute(sql) + assert parse_uniques(sql) == [ + Unique( + (UniqueColumn("email", collation="RTRIM"),), + name="unique_email", + column="email", + conflict="IGNORE", + ), + Unique( + ( + UniqueColumn("account_id", order="DESC"), + UniqueColumn("email", collation="NOCASE", order="ASC"), + ), + name="unique_membership", + conflict="REPLACE", + ), + ] + uniques = parse_uniques(sql) + assert uniques[0].sql == "CONSTRAINT unique_email UNIQUE ON CONFLICT IGNORE" + assert sql[uniques[1].start : uniques[1].end] == uniques[1].sql + + +def test_unique_like_text_in_comments_and_checks_is_ignored(): + sql = """ + CREATE TABLE t ( + value TEXT /* UNIQUE ON CONFLICT REPLACE */ + CHECK(value != 'UNIQUE(other)'), + other TEXT + ) + """ + sqlite3.connect(":memory:").execute(sql) + assert parse_uniques(sql) == [] + + comment_or_space = st.sampled_from( [ " ", diff --git a/tests/test_transform.py b/tests/test_transform.py index e6096cd..3be6c6f 100644 --- a/tests/test_transform.py +++ b/tests/test_transform.py @@ -1033,24 +1033,78 @@ def test_transform_with_unique_constraint_implicit_index(fresh_db): fresh_db.execute(""" CREATE TABLE dogs ( id INTEGER PRIMARY KEY, - name TEXT UNIQUE, + name TEXT UNIQUE ON CONFLICT IGNORE, age INTEGER ); """) dogs.insert({"id": 1, "name": "Cleo", "age": 5}) - # Attempt to transform the table without modifying 'name' - with pytest.raises(TransformError) as excinfo: - dogs.transform(types={"age": str}) + dogs.transform(types={"age": str}, rename={"name": "dog_name"}) + + assert 'dog_name" TEXT UNIQUE ON CONFLICT IGNORE' in dogs.schema + dogs.insert({"id": 2, "dog_name": "Cleo", "age": "6"}) + assert list(dogs.rows) == [{"id": 1, "dog_name": "Cleo", "age": "5"}] + + +def test_transform_preserves_composite_unique_constraint(fresh_db): + fresh_db.execute(""" + CREATE TABLE memberships ( + account_id INTEGER, + email TEXT, + note TEXT, + CONSTRAINT unique_membership + UNIQUE (account_id DESC, email COLLATE NOCASE) + ON CONFLICT ABORT + ) + """) + memberships = fresh_db.table("memberships") + memberships.insert({"account_id": 1, "email": "one@example.com", "note": "x"}) + + memberships.transform(rename={"account_id": "organization_id"}, types={"note": str}) assert ( - "Index 'sqlite_autoindex_dogs_1' on table 'dogs' does not have a CREATE INDEX statement." - in str(excinfo.value) - ) - assert ( - "You must manually drop this index prior to running this transformation and manually recreate the new index after running this transformation." - in str(excinfo.value) + 'CONSTRAINT "unique_membership" UNIQUE ' + '("organization_id" DESC, "email" COLLATE "NOCASE") ON CONFLICT ABORT' + in memberships.schema ) + with pytest.raises(sqlite3.IntegrityError): + memberships.insert( + {"organization_id": 1, "email": "ONE@example.com", "note": "y"} + ) + + +def test_transform_preserves_column_unique_collation(fresh_db): + fresh_db.execute(""" + CREATE TABLE people ( + id INTEGER PRIMARY KEY, + name TEXT COLLATE NOCASE UNIQUE + ) + """) + people = fresh_db.table("people") + people.insert({"id": 1, "name": "Cleo"}) + + people.transform(rename={"name": "full_name"}) + + assert 'UNIQUE ("full_name" COLLATE "NOCASE")' in people.schema + with pytest.raises(sqlite3.IntegrityError): + people.insert({"id": 2, "full_name": "cleo"}) + + +def test_transform_drops_entire_composite_unique_constraint(fresh_db): + fresh_db.execute(""" + CREATE TABLE memberships ( + account_id INTEGER, + email TEXT, + UNIQUE (account_id, email) + ) + """) + memberships = fresh_db.table("memberships") + memberships.insert({"account_id": 1, "email": "one@example.com"}) + + memberships.transform(drop={"email"}) + + assert "UNIQUE" not in memberships.schema + memberships.insert({"account_id": 1}) def test_transform_preserves_autoincrement_and_sequence(fresh_db):