Improve coverage

src/psycopack/_commands.py

@@ -157,8 +157,8 @@ def _get_first_partition_start_date(
         elif strategy.partition_by == _partition.PartitionInterval.MONTH:
             # Align to start of month
             return min_value.replace(day=1)
-        else:
-            raise ValueError(f"Unsupported partition_by: {strategy.partition_by}")
+        else:  # pragma: no cover
+            raise NotImplementedError
 
     def _get_last_partition_end_date(
         self,
@@ -186,8 +186,8 @@ def _get_last_partition_end_date(
             for _ in range(1 + num_of_extra_partitions):
                 temp_date = (temp_date + datetime.timedelta(days=32)).replace(day=1)
             return temp_date
-        else:
-            raise ValueError(f"Unsupported partition_by: {strategy.partition_by}")
+        else:  # pragma: no cover
+            raise NotImplementedError
 
     def _get_partition_end_boundary(
         self,
@@ -207,8 +207,8 @@ def _get_partition_end_boundary(
             return (current_partition_start + datetime.timedelta(days=32)).replace(
                 day=1
             )
-        else:
-            raise ValueError(f"Unsupported partition_by: {strategy.partition_by}")
+        else:  # pragma: no cover
+            raise NotImplementedError
 
     def _get_partition_suffix(
         self,
@@ -227,8 +227,8 @@ def _get_partition_suffix(
         elif strategy.partition_by == _partition.PartitionInterval.MONTH:
             # Format: p202501 (YYYYMM)
             return f"p{current_partition_start.strftime('%Y%m')}"
-        else:
-            raise ValueError(f"Unsupported partition_by: {strategy.partition_by}")
+        else:  # pragma: no cover
+            raise NotImplementedError
 
     def _create_datetime_partition(
         self,

src/psycopack/_introspect.py

@@ -765,14 +765,6 @@ def is_table_owner(self, *, table: str, schema: str) -> bool:
         assert result is not None
         return bool(result[0])
 
-    def get_current_date(self) -> datetime.date:
-        self.cur.execute("SELECT CURRENT_DATE;")
-        result = self.cur.fetchone()
-        assert result is not None
-        current_date = result[0]
-        assert isinstance(current_date, datetime.date)
-        return current_date
-
     def get_min_partition_date_value(self, *, table: str, column: str) -> datetime.date:
         """
         Get the minimum value of the partition column from the table.

src/psycopack/_repack.py

@@ -180,10 +180,10 @@ def __init__(
         # after checking the table exists so we can safely introspect it.
         # The original table always has a single-column PK; partitioning
         # only affects the copy table until swap() is called.
-        self._pk_column = ""
+        self.pk_column = ""
         pk_info = self.introspector.get_primary_key_info(table=self.table)
         if pk_info and len(pk_info.columns) == 1:
-            self._pk_column = pk_info.columns[0]
+            self.pk_column = pk_info.columns[0]
 
         if not skip_permissions_check:
             self._check_user_permissions()
@@ -240,26 +240,6 @@ def __init__(
             schema=schema,
         )
 
-    @property
-    def pk_column(self) -> str:
-        """
-        Method to cache the name of the pk column in the instance as to avoid
-        calling introspection queries multiple times.
-
-        When partitioning is enabled, the table's PK becomes composite after
-        setup_repacking(). In that case, we return the first column which is
-        always the original single-column PK.
-        """
-        if self._pk_column:
-            return self._pk_column
-        pk_info = self.introspector.get_primary_key_info(table=self.table)
-        assert pk_info is not None
-        # For partitioned tables, PK becomes composite (original PK + partition column)
-        # We always want the first column, which is the original PK
-        assert len(pk_info.columns) >= 1
-        self._pk_column = pk_info.columns[0]
-        return self._pk_column
-
     def full(self) -> None:
         """
         Process a full table repack from beginning to end.
@@ -483,7 +463,7 @@ def post_sync_update(self) -> None:
                 return
             elif self.sync_strategy == _sync_strategy.SyncStrategy.CHANGE_LOG:
                 return self._post_sync_update_for_change_log()
-            else:
+            else:  # pragma: no cover
                 raise NotImplementedError
 
     def _post_sync_update_for_change_log(self) -> None:

src/psycopack/_tracker.py

@@ -201,7 +201,7 @@ def _validate_stage_dependencies(self, *, stage: Stage) -> None:
                 Stage.SWAP: [self.table, self.copy_table],
                 Stage.CLEAN_UP: [self.repacked_name, self.table],
             }
-        else:
+        else:  # pragma: no cover
             raise NotImplementedError
 
         if stage in table_dependencies:
@@ -223,7 +223,7 @@ def _validate_stage_dependencies(self, *, stage: Stage) -> None:
                 Stage.SWAP: [self.trigger],
                 Stage.CLEAN_UP: [self.repacked_trigger],
             }
-        else:
+        else:  # pragma: no cover
             raise NotImplementedError
 
         if stage in trigger_dependencies:

tests/test_repack.py

@@ -3080,3 +3080,161 @@ def test_partition_calling_stages_individually(
         """)
         partitions = cur.fetchall()
         assert len(partitions) > 0, "Table should have partitions"
+
+
+def test_pk_column_property_accessed_before_setup(
+    connection: _psycopg.Connection,
+) -> None:
+    """
+    Test pk_column property when accessed before setup_repacking.
+    """
+    with _cur.get_cursor(connection, logged=True) as cur:
+        # Create a simple table
+        cur.execute(
+            """
+            CREATE TABLE test_pk_early_access (
+                id bigint PRIMARY KEY,
+                data text
+            );
+            """
+        )
+        cur.execute("INSERT INTO test_pk_early_access (id, data) VALUES (1, 'test');")
+        connection.commit()
+
+        try:
+            # Create Psycopack instance
+            repack = Psycopack(
+                table="test_pk_early_access",
+                batch_size=1,
+                conn=connection,
+                cur=cur,
+            )
+            # Access pk_column BEFORE calling setup_repacking
+            pk_col = repack.pk_column
+            assert pk_col == "id"
+            # Verify _pk_column was cached
+            assert repack.pk_column == "id"
+        finally:
+            connection.rollback()
+            cur.execute("DROP TABLE IF EXISTS test_pk_early_access CASCADE;")
+            connection.commit()
+
+
+def test_pk_column_property_for_partitioned_table(
+    connection: _psycopg.Connection,
+) -> None:
+    """
+    Test pk_column property when _pk_column is None for partitioned tables.
+
+    This tests the code path where pk_column is accessed before setup_repacking
+    and handles composite primary keys (original PK + partition column).
+    """
+    with _cur.get_cursor(connection, logged=True) as cur:
+        # Create a simple table with date column for partitioning
+        cur.execute(
+            """
+            CREATE TABLE test_pk_column (
+                id bigint PRIMARY KEY,
+                datetime_field date NOT NULL,
+                data text
+            );
+            """
+        )
+        cur.execute(
+            "INSERT INTO test_pk_column (id, datetime_field, data) "
+            "VALUES (1, '2025-01-01', 'test');"
+        )
+        connection.commit()
+        try:
+            # Create Psycopack instance with partitioning
+            repack = Psycopack(
+                table="test_pk_column",
+                batch_size=1,
+                conn=connection,
+                cur=cur,
+                partition_config=_partition.PartitionConfig(
+                    column="datetime_field",
+                    num_of_extra_partitions_ahead=1,
+                    strategy=_partition.DateRangeStrategy(
+                        partition_by=_partition.PartitionInterval.DAY
+                    ),
+                ),
+                sync_strategy=SyncStrategy.CHANGE_LOG,
+            )
+
+            # Validate and setup to create the partitioned table
+            repack.pre_validate()
+            repack.setup_repacking()
+
+            # Now test pk_column property - for partitioned tables,
+            # the PK becomes composite (original PK + partition column)
+            # but pk_column should return the first column (original PK)
+            pk_col = repack.pk_column
+            assert pk_col == "id"
+
+            # Verify the copy table was created as partitioned
+            cur.execute(
+                f"""
+                SELECT relkind FROM pg_class
+                WHERE relname = '{repack.copy_table}'
+                AND relnamespace = 'public'::regnamespace;
+                """
+            )
+            result = cur.fetchone()
+            assert result is not None
+            assert result[0] == "p"  # 'p' means partitioned table
+
+        finally:
+            connection.rollback()
+            cur.execute("DROP TABLE IF EXISTS test_pk_column CASCADE;")
+            connection.commit()
+
+
+def test_full_method_resume_from_swap_stage(
+    connection: _psycopg.Connection,
+) -> None:
+    """
+    Test the full() method when resuming from SWAP stage.
+    """
+    with _cur.get_cursor(connection, logged=True) as cur:
+        # Create a test table
+        table_name = "test_resume_swap"
+        cur.execute(f"CREATE TABLE {table_name} (id bigint PRIMARY KEY, data text);")
+        cur.execute(f"INSERT INTO {table_name} (id, data) VALUES (1, 'test');")
+        connection.commit()
+
+        try:
+            repack = Psycopack(
+                table=table_name,
+                batch_size=100,
+                conn=connection,
+                cur=cur,
+            )
+
+            # Run through to POST_SYNC_UPDATE stage (but not SWAP)
+            repack.pre_validate()
+            repack.setup_repacking()
+            repack.backfill()
+            repack.sync_schemas()
+            repack.post_sync_update()
+
+            # Verify we're at SWAP stage (next to be completed)
+            current_stage = repack.tracker.get_current_stage()
+            assert current_stage == _tracker.Stage.SWAP
+
+            # Now call full() - it should run swap and clean_up
+            repack.full()
+
+            # Verify we've completed everything - after clean_up, the tracker row
+            # is deleted, so we're back at PRE_VALIDATION stage
+            current_stage = repack.tracker.get_current_stage()
+            assert current_stage == _tracker.Stage.PRE_VALIDATION
+
+        finally:
+            connection.rollback()
+            # Clean up any remaining tables
+            cur.execute(f"DROP TABLE IF EXISTS {table_name} CASCADE;")
+            cur.execute(
+                f"DROP TABLE IF EXISTS {table_name}_psycopack_repacked CASCADE;"
+            )
+            connection.commit()