Move updating sequence value to clean_up stage

Author: timb07

src/psycopack/_commands.py

@@ -476,20 +476,11 @@ 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:
+    def transfer_pk_sequence_value(self, *, source_table: str, dest_table: str) -> 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(
@@ -500,6 +491,27 @@ def transfer_pk_sequence_value(
             .as_string(self.conn)
         )
 
+    def update_pk_sequence_value(self, *, table: str) -> None:
+        """
+        Update the sequence value if it was negative (for use in bigint conversions).
+        """
+        seq = self.introspector.get_pk_sequence_name(table=table)
+        value = self.introspector.get_pk_sequence_value(seq=seq)
+
+        if 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
+
+            self.cur.execute(
+                psycopg.sql.SQL("SELECT setval('{schema}.{sequence}', {value});")
+                .format(
+                    schema=psycopg.sql.Identifier(self.schema),
+                    sequence=psycopg.sql.Identifier(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/_repack.py

@@ -403,7 +403,6 @@ def swap(self) -> None:
                     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
@@ -461,7 +460,6 @@ def revert_swap(self) -> None:
                 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)
@@ -507,6 +505,11 @@ def clean_up(self) -> None:
                 )
                 self.command.drop_function_if_exists(function=self.repacked_function)
 
+                if self.convert_pk_to_bigint and self.introspector.get_pk_sequence_name(
+                    table=self.table
+                ):
+                    self.command.update_pk_sequence_value(table=self.table)
+
                 for idx_sql in indexes:
                     for index_data in indexes[idx_sql]:
                         self.command.rename_index(

tests/test_repack.py

@@ -3,8 +3,6 @@
 from typing import Tuple, Union
 from unittest import mock
 
-from psycopg.errors import NumericValueOutOfRange
-
 import pytest
 
 from psycopack import (
@@ -1399,15 +1397,6 @@ def test_when_table_has_negative_pk_values(
         )
 
 
-@pytest.mark.xfail(
-    raises=NumericValueOutOfRange,
-    reason="""
-E           psycopg.errors.NumericValueOutOfRange: integer out of range
-E           CONTEXT:  SQL function "psycopack_repacked_6723301_fun" statement 1
-E           SQL statement "SELECT "public"."psycopack_repacked_6723301_fun"(NEW."id", NEW."id")"
-E           PL/pgSQL function psycopack_repacked_6723301_tgr() line 4 at PERFORM
-""",
-)
 def test_with_writes_when_table_has_negative_pk_values(
     connection: _psycopg.Connection,
 ) -> None: