Preserve composite UNIQUE constraints in transforms

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

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)

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 "

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

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'
+    dogs.transform(types={"age": str}, rename={"name": "dog_name"})
-    with pytest.raises(TransformError) as excinfo:
+
-        dogs.transform(types={"age": str})
+    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."
+        'CONSTRAINT "unique_membership" UNIQUE '
-        in str(excinfo.value)
+        '("organization_id" DESC, "email" COLLATE "NOCASE") ON CONFLICT ABORT'
-    )
+        in memberships.schema
-    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)
     )
+    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):