mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-09-13 03:54:23 +02:00
Use .transform() instead of modifying sqlite_master for add_foreign_keys(), refs #577
This commit is contained in:
Simon Willison
parent
1dc6b5aa64
commit
842b61321f
4 changed files with 133 additions and 108 deletions
|
|
@ -747,9 +747,16 @@ class Database:
|
|||
def resolve_foreign_keys(
|
||||
self, name: str, foreign_keys: ForeignKeysType
|
||||
) -> List[ForeignKey]:
|
||||
# foreign_keys may be a list of column names, a list of ForeignKey tuples,
|
||||
# a list of tuple-pairs or a list of tuple-triples. We want to turn
|
||||
# it into a list of ForeignKey tuples
|
||||
"""
|
||||
Given a list of differing foreign_keys definitions, return a list of
|
||||
fully resolved ForeignKey() named tuples.
|
||||
|
||||
:param name: Name of table that foreign keys are being defined for
|
||||
:param foreign_keys: List of foreign keys, each of which can be a
|
||||
string, a ForeignKey() named tuple, a tuple of (column, other_table),
|
||||
or a tuple of (column, other_table, other_column), or a tuple of
|
||||
(table, column, other_table, other_column)
|
||||
"""
|
||||
table = cast(Table, self[name])
|
||||
if all(isinstance(fk, ForeignKey) for fk in foreign_keys):
|
||||
return cast(List[ForeignKey], foreign_keys)
|
||||
|
|
@ -767,6 +774,11 @@ class Database:
|
|||
), "foreign_keys= should be a list of tuples"
|
||||
fks = []
|
||||
for tuple_or_list in foreign_keys:
|
||||
if len(tuple_or_list) == 4:
|
||||
assert (
|
||||
tuple_or_list[0] == name
|
||||
), "First item in {} should have been {}".format(tuple_or_list, name)
|
||||
tuple_or_list = tuple(tuple_or_list[1:])
|
||||
assert len(tuple_or_list) in (
|
||||
2,
|
||||
3,
|
||||
|
|
@ -1148,32 +1160,14 @@ class Database:
|
|||
(table, column, other_table, other_column)
|
||||
)
|
||||
|
||||
# Construct SQL for use with "UPDATE sqlite_master SET sql = ? WHERE name = ?"
|
||||
table_sql: Dict[str, str] = {}
|
||||
for table, column, other_table, other_column in foreign_keys_to_create:
|
||||
old_sql = table_sql.get(table, self[table].schema)
|
||||
extra_sql = ",\n FOREIGN KEY([{column}]) REFERENCES [{other_table}]([{other_column}])\n".format(
|
||||
column=column, other_table=other_table, other_column=other_column
|
||||
)
|
||||
# Stick that bit in at the very end just before the closing ')'
|
||||
last_paren = old_sql.rindex(")")
|
||||
new_sql = old_sql[:last_paren].strip() + extra_sql + old_sql[last_paren:]
|
||||
table_sql[table] = new_sql
|
||||
# Group them by table
|
||||
by_table = {}
|
||||
for fk in foreign_keys_to_create:
|
||||
by_table.setdefault(fk[0], []).append(fk)
|
||||
|
||||
for table, fks in by_table.items():
|
||||
self[table].transform(add_foreign_keys=fks)
|
||||
|
||||
# And execute it all within a single transaction
|
||||
with self.conn:
|
||||
cursor = self.conn.cursor()
|
||||
schema_version = cursor.execute("PRAGMA schema_version").fetchone()[0]
|
||||
cursor.execute("PRAGMA writable_schema = 1")
|
||||
for table_name, new_sql in table_sql.items():
|
||||
cursor.execute(
|
||||
"UPDATE sqlite_master SET sql = ? WHERE name = ?",
|
||||
(new_sql, table_name),
|
||||
)
|
||||
cursor.execute("PRAGMA schema_version = %d" % (schema_version + 1))
|
||||
cursor.execute("PRAGMA writable_schema = 0")
|
||||
# Have to VACUUM outside the transaction to ensure .foreign_keys property
|
||||
# can see the newly created foreign key.
|
||||
self.vacuum()
|
||||
|
||||
def index_foreign_keys(self):
|
||||
|
|
@ -1704,7 +1698,9 @@ class Table(Queryable):
|
|||
pk: Optional[Any] = DEFAULT,
|
||||
not_null: Optional[Iterable[str]] = None,
|
||||
defaults: Optional[Dict[str, Any]] = None,
|
||||
drop_foreign_keys: Optional[Iterable] = None,
|
||||
drop_foreign_keys: Optional[Iterable[str]] = None,
|
||||
add_foreign_keys: Optional[ForeignKeysType] = None,
|
||||
foreign_keys: Optional[ForeignKeysType] = None,
|
||||
column_order: Optional[List[str]] = None,
|
||||
keep_table: Optional[str] = None,
|
||||
) -> "Table":
|
||||
|
|
@ -1721,6 +1717,8 @@ class Table(Queryable):
|
|||
:param not_null: Columns to set as ``NOT NULL``
|
||||
:param defaults: Default values for columns
|
||||
:param drop_foreign_keys: Names of columns that should have their foreign key constraints removed
|
||||
:param add_foreign_keys: List of foreign keys to add to the table
|
||||
:param foreign_keys: List of foreign keys to set for the table, replacing any existing foreign keys
|
||||
:param column_order: List of strings specifying a full or partial column order
|
||||
to use when creating the table
|
||||
:param keep_table: If specified, the existing table will be renamed to this and will not be
|
||||
|
|
@ -1735,6 +1733,8 @@ class Table(Queryable):
|
|||
not_null=not_null,
|
||||
defaults=defaults,
|
||||
drop_foreign_keys=drop_foreign_keys,
|
||||
add_foreign_keys=add_foreign_keys,
|
||||
foreign_keys=foreign_keys,
|
||||
column_order=column_order,
|
||||
keep_table=keep_table,
|
||||
)
|
||||
|
|
@ -1765,6 +1765,8 @@ class Table(Queryable):
|
|||
not_null: Optional[Iterable[str]] = None,
|
||||
defaults: Optional[Dict[str, Any]] = None,
|
||||
drop_foreign_keys: Optional[Iterable] = None,
|
||||
add_foreign_keys: Optional[ForeignKeysType] = None,
|
||||
foreign_keys: Optional[ForeignKeysType] = None,
|
||||
column_order: Optional[List[str]] = None,
|
||||
tmp_suffix: Optional[str] = None,
|
||||
keep_table: Optional[str] = None,
|
||||
|
|
@ -1779,6 +1781,8 @@ class Table(Queryable):
|
|||
:param not_null: Columns to set as ``NOT NULL``
|
||||
:param defaults: Default values for columns
|
||||
:param drop_foreign_keys: Names of columns that should have their foreign key constraints removed
|
||||
:param add_foreign_keys: List of foreign keys to add to the table
|
||||
:param foreign_keys: List of foreign keys to set for the table, replacing any existing foreign keys
|
||||
:param column_order: List of strings specifying a full or partial column order
|
||||
to use when creating the table
|
||||
:param tmp_suffix: Suffix to use for the temporary table name
|
||||
|
|
@ -1788,6 +1792,43 @@ class Table(Queryable):
|
|||
types = types or {}
|
||||
rename = rename or {}
|
||||
drop = drop or set()
|
||||
|
||||
if foreign_keys is not None:
|
||||
if add_foreign_keys is not None:
|
||||
raise ValueError(
|
||||
"Cannot specify both foreign_keys and add_foreign_keys"
|
||||
)
|
||||
if drop_foreign_keys is not None:
|
||||
raise ValueError(
|
||||
"Cannot specify both foreign_keys and drop_foreign_keys"
|
||||
)
|
||||
create_table_foreign_keys = foreign_keys
|
||||
else:
|
||||
# Construct foreign_keys from current, plus add_foreign_keys, minus drop_foreign_keys
|
||||
create_table_foreign_keys = []
|
||||
for table, column, other_table, other_column in self.foreign_keys:
|
||||
# Copy over old foreign keys, unless we are dropping them
|
||||
if (drop_foreign_keys is None) or (column not in drop_foreign_keys):
|
||||
create_table_foreign_keys.append(
|
||||
ForeignKey(
|
||||
table,
|
||||
rename.get(column) or column,
|
||||
other_table,
|
||||
other_column,
|
||||
)
|
||||
)
|
||||
# Add new foreign keys
|
||||
if add_foreign_keys is not None:
|
||||
for fk in self.db.resolve_foreign_keys(self.name, add_foreign_keys):
|
||||
create_table_foreign_keys.append(
|
||||
ForeignKey(
|
||||
self.name,
|
||||
rename.get(fk.column) or fk.column,
|
||||
fk.other_table,
|
||||
fk.other_column,
|
||||
)
|
||||
)
|
||||
|
||||
new_table_name = "{}_new_{}".format(
|
||||
self.name, tmp_suffix or os.urandom(6).hex()
|
||||
)
|
||||
|
|
@ -1847,14 +1888,6 @@ class Table(Queryable):
|
|||
{rename.get(c) or c: v for c, v in defaults.items()}
|
||||
)
|
||||
|
||||
# foreign_keys
|
||||
create_table_foreign_keys = []
|
||||
for table, column, other_table, other_column in self.foreign_keys:
|
||||
if (drop_foreign_keys is None) or (column not in drop_foreign_keys):
|
||||
create_table_foreign_keys.append(
|
||||
(rename.get(column) or column, other_table, other_column)
|
||||
)
|
||||
|
||||
if column_order is not None:
|
||||
column_order = [rename.get(col) or col for col in column_order]
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue