mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-08-14 13:24:11 +02:00
Preserve composite UNIQUE constraints in transforms
This commit is contained in:
Simon Willison
parent
2b52b5ed6f
commit
75ba588462
5 changed files with 425 additions and 10 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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 "
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue