Add self-referential FK support and additional tests

- Fix transform_sql to update other_column for self-referential FKs
  when the referenced column is being renamed
- Skip FK validation during transform when update_incoming_fks=True
  since self-referential FKs temporarily reference non-existent columns
- Add tests for multiple tables referencing renamed column
- Add test for self-referential FK handling
This commit is contained in:
Claude 2026-01-21 14:56:06 +00:00
commit e22b34ee24
No known key found for this signature in database
2 changed files with 128 additions and 13 deletions

View file

@ -1944,19 +1944,29 @@ class Table(Queryable):
finally:
self.db._skip_fk_validation = False
sqls = self.transform_sql(
types=types,
rename=rename,
drop=drop,
pk=pk,
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,
)
# Skip FK validation for main transform if update_incoming_fks is True
# because self-referential FKs will reference the new column name
# that only exists after the transform completes
if update_incoming_fks and rename:
self.db._skip_fk_validation = True
try:
sqls = self.transform_sql(
types=types,
rename=rename,
drop=drop,
pk=pk,
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,
)
finally:
if update_incoming_fks and rename:
self.db._skip_fk_validation = False
pragma_foreign_keys_was_on = self.db.execute("PRAGMA foreign_keys").fetchone()[
0
]
@ -2034,6 +2044,9 @@ class Table(Queryable):
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):
# For self-referential FKs, also update the referenced column if renamed
if other_table == self.name:
other_column = rename.get(other_column) or other_column
create_table_foreign_keys.append(
ForeignKey(
table,

View file

@ -708,3 +708,105 @@ def test_transform_update_incoming_fks_on_column_rename(fresh_db):
assert fresh_db.execute("PRAGMA foreign_keys").fetchone()[0] == 1
violations = list(fresh_db.execute("PRAGMA foreign_key_check").fetchall())
assert violations == []
def test_transform_update_incoming_fks_multiple_tables(fresh_db):
"""
Test that update_incoming_fks=True updates FK constraints in multiple tables
when a referenced column is renamed.
"""
fresh_db.execute("PRAGMA foreign_keys=ON")
# Create authors table with id as PK
fresh_db["authors"].insert({"id": 1, "name": "Alice"}, pk="id")
# Create multiple tables with FKs to authors.id
fresh_db["books"].insert(
{"id": 1, "title": "Book A", "author_id": 1},
pk="id",
foreign_keys=[("author_id", "authors", "id")],
)
fresh_db["articles"].insert(
{"id": 1, "headline": "Article A", "writer_id": 1},
pk="id",
foreign_keys=[("writer_id", "authors", "id")],
)
fresh_db["quotes"].insert(
{"id": 1, "text": "Quote A", "speaker_id": 1},
pk="id",
foreign_keys=[("speaker_id", "authors", "id")],
)
# Rename authors.id to authors.author_pk with update_incoming_fks=True
fresh_db["authors"].transform(
rename={"id": "author_pk"},
update_incoming_fks=True,
)
# Verify authors column was renamed
assert "author_pk" in fresh_db["authors"].columns_dict
assert "id" not in fresh_db["authors"].columns_dict
# Verify all FKs were updated
assert fresh_db["books"].foreign_keys == [
ForeignKey(table="books", column="author_id", other_table="authors", other_column="author_pk")
]
assert fresh_db["articles"].foreign_keys == [
ForeignKey(table="articles", column="writer_id", other_table="authors", other_column="author_pk")
]
assert fresh_db["quotes"].foreign_keys == [
ForeignKey(table="quotes", column="speaker_id", other_table="authors", other_column="author_pk")
]
# Verify FK enforcement still works
violations = list(fresh_db.execute("PRAGMA foreign_key_check").fetchall())
assert violations == []
def test_transform_update_incoming_fks_self_referential(fresh_db):
"""
Test that update_incoming_fks=True handles self-referential FK constraints.
"""
fresh_db.execute("PRAGMA foreign_keys=ON")
# Create employees table with self-referential FK (manager_id -> id)
fresh_db.execute("""
CREATE TABLE employees (
id INTEGER PRIMARY KEY,
name TEXT,
manager_id INTEGER REFERENCES employees(id)
)
""")
fresh_db["employees"].insert_all([
{"id": 1, "name": "CEO", "manager_id": None},
{"id": 2, "name": "VP", "manager_id": 1},
{"id": 3, "name": "Dev", "manager_id": 2},
])
# Verify initial FK
assert fresh_db["employees"].foreign_keys == [
ForeignKey(table="employees", column="manager_id", other_table="employees", other_column="id")
]
# Rename employees.id to employees.emp_id with update_incoming_fks=True
fresh_db["employees"].transform(
rename={"id": "emp_id"},
update_incoming_fks=True,
)
# Verify column was renamed
assert "emp_id" in fresh_db["employees"].columns_dict
assert "id" not in fresh_db["employees"].columns_dict
# Verify self-referential FK was updated
assert fresh_db["employees"].foreign_keys == [
ForeignKey(table="employees", column="manager_id", other_table="employees", other_column="emp_id")
]
# Verify data integrity
rows = list(fresh_db.execute("SELECT * FROM employees ORDER BY emp_id").fetchall())
assert rows == [(1, "CEO", None), (2, "VP", 1), (3, "Dev", 2)]
# Verify FK enforcement still works
violations = list(fresh_db.execute("PRAGMA foreign_key_check").fetchall())
assert violations == []