Fix handling of PK sequence value

Prior to this change, the current value of the table's PK sequence
wasn't transferred to the repacked table's PK sequence. Without that,
new rows would be inserted into the repacked table starting with a PK
value of 1 (which would likely conflict with existing rows).

This change adds a command during swap to transfer the sequence value
from one table to the other. There is also special handling for the edge
case of a bigint conversion where the PK sequence has been set to negative
(presumably after positive PK values were exhausted); in that case, the
sequence is reset to the first positive bigint value not within the integer
range, which is 2**31. Note that this edge case won't be handled correctly
during a revert_swap(); addressing that is left as a TODO for later.

Author: timb07

src/psycopack/_commands.py

@@ -486,6 +486,30 @@ def swap_pk_sequence_name(self, *, first_table: str, second_table: str) -> None:
         self.rename_sequence(seq_from=second_seq, seq_to=first_seq)
         self.rename_sequence(seq_from=temp_seq, seq_to=second_seq)
 
+    def transfer_pk_sequence_value(
+        self, *, source_table: str, dest_table: str, convert_pk_to_bigint: bool
+    ) -> None:
+        source_seq = self.introspector.get_pk_sequence_name(table=source_table)
+        dest_seq = self.introspector.get_pk_sequence_name(table=dest_table)
+        value = self.introspector.get_pk_sequence_value(seq=source_seq)
+
+        if convert_pk_to_bigint and value < 0:
+            # special case handling where negative PK values were used before bigint conversion
+            value = 2**31  # reset to positive, specifically the first bigint value
+
+        # TODO: try to correctly restore a negative PK sequence value if we revert swap
+        # while doing a bigint conversion
+
+        self.cur.execute(
+            psycopg.sql.SQL("SELECT setval('{schema}.{sequence}', {value});")
+            .format(
+                schema=psycopg.sql.Identifier(self.schema),
+                sequence=psycopg.sql.Identifier(dest_seq),
+                value=psycopg.sql.SQL(str(value)),
+            )
+            .as_string(self.conn)
+        )
+
     def acquire_access_exclusive_lock(self, *, table: str) -> None:
         self.cur.execute(
             psycopg.sql.SQL("LOCK TABLE {schema}.{table} IN ACCESS EXCLUSIVE MODE;")

src/psycopack/_introspect.py

@@ -613,6 +613,21 @@ def get_pk_sequence_name(self, *, table: str) -> str:
             assert isinstance(seq, str)
             return seq
 
+    def get_pk_sequence_value(self, *, seq: str) -> int:
+        self.cur.execute(
+            psycopg.sql.SQL("SELECT last_value FROM {schema}.{sequence};")
+            .format(
+                schema=psycopg.sql.Identifier(self.schema),
+                sequence=psycopg.sql.Identifier(seq),
+            )
+            .as_string(self.conn)
+        )
+        result = self.cur.fetchone()
+        assert result is not None
+        value = result[0]
+        assert isinstance(value, int)
+        return value
+
     def get_backfill_batch(self, *, table: str) -> BackfillBatch | None:
         self.cur.execute(
             psycopg.sql.SQL(

src/psycopack/_repack.py

@@ -381,6 +381,11 @@ def swap(self) -> None:
                     self.command.swap_pk_sequence_name(
                         first_table=self.table, second_table=self.copy_table
                     )
+                    self.command.transfer_pk_sequence_value(
+                        source_table=self.table,
+                        dest_table=self.copy_table,
+                        convert_pk_to_bigint=self.convert_pk_to_bigint,
+                    )
                 self.command.rename_table(
                     table_from=self.table, table_to=self.repacked_name
                 )
@@ -434,6 +439,12 @@ def revert_swap(self) -> None:
                 self.command.swap_pk_sequence_name(
                     first_table=self.table, second_table=self.repacked_name
                 )
+                self.command.transfer_pk_sequence_value(
+                    source_table=self.table,
+                    dest_table=self.repacked_name,
+                    convert_pk_to_bigint=self.convert_pk_to_bigint,
+                )
+
             self.command.rename_table(table_from=self.table, table_to=self.copy_table)
             self.command.rename_table(
                 table_from=self.repacked_name, table_to=self.table

tests/factories.py

@@ -13,6 +13,7 @@ def create_table_for_repacking(
     with_exclusion_constraint: bool = False,
     pk_type: str = "SERIAL",
     pk_name: str = "id",
+    pk_start: int = 1,
     ommit_sequence: bool = False,
     schema: str = "public",
 ) -> None:
@@ -45,7 +46,9 @@ def create_table_for_repacking(
     ):
         # Create a sequence manually.
         seq = f"{table_name}_seq"
-        cur.execute(f"CREATE SEQUENCE {schema}.{seq};")
+        cur.execute(
+            f"CREATE SEQUENCE {schema}.{seq} MINVALUE {pk_start} START WITH {pk_start};"
+        )
         pk_type = f"{pk_type} DEFAULT NEXTVAL('{schema}.{seq}')"
 
     cur.execute(
@@ -191,7 +194,7 @@ def create_table_for_repacking(
     cur.execute(
         dedent(f"""
         INSERT INTO {schema}.referring_table ({table_name}_{pk_name})
-        SELECT generate_series(1, {referring_table_rows});
+        SELECT generate_series({pk_start}, {pk_start + referring_table_rows - 1});
     """)
     )
     cur.execute(
@@ -213,7 +216,7 @@ def create_table_for_repacking(
     cur.execute(
         dedent(f"""
         INSERT INTO {schema}.not_valid_referring_table ({table_name}_{pk_name})
-        SELECT generate_series(1, {referring_table_rows});
+        SELECT generate_series({pk_start}, {pk_start + referring_table_rows - 1});
     """)
     )
 

tests/test_repack.py

@@ -39,6 +39,7 @@ class _TableInfo:
     referring_fks: list[_introspect.ReferringForeignKey]
     constraints: list[_introspect.Constraint]
     pk_seq: str
+    pk_seq_val: int | None
 
 
 def _collect_table_info(
@@ -60,13 +61,18 @@ def _collect_table_info(
             table=table, types=["c", "f", "n", "p", "u", "t", "x"]
         )
         pk_seq = introspector.get_pk_sequence_name(table=table)
+        if pk_seq:
+            pk_seq_val = introspector.get_pk_sequence_value(seq=pk_seq)
+        else:
+            pk_seq_val = None
 
     return _TableInfo(
         oid=oid,
         indexes=indexes,
         referring_fks=referring_fks,
         constraints=constraints,
         pk_seq=pk_seq,
+        pk_seq_val=pk_seq_val,
     )
 
 
@@ -127,6 +133,10 @@ def _assert_repack(
     assert table_before.referring_fks == table_after.referring_fks
     assert table_before.constraints == table_after.constraints
     assert table_before.pk_seq == table_after.pk_seq
+    if table_before.pk_seq_val is None or table_before.pk_seq_val > 0:
+        assert table_before.pk_seq_val == table_after.pk_seq_val
+    else:
+        assert table_after.pk_seq_val == 2**31
 
     # All functions and triggers are removed.
     trigger_info = _get_trigger_info(repack, cur)
@@ -1283,6 +1293,36 @@ def test_repack_full_with_serial_pk(
         )
 
 
+def test_when_table_has_negative_pk_values(
+    connection: _psycopg.Connection,
+) -> None:
+    with _cur.get_cursor(connection, logged=True) as cur:
+        factories.create_table_for_repacking(
+            connection=connection,
+            cur=cur,
+            table_name="to_repack",
+            rows=100,
+            pk_type="integer",
+            pk_start=-200,
+        )
+        table_before = _collect_table_info(table="to_repack", connection=connection)
+        repack = Psycopack(
+            table="to_repack",
+            batch_size=1,
+            conn=connection,
+            cur=cur,
+            convert_pk_to_bigint=True,
+        )
+        repack.full()
+        table_after = _collect_table_info(table="to_repack", connection=connection)
+        _assert_repack(
+            table_before=table_before,
+            table_after=table_after,
+            repack=repack,
+            cur=cur,
+        )
+
+
 def test_when_table_has_large_value_being_inserted(
     connection: _psycopg.Connection,
 ) -> None: