diff --git a/src/psycopack/_commands.py b/src/psycopack/_commands.py index 3d248ba..0432777 100644 --- a/src/psycopack/_commands.py +++ b/src/psycopack/_commands.py @@ -327,12 +327,36 @@ def create_change_log_copy_function( self.cur.execute( psycopg.sql.SQL( dedent(""" - CREATE OR REPLACE FUNCTION {schema}.{function}(VARIADIC pks BIGINT[]) + CREATE OR REPLACE FUNCTION {schema}.{function}(batch_size BIGINT) RETURNS VOID LANGUAGE plpgsql SECURITY DEFINER AS $$ + DECLARE + pks BIGINT[]; BEGIN + -- Lock change log rows FIRST, then aggregate + SELECT + ARRAY_AGG(src_pk) + INTO + pks + FROM ( + SELECT + src_pk + FROM + {schema}.{change_log} + ORDER BY + src_pk + LIMIT + batch_size + FOR UPDATE SKIP LOCKED + ) locked_rows; + + -- Nothing to do + IF pks IS NULL THEN + RETURN; + END IF; + -- Lock source rows PERFORM 1 FROM {schema}.{table_from} @@ -345,12 +369,6 @@ def create_change_log_copy_function( WHERE {pk_column} = ANY (pks) FOR UPDATE; - -- Lock change log rows - PERFORM 1 - FROM {schema}.{change_log} - WHERE src_pk = ANY (pks) - FOR UPDATE; - -- Delete destination rows DELETE FROM {schema}.{table_to} WHERE {pk_column} = ANY (pks); @@ -713,17 +731,14 @@ def execute_copy_function( ) def execute_change_log_copy_function( - self, - *, - function: str, - pks: list[int], + self, *, function: str, batch_size: int ) -> None: self.cur.execute( - psycopg.sql.SQL("SELECT {schema}.{function}({pks});") + psycopg.sql.SQL("SELECT {schema}.{function}({batch_size});") .format( function=psycopg.sql.Identifier(function), schema=psycopg.sql.Identifier(self.schema), - pks=psycopg.sql.SQL(", ").join(map(psycopg.sql.Literal, pks)), + batch_size=psycopg.sql.Literal(batch_size), ) .as_string(self.conn) ) diff --git a/src/psycopack/_repack.py b/src/psycopack/_repack.py index 1b7d24b..2366398 100644 --- a/src/psycopack/_repack.py +++ b/src/psycopack/_repack.py @@ -466,7 +466,7 @@ def _post_sync_update_for_change_log(self) -> None: self.command.execute_change_log_copy_function( function=self.change_log_copy_function, - pks=[change.src_pk for change in change_log_batch], + batch_size=self.change_log_batch_size, ) def swap(self) -> None: @@ -581,13 +581,41 @@ def clean_up(self) -> None: idx_data["idx_to"] = idx.name # Rename foreign keys from other tables using a dict data structure to - # hold names from/to. - table_to_fk: dict[str, dict[str, str]] = {} + # hold names from/to. Support multiple FKs from the same referring table. + # Match FKs by their definition (columns), similar to how indexes are matched. + fk_definitions: dict[str, list[dict[str, str]]] = defaultdict(list) + for fk in self.introspector.get_referring_fks(table=self.repacked_name): - table_to_fk[fk.referring_table] = {"cons_from": fk.name} + # Normalize the FK definition by replacing the old table name with the new one + # This allows us to match FKs based on their structure. + fk_def = fk.definition.replace( + f"REFERENCES {self.schema}.{self.repacked_name}", + f"REFERENCES {self.schema}.{self.table}", + ).replace( + f"REFERENCES {self.repacked_name}", + f"REFERENCES {self.table}", + ) + fk_definitions[fk_def].append( + { + "cons_from": fk.name, + "referring_table": fk.referring_table, + } + ) for fk in self.introspector.get_referring_fks(table=self.table): - table_to_fk[fk.referring_table]["cons_to"] = fk.name + # Use the FK definition as the key to match with the old FK. + fk_def = fk.definition + if fk_def in fk_definitions: + # Find the first unmatched FK with this definition. + fk_pair = next( + ( + pair + for pair in fk_definitions[fk_def] + if "cons_to" not in pair + ), + ) + if fk_pair: + fk_pair["cons_to"] = fk.name with ( self.command.db_transaction(), @@ -616,17 +644,17 @@ def clean_up(self) -> None: idx_to=index_data["idx_from"], ) - for table in table_to_fk: - fk_data = table_to_fk[table] - self.command.drop_constraint( - table=table, - constraint=fk_data["cons_from"], - ) - self.command.rename_constraint( - table=table, - cons_from=fk_data["cons_to"], - cons_to=fk_data["cons_from"], - ) + for fk_def in fk_definitions: + for fk_data in fk_definitions[fk_def]: + self.command.drop_constraint( + table=fk_data["referring_table"], + constraint=fk_data["cons_from"], + ) + self.command.rename_constraint( + table=fk_data["referring_table"], + cons_from=fk_data["cons_to"], + cons_to=fk_data["cons_from"], + ) self.command.drop_table_if_exists(table=self.repacked_name) self.command.drop_table_if_exists(table=self.backfill_log) diff --git a/tests/test_repack.py b/tests/test_repack.py index 30aaf83..563fd78 100644 --- a/tests/test_repack.py +++ b/tests/test_repack.py @@ -2603,3 +2603,135 @@ def test_repack_with_change_log_strategy( repack=repack, cur=cur, ) + + +@pytest.mark.parametrize( + "sync_strategy", + [SyncStrategy.DIRECT_TRIGGER, SyncStrategy.CHANGE_LOG], +) +def test_multiple_foreign_keys_from_same_referring_table( + connection: _psycopg.Connection, + sync_strategy: SyncStrategy, +) -> None: + """ + Test that multiple foreign keys from the same referring table are correctly + handled during clean_up(). This verifies the fix for the bug where only one + FK would be renamed correctly when multiple FKs from the same table + existed. + """ + with _cur.get_cursor(connection, logged=True) as cur: + # Create the table to be repacked (simulating a users table) + cur.execute( + dedent(""" + CREATE TABLE person ( + id SERIAL PRIMARY KEY, + name VARCHAR(100) + ); + """) + ) + cur.execute("INSERT INTO person (name) VALUES ('Anna'), ('Bob'), ('Carlos');") + + # Create a referring table with TWO foreign keys to the same table + # (simulating created_by and updated_by columns) + cur.execute( + dedent(""" + CREATE TABLE posts ( + id SERIAL PRIMARY KEY, + title VARCHAR(200), + created_by_id INTEGER REFERENCES person(id), + updated_by_id INTEGER REFERENCES person(id) + ); + """) + ) + cur.execute( + dedent(""" + INSERT INTO posts (title, created_by_id, updated_by_id) + VALUES + ('Post 1', 1, 1), + ('Post 2', 1, 2), + ('Post 3', 2, 3); + """) + ) + + # Collect FK info before repacking + introspector = _introspect.Introspector( + conn=connection, + cur=cur, + schema="public", + ) + fks_before = introspector.get_referring_fks(table="person") + assert len(fks_before) == 2, "Expected 2 foreign keys from posts table" + fk_names_before = {fk.name for fk in fks_before} + + # Run the full repack + repack = Psycopack( + table="person", + batch_size=10, + conn=connection, + cur=cur, + sync_strategy=sync_strategy, + change_log_batch_size=100, + ) + repack.full() + + # Verify both foreign keys still exist and point to the person table + fks_after = introspector.get_referring_fks(table="person") + assert len(fks_after) == 2, "Expected 2 foreign keys after repack" + fk_names_after = {fk.name for fk in fks_after} + + # FK names should be preserved + assert fk_names_before == fk_names_after, ( + f"FK names changed: before={fk_names_before}, after={fk_names_after}" + ) + + # Verify both FKs are still valid by checking constraints + cur.execute( + dedent(""" + SELECT + conname, + conrelid::regclass, + confrelid::regclass + FROM + pg_constraint + WHERE + confrelid = 'person'::regclass + AND contype = 'f' + ORDER BY + conname; + """) + ) + constraints = cur.fetchall() + assert len(constraints) == 2 + + # Verify both FKs point to the person table (not the old repacked table) + for constraint in constraints: + constraint_name, referring_table, referenced_table = constraint + assert str(referenced_table) == "person" + assert str(referring_table) == "posts" + + # Verify data integrity: FK constraints should still work + cur.execute("SELECT COUNT(*) FROM posts;") + row = cur.fetchone() + assert row is not None + assert row[0] == 3 + + # Verify we can still query using the foreign keys + cur.execute( + dedent(""" + SELECT + p.title, + u1.name as creator, + u2.name as updater + FROM posts p + JOIN person u1 + ON p.created_by_id = u1.id + JOIN person u2 + ON p.updated_by_id = u2.id + ORDER BY p.id; + """) + ) + rows = cur.fetchall() + assert len(rows) == 3 + assert rows[0] == ("Post 1", "Anna", "Anna") + assert rows[1] == ("Post 2", "Anna", "Bob") + assert rows[2] == ("Post 3", "Bob", "Carlos")