Multi fk fixes and better batch inserts

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

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)

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