Fix when referring table has multiple FKs

Prior to this change, a part of the clean_up function that changes the
name of foreign keys back to their original wasn't handling referring
tables that had more than one FK to the original table.

This commit changes that, so that no matter how many FKs a referring
table has to the original table there won't be any problems.

Author: marcelofern

src/psycopack/_repack.py

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

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