diff --git a/src/psycopack/_commands.py b/src/psycopack/_commands.py index 74de1dd..00259ba 100644 --- a/src/psycopack/_commands.py +++ b/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, diff --git a/src/psycopack/_introspect.py b/src/psycopack/_introspect.py index a030c34..7d88659 100644 --- a/src/psycopack/_introspect.py +++ b/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. diff --git a/src/psycopack/_repack.py b/src/psycopack/_repack.py index 7411ac6..a23c87f 100644 --- a/src/psycopack/_repack.py +++ b/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: diff --git a/src/psycopack/_tracker.py b/src/psycopack/_tracker.py index 020f432..9d87e9a 100644 --- a/src/psycopack/_tracker.py +++ b/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: diff --git a/tests/test_repack.py b/tests/test_repack.py index 30c9079..84c67a9 100644 --- a/tests/test_repack.py +++ b/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()