From 9a51968f79d7e9b75892ff834c7d7e329c0145ca Mon Sep 17 00:00:00 2001 From: Marcelo Fernandes Date: Thu, 2 Apr 2026 20:18:14 +1300 Subject: [PATCH] Add support for partitioning by date --- README.md | 8 + src/psycopack/__init__.py | 8 +- src/psycopack/_commands.py | 250 +++++++++++++++++++++++-- src/psycopack/_introspect.py | 57 +++++- src/psycopack/_partition.py | 20 ++ src/psycopack/_repack.py | 99 ++++++++-- tests/factories.py | 9 +- tests/test_repack.py | 345 +++++++++++++++++++++++++++++++++++ 8 files changed, 769 insertions(+), 27 deletions(-) create mode 100644 src/psycopack/_partition.py diff --git a/README.md b/README.md index b7de3bc..e8ccda9 100644 --- a/README.md +++ b/README.md @@ -125,6 +125,14 @@ The following types of tables aren't currently supported: - Tables with deferrable unique constraints. - Referring foreign keys on a different schema than the original table. +### Limitations when partitioning: + +- The CHANGE_LOG strategy must be used when doing partitioning. +- Unique constraints are skipped, as they are not supported on partitioned + tables. The user should add them manually after the process is complete or + after the setup step. +- No support for referring fks yet. + ## Required user permissions (or privileges) Unless the user is a superuser, they may lack certain privileges to run diff --git a/src/psycopack/__init__.py b/src/psycopack/__init__.py index ba8f022..fc38918 100644 --- a/src/psycopack/__init__.py +++ b/src/psycopack/__init__.py @@ -5,6 +5,7 @@ from ._conn import get_db_connection from ._cur import get_cursor from ._introspect import BackfillBatch +from ._partition import DateRangeStrategy, PartitionConfig, PartitionInterval from ._registry import RegistryException, UnexpectedSyncStrategy from ._repack import ( BasePsycopackError, @@ -18,6 +19,7 @@ NoReferencesPrivilege, NoReferringTableOwnership, NotTableOwner, + PartitioningForTableWithReferringFKs, PostBackfillBatchCallback, PrimaryKeyNotFound, Psycopack, @@ -35,6 +37,7 @@ "BackfillBatch", "BasePsycopackError", "CompositePrimaryKey", + "DateRangeStrategy", "DeferrableUniqueConstraint", "FailureDueToLockTimeout", "InheritedTable", @@ -45,11 +48,14 @@ "NoReferencesPrivilege", "NoReferringTableOwnership", "NotTableOwner", + "PartitionConfig", + "PartitionInterval", + "PartitioningForTableWithReferringFKs", "PostBackfillBatchCallback", "PrimaryKeyNotFound", + "Psycopack", "ReferringForeignKeyInDifferentSchema", "RegistryException", - "Psycopack", "Stage", "SyncStrategy", "TableDoesNotExist", diff --git a/src/psycopack/_commands.py b/src/psycopack/_commands.py index 0432777..74de1dd 100644 --- a/src/psycopack/_commands.py +++ b/src/psycopack/_commands.py @@ -4,7 +4,7 @@ from textwrap import dedent from typing import Iterator -from . import _cur, _introspect +from . import _cur, _introspect, _partition from . import _psycopg as psycopg @@ -16,11 +16,13 @@ def __init__( cur: _cur.Cursor, introspector: _introspect.Introspector, schema: str, + partition_config: _partition.PartitionConfig | None = None, ) -> None: self.conn = conn self.cur = cur self.introspector = introspector self.schema = schema + self.partition_config = partition_config def drop_constraint(self, *, table: str, constraint: str) -> None: self.cur.execute( @@ -49,6 +51,17 @@ def drop_table_if_exists(self, *, table: str) -> None: ) def create_copy_table(self, *, base_table: str, copy_table: str) -> None: + if self.partition_config: + return self._create_partitioned_table( + base_table=base_table, copy_table=copy_table + ) + + # Create a non-partitioned table (default behavior) + self._create_non_partitioned_table(base_table=base_table, copy_table=copy_table) + + def _create_non_partitioned_table( + self, *, base_table: str, copy_table: str + ) -> None: self.cur.execute( psycopg.sql.SQL( dedent(""" @@ -64,6 +77,189 @@ def create_copy_table(self, *, base_table: str, copy_table: str) -> None: .as_string(self.conn) ) + def _create_partitioned_table(self, *, base_table: str, copy_table: str) -> None: + assert self.partition_config is not None + assert isinstance(self.partition_config.strategy, _partition.DateRangeStrategy) + + # Create the parent partitioned table + self.cur.execute( + psycopg.sql.SQL( + dedent(""" + CREATE TABLE {schema}.{copy_table} + (LIKE {schema}.{table} INCLUDING DEFAULTS) + PARTITION BY RANGE ({partition_column}); + """) + ) + .format( + table=psycopg.sql.Identifier(base_table), + copy_table=psycopg.sql.Identifier(copy_table), + schema=psycopg.sql.Identifier(self.schema), + partition_column=psycopg.sql.Identifier(self.partition_config.column), + ) + .as_string(self.conn) + ) + + # Create partitions ahead of time + self._create_partitions(base_table=base_table, copy_table=copy_table) + + def _create_partitions(self, *, base_table: str, copy_table: str) -> None: + assert self.partition_config is not None + strategy = self.partition_config.strategy + assert isinstance(strategy, _partition.DateRangeStrategy) + + num_of_extra_partitions = self.partition_config.num_of_extra_partitions_ahead + + min_value = self.introspector.get_min_partition_date_value( + table=base_table, column=self.partition_config.column + ) + max_value = self.introspector.get_max_partition_date_value( + table=base_table, column=self.partition_config.column + ) + partition_start = self._get_first_partition_start_date( + min_value=min_value, strategy=strategy + ) + partition_end = self._get_last_partition_end_date( + max_value=max_value, + strategy=strategy, + num_of_extra_partitions=num_of_extra_partitions, + ) + + # Create partitions from partition_start to partition_end + current_partition_start = partition_start + + while current_partition_start < partition_end: + partition_suffix = self._get_partition_suffix( + current_partition_start=current_partition_start, strategy=strategy + ) + + current_partition_end = self._get_partition_end_boundary( + current_partition_start=current_partition_start, strategy=strategy + ) + self._create_datetime_partition( + base_table=base_table, + copy_table=copy_table, + partition_suffix=partition_suffix, + start=current_partition_start, + end=current_partition_end, + ) + current_partition_start = current_partition_end + + def _get_first_partition_start_date( + self, *, min_value: datetime.date, strategy: _partition.DateRangeStrategy + ) -> datetime.date: + """ + Align the minimum value to partition boundaries. + For DAY: uses the exact min_value + For MONTH: aligns to the first day of the month + """ + if strategy.partition_by == _partition.PartitionInterval.DAY: + return min_value + 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}") + + def _get_last_partition_end_date( + self, + *, + max_value: datetime.date, + strategy: _partition.DateRangeStrategy, + num_of_extra_partitions: int, + ) -> datetime.date: + """ + Calculate the end date for partitioning. + Always adds 1 interval unit to cover max_value (since range partitions are + exclusive on the upper bound), plus num_of_extra_partitions. + For DAY: adds 1 + num_of_extra_partitions days + For MONTH: adds 1 + num_of_extra_partitions months + """ + if strategy.partition_by == _partition.PartitionInterval.DAY: + # Add 1 day to ensure max_value is covered, plus extra partitions + return max_value + datetime.timedelta(days=1 + num_of_extra_partitions) + elif strategy.partition_by == _partition.PartitionInterval.MONTH: + # Add 1 month to ensure max_value is covered, plus extra partitions + # Add months by advancing to first of month and adding 32*months, + # then normalising. This is because timedelta doesn't accept + # "months" as argument. + temp_date = max_value.replace(day=1) + 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}") + + def _get_partition_end_boundary( + self, + *, + current_partition_start: datetime.date, + strategy: _partition.DateRangeStrategy, + ) -> datetime.date: + """ + Calculate the end boundary for a single partition. + For DAY: adds 1 day + For MONTH: advances to the first day of the next month + """ + if strategy.partition_by == _partition.PartitionInterval.DAY: + return current_partition_start + datetime.timedelta(days=1) + elif strategy.partition_by == _partition.PartitionInterval.MONTH: + # Next month boundary + return (current_partition_start + datetime.timedelta(days=32)).replace( + day=1 + ) + else: + raise ValueError(f"Unsupported partition_by: {strategy.partition_by}") + + def _get_partition_suffix( + self, + *, + current_partition_start: datetime.date, + strategy: _partition.DateRangeStrategy, + ) -> str: + """ + Generate a date-based partition suffix. + For DAY: returns p20250101 (YYYYMMDD format) + For MONTH: returns p202501 (YYYYMM format) + """ + if strategy.partition_by == _partition.PartitionInterval.DAY: + # Format: p20250101 (YYYYMMDD) + return f"p{current_partition_start.strftime('%Y%m%d')}" + 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}") + + def _create_datetime_partition( + self, + *, + base_table: str, + copy_table: str, + partition_suffix: str, + start: datetime.date, + end: datetime.date, + ) -> None: + """Create a single datetime range partition.""" + self.cur.execute( + psycopg.sql.SQL( + dedent(""" + CREATE TABLE {schema}.{partition_name} + PARTITION OF {schema}.{copy_table} + FOR VALUES FROM ({start}) TO ({end}); + """) + ) + .format( + schema=psycopg.sql.Identifier(self.schema), + partition_name=psycopg.sql.Identifier( + f"{base_table}_{partition_suffix}" + ), + copy_table=psycopg.sql.Identifier(copy_table), + start=psycopg.sql.Literal(start), + end=psycopg.sql.Literal(end), + ) + .as_string(self.conn) + ) + def drop_sequence_if_exists(self, *, seq: str) -> None: self.cur.execute( psycopg.sql.SQL("DROP SEQUENCE IF EXISTS {schema}.{seq};") @@ -109,17 +305,37 @@ def set_table_id_seq(self, *, table: str, seq: str, pk_column: str) -> None: ) def add_pk(self, *, table: str, pk_column: str) -> None: - self.cur.execute( - psycopg.sql.SQL( - "ALTER TABLE {schema}.{table} ADD PRIMARY KEY ({pk_column});" + # For partitioned tables, the PK must include all partitioning columns + if self.partition_config: + pk_columns = psycopg.sql.SQL(", ").join( + [ + psycopg.sql.Identifier(pk_column), + psycopg.sql.Identifier(self.partition_config.column), + ] ) - .format( - table=psycopg.sql.Identifier(table), - pk_column=psycopg.sql.Identifier(pk_column), - schema=psycopg.sql.Identifier(self.schema), + self.cur.execute( + psycopg.sql.SQL( + "ALTER TABLE {schema}.{table} ADD PRIMARY KEY ({pk_columns});" + ) + .format( + table=psycopg.sql.Identifier(table), + pk_columns=pk_columns, + schema=psycopg.sql.Identifier(self.schema), + ) + .as_string(self.conn) + ) + else: + self.cur.execute( + psycopg.sql.SQL( + "ALTER TABLE {schema}.{table} ADD PRIMARY KEY ({pk_column});" + ) + .format( + table=psycopg.sql.Identifier(table), + pk_column=psycopg.sql.Identifier(pk_column), + schema=psycopg.sql.Identifier(self.schema), + ) + .as_string(self.conn) ) - .as_string(self.conn) - ) def create_copy_function( self, @@ -511,12 +727,24 @@ def create_unique_constraint_using_idx( def create_not_valid_constraint_from_def( self, *, table: str, constraint: str, definition: str, is_validated: bool ) -> None: + # For partitioned tables, we can't use NOT VALID on foreign keys + # So we need to remove it from the definition + is_fk = "FOREIGN KEY" in definition.upper() + if self.partition_config and is_fk and not is_validated: + # Remove NOT VALID from the definition for partitioned tables + definition = definition.replace(" NOT VALID", "").replace("NOT VALID", "") + add_constraint_sql = dedent(""" ALTER TABLE {schema}.{table} ADD CONSTRAINT {constraint} {definition} """) - if is_validated: + # Only add NOT VALID if: + # 1. The constraint is validated (so we make it NOT VALID temporarily) + # 2. AND it's not a FK on a partitioned table (which doesn't support NOT VALID) + should_add_not_valid = is_validated and not (self.partition_config and is_fk) + + if should_add_not_valid: # If the definition is for a valid constraint, alter it to be not # valid manually so that it can be created ONLINE. add_constraint_sql += " NOT VALID" diff --git a/src/psycopack/_introspect.py b/src/psycopack/_introspect.py index 9128794..a030c34 100644 --- a/src/psycopack/_introspect.py +++ b/src/psycopack/_introspect.py @@ -1,4 +1,5 @@ import dataclasses +import datetime from textwrap import dedent from . import _const, _cur @@ -539,7 +540,9 @@ def get_primary_key_info(self, *, table: str) -> PrimaryKey | None: def get_pk_sequence_name(self, *, table: str) -> str: pk_info = self.get_primary_key_info(table=table) assert pk_info is not None - assert len(pk_info.columns) == 1 + # Return empty string for composite PKs (e.g., partitioned tables) + if len(pk_info.columns) != 1: + return "" if pk_info.identity_type: self.cur.execute( @@ -761,3 +764,55 @@ def is_table_owner(self, *, table: str, schema: str) -> bool: result = self.cur.fetchone() 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. + If the table is empty or the column is NULL, returns the current date. + """ + self.cur.execute( + psycopg.sql.SQL( + "SELECT COALESCE(MIN({column})::DATE, CURRENT_DATE) FROM {schema}.{table};" + ) + .format( + column=psycopg.sql.Identifier(column), + schema=psycopg.sql.Identifier(self.schema), + table=psycopg.sql.Identifier(table), + ) + .as_string(self.conn) + ) + result = self.cur.fetchone() + assert result is not None + min_value = result[0] + assert isinstance(min_value, datetime.date) + return min_value + + def get_max_partition_date_value(self, *, table: str, column: str) -> datetime.date: + """ + Get the maximum value of the partition column from the table. + If the table is empty or the column is NULL, returns the current date. + """ + self.cur.execute( + psycopg.sql.SQL( + "SELECT COALESCE(MAX({column})::DATE, CURRENT_DATE) FROM {schema}.{table};" + ) + .format( + column=psycopg.sql.Identifier(column), + schema=psycopg.sql.Identifier(self.schema), + table=psycopg.sql.Identifier(table), + ) + .as_string(self.conn) + ) + result = self.cur.fetchone() + assert result is not None + max_value = result[0] + assert isinstance(max_value, datetime.date) + return max_value diff --git a/src/psycopack/_partition.py b/src/psycopack/_partition.py new file mode 100644 index 0000000..9130891 --- /dev/null +++ b/src/psycopack/_partition.py @@ -0,0 +1,20 @@ +import dataclasses +import enum + + +class PartitionInterval(enum.Enum): + DAY = "DAY" + MONTH = "MONTH" + + +@dataclasses.dataclass +class DateRangeStrategy: + partition_by: PartitionInterval + + +@dataclasses.dataclass +class PartitionConfig: + column: str + num_of_extra_partitions_ahead: int + # Todo: Add support for other types of partitioning. + strategy: DateRangeStrategy diff --git a/src/psycopack/_repack.py b/src/psycopack/_repack.py index 2366398..7411ac6 100644 --- a/src/psycopack/_repack.py +++ b/src/psycopack/_repack.py @@ -12,6 +12,7 @@ _cur, _identifiers, _introspect, + _partition, _registry, _sync_strategy, _tracker, @@ -71,6 +72,10 @@ class DeferrableUniqueConstraint(BasePsycopackError): pass +class PartitioningForTableWithReferringFKs(BasePsycopackError): + pass + + class NoCreateAndUsagePrivilegeOnSchema(BasePsycopackError): pass @@ -149,6 +154,7 @@ def __init__( skip_permissions_check: bool = False, sync_strategy: _sync_strategy.SyncStrategy = _sync_strategy.SyncStrategy.DIRECT_TRIGGER, change_log_batch_size: int | None = None, + partition_config: _partition.PartitionConfig | None = None, ) -> None: self.conn = conn self.cur = cur @@ -162,11 +168,23 @@ def __init__( cur=self.cur, introspector=self.introspector, schema=schema, + partition_config=partition_config, ) self.table = table self.schema = schema self._check_table_exists() + + # Cache the pk_column early to avoid issues when the table is + # partitioned later (which creates a composite PK). We do this + # 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 = "" + 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] + if not skip_permissions_check: self._check_user_permissions() @@ -176,6 +194,7 @@ def __init__( self.convert_pk_to_bigint = convert_pk_to_bigint self.allow_empty = allow_empty self.sync_strategy = sync_strategy + self.partition_config = partition_config # Names for psycopack objects are stored in the Registry self.registry = _registry.Registry( @@ -220,19 +239,24 @@ def __init__( command=self.command, schema=schema, ) - self._pk_column = "" @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 - assert len(pk_info.columns) == 1 + # 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 @@ -382,6 +406,19 @@ def pre_validate(self) -> None: f"{deferrable_unique_constraints}." ) + if self.partition_config: + referring_fks = self.introspector.get_referring_fks(table=self.table) + if referring_fks: + fk_names = [ + f"{fk.schema}.{fk.referring_table}.{fk.name}" + for fk in referring_fks + ] + raise PartitioningForTableWithReferringFKs( + f"Psycopack does not support partitioning tables with " + f"referring foreign keys. The table {self.table} has " + f"the following referring foreign keys: {fk_names}." + ) + def setup_repacking(self) -> None: with ( self.tracker.track(_tracker.Stage.SETUP), @@ -406,11 +443,11 @@ def backfill(self) -> None: def sync_schemas(self) -> None: with self.tracker.track(_tracker.Stage.SYNC_SCHEMAS): - self._create_indexes( - concurrently=bool( - self.sync_strategy != _sync_strategy.SyncStrategy.CHANGE_LOG - ) - ) + # Partitioned tables don't support CONCURRENTLY for index creation + concurrently = ( + self.sync_strategy != _sync_strategy.SyncStrategy.CHANGE_LOG + ) and not self.partition_config + self._create_indexes(concurrently=concurrently) with self.command.lock_timeout(self.lock_timeout): self._create_unique_constraints() self._create_check_and_fk_constraints() @@ -489,7 +526,9 @@ def swap(self) -> None: table=self.table, trigger=self.trigger ) self.command.drop_function_if_exists(function=self.function) - if self.introspector.get_pk_sequence_name(table=self.table): + if self.introspector.get_pk_sequence_name( + table=self.table + ) and self.introspector.get_pk_sequence_name(table=self.copy_table): self.command.swap_pk_sequence_name( first_table=self.table, second_table=self.copy_table ) @@ -546,7 +585,9 @@ def revert_swap(self) -> None: trigger=self.repacked_trigger, ) self.command.drop_function_if_exists(function=self.repacked_function) - if self.introspector.get_pk_sequence_name(table=self.table): + if self.introspector.get_pk_sequence_name( + table=self.table + ) and self.introspector.get_pk_sequence_name(table=self.repacked_name): self.command.swap_pk_sequence_name( first_table=self.table, second_table=self.repacked_name ) @@ -577,8 +618,13 @@ def clean_up(self) -> None: sql = idx.definition sql_arr = sql.split(" ON") sql = sql_arr[1] - idx_data = next(idx for idx in indexes[sql] if "idx_to" not in idx) - idx_data["idx_to"] = idx.name + # Find matching index. Skip if not found (e.g., partitioned table with different PK) + matching_idx = next( + (idx for idx in indexes.get(sql, []) if "idx_to" not in idx), + None, + ) + if matching_idx: + matching_idx["idx_to"] = idx.name # Rename foreign keys from other tables using a dict data structure to # hold names from/to. Support multiple FKs from the same referring table. @@ -633,6 +679,9 @@ def clean_up(self) -> None: for idx_sql in indexes: for index_data in indexes[idx_sql]: + # Skip if no matching index was found (e.g., partitioned table with different PK) + if "idx_to" not in index_data: + continue self.command.rename_index( idx_from=index_data["idx_from"], idx_to=_identifiers.build_postgres_identifier( @@ -646,6 +695,9 @@ def clean_up(self) -> None: for fk_def in fk_definitions: for fk_data in fk_definitions[fk_def]: + # Skip if no matching FK was found (e.g., partitioned table) + if "cons_to" not in fk_data: + continue self.command.drop_constraint( table=fk_data["referring_table"], constraint=fk_data["cons_from"], @@ -656,6 +708,15 @@ def clean_up(self) -> None: cons_to=fk_data["cons_from"], ) + # For partitioned tables, drop referring FKs before dropping the table + if self.partition_config: + for fk in self.introspector.get_referring_fks( + table=self.repacked_name + ): + self.command.drop_constraint( + table=fk.referring_table, constraint=fk.name + ) + self.command.drop_table_if_exists(table=self.repacked_name) self.command.drop_table_if_exists(table=self.backfill_log) self.registry.delete_row_for(table=self.table) @@ -923,6 +984,12 @@ def _create_indexes(self, concurrently: bool) -> None: for index in indexes: name = index.name sql = index.definition + + # For partitioned tables, convert UNIQUE indexes to regular indexes + # since unique constraints must include all partitioning columns + if self.partition_config and "UNIQUE" in sql: + sql = sql.replace("CREATE UNIQUE INDEX", "CREATE INDEX") + if concurrently: sql = sql.replace( "CREATE INDEX", "CREATE INDEX CONCURRENTLY IF NOT EXISTS" @@ -955,6 +1022,11 @@ def _create_indexes(self, concurrently: bool) -> None: self.cur.execute(sql) def _create_unique_constraints(self) -> None: + # Skip unique constraints for partitioned tables since they must + # include all partitioning columns + if self.partition_config: + return + table_constraints = self.introspector.get_constraints( table=self.table, types=["u"] ) @@ -1008,6 +1080,11 @@ def _create_check_and_fk_constraints(self) -> None: ) def _create_referring_fks(self) -> None: + # Skip referring FKs for partitioned tables since the PK now includes + # the partition column, making it incompatible with existing FKs + if self.partition_config: + return + table_fks = self.introspector.get_referring_fks(table=self.table) copy_fks = self.introspector.get_referring_fks(table=self.copy_table) diff --git a/tests/factories.py b/tests/factories.py index b871ec8..0085acb 100644 --- a/tests/factories.py +++ b/tests/factories.py @@ -66,7 +66,8 @@ def create_table_for_repacking( not_valid_fk INTEGER, {table_name} INTEGER, var_maybe_with_exclusion VARCHAR(255), - var_with_multiple_idx VARCHAR(10) + var_with_multiple_idx VARCHAR(10), + datetime_field TIMESTAMP ); """) ) @@ -150,7 +151,8 @@ def create_table_for_repacking( not_valid_fk, {table_name}, var_maybe_with_exclusion, - var_with_multiple_idx + var_with_multiple_idx, + datetime_field ) SELECT {"gs," if ommit_sequence else ""} @@ -165,7 +167,8 @@ def create_table_for_repacking( (floor(random() * {referred_table_rows}) + 1)::int, (floor(random() * 10) + 1)::int, substring(md5(random()::text), 1, 10), - substring(md5(random()::text), 1, 10) + substring(md5(random()::text), 1, 10), + TIMESTAMP '2025-01-01' + (RANDOM() * INTERVAL '1 month') FROM generate_series(1, {rows}) AS gs; """) ) diff --git a/tests/test_repack.py b/tests/test_repack.py index 563fd78..30c9079 100644 --- a/tests/test_repack.py +++ b/tests/test_repack.py @@ -18,6 +18,7 @@ NoReferencesPrivilege, NoReferringTableOwnership, NotTableOwner, + PartitioningForTableWithReferringFKs, PrimaryKeyNotFound, Psycopack, ReferringForeignKeyInDifferentSchema, @@ -30,6 +31,7 @@ _const, _cur, _introspect, + _partition, _psycopg, _tracker, ) @@ -2050,6 +2052,46 @@ def test_with_deferrable_unique_constraint(connection: _psycopg.Connection) -> N repack.full() +def test_partition_with_referring_fks(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", + ) + # Create a referring table with FK to to_repack + cur.execute("DROP TABLE IF EXISTS referring_table;") + cur.execute( + dedent(""" + CREATE TABLE referring_table ( + id SERIAL PRIMARY KEY, + to_repack_id INTEGER REFERENCES to_repack(id) + ); + """) + ) + repack = Psycopack( + table="to_repack", + batch_size=1, + conn=connection, + cur=cur, + partition_config=_partition.PartitionConfig( + column="datetime_field", + num_of_extra_partitions_ahead=10, + strategy=_partition.DateRangeStrategy( + partition_by=_partition.PartitionInterval.DAY + ), + ), + ) + + with pytest.raises( + PartitioningForTableWithReferringFKs, + match="referring foreign keys", + ): + repack.full() + + def test_without_schema_privileges(connection: _psycopg.Connection) -> None: schema = "sweet_schema" with _cur.get_cursor(connection, logged=True) as cur: @@ -2735,3 +2777,306 @@ def test_multiple_foreign_keys_from_same_referring_table( assert rows[0] == ("Post 1", "Anna", "Anna") assert rows[1] == ("Post 2", "Anna", "Bob") assert rows[2] == ("Post 3", "Bob", "Carlos") + + +def test_repack_with_day_range_partition(connection: _psycopg.Connection) -> None: + with _cur.get_cursor(connection) as cur: + factories.create_table_for_repacking( + connection=connection, + cur=cur, + table_name="to_repack", + rows=100, + ) + # Clean up any referring tables from other tests + cur.execute("DROP TABLE IF EXISTS referring_table CASCADE;") + cur.execute("DROP TABLE IF EXISTS not_valid_referring_table CASCADE;") + cur.execute("DROP TABLE IF EXISTS other_referring_table CASCADE;") + # Insert a row with a datetime field for the last day of January to + # guarantee all of January will be covered by partitions. + create_row_sql = dedent( + """ + INSERT INTO + to_repack + (datetime_field) + VALUES + ('2025-01-31 23:59:59'); + """ + ) + # Insert a row in the table to_repack, to verify the trigger places the + # pk of the row being changed in the change log. + cur.execute(create_row_sql) + repack = Psycopack( + table="to_repack", + batch_size=1, + conn=connection, + cur=cur, + partition_config=_partition.PartitionConfig( + column="datetime_field", + num_of_extra_partitions_ahead=10, + strategy=_partition.DateRangeStrategy( + partition_by=_partition.PartitionInterval.DAY + ), + ), + ) + repack.full() + + # Verify the table is now partitioned + cur.execute(""" + SELECT + pt.relname AS partition_name + FROM + pg_inherits + JOIN + pg_class pt + ON pt.oid = inhrelid + JOIN + pg_class p + ON p.oid = inhparent + WHERE + p.relname = 'to_repack' + ORDER BY + partition_name; + """) + partitions = cur.fetchall() + assert len(partitions) > 0, "Table should have partitions" + + # Verify partition names follow the date-based format + # Example: to_repack_p20250106 (table_name_pYYYYMMDD for DAY) + # We should have partitions from MIN(datetime_field) to MAX(datetime_field) + 10 extra days + # Since we inserted a row with 2025-01-31, we know that will be included + partition_names = [p[0] for p in partitions] + + # All partitions should follow the naming format: to_repack_pYYYYMMDD + for partition_name in partition_names: + assert partition_name.startswith("to_repack_p"), ( + f"Invalid partition name: {partition_name}" + ) + date_part = partition_name.replace("to_repack_p", "") + assert len(date_part) == 8, ( + f"Date part should be 8 digits (YYYYMMDD): {partition_name}" + ) + assert date_part.isdigit(), f"Date part should be numeric: {partition_name}" + + # Verify we have the partition for the explicitly inserted row (2025-01-31) + assert "to_repack_p20250131" in partition_names, ( + "Should have partition for 2025-01-31" + ) + + # Verify we have 10 extra partitions after the max date (into February) + february_partitions = [ + p for p in partition_names if p.startswith("to_repack_p202502") + ] + assert len(february_partitions) == 10, ( + f"Should have exactly 10 partitions in February, got {len(february_partitions)}" + ) + + # Verify the primary key includes both id and datetime_field (partition column) + cur.execute( + """ + SELECT + a.attname AS column_name + FROM + pg_index i + JOIN + pg_attribute a + ON a.attrelid = i.indrelid + AND a.attnum = ANY(i.indkey) + WHERE + i.indrelid = 'to_repack'::regclass + AND i.indisprimary + ORDER BY + array_position(i.indkey, a.attnum); + """ + ) + pk_columns = [row[0] for row in cur.fetchall()] + assert pk_columns == ["id", "datetime_field"] + + # Verify data integrity + cur.execute("SELECT COUNT(*) FROM to_repack;") + count = cur.fetchone() + assert count is not None + assert count[0] == 101, "All rows should be present (100 + 1 inserted)" + + +def test_repack_with_month_range_partition(connection: _psycopg.Connection) -> None: + with _cur.get_cursor(connection) as cur: + factories.create_table_for_repacking( + connection=connection, + cur=cur, + table_name="to_repack", + rows=100, + ) + # Clean up any referring tables from other tests + cur.execute("DROP TABLE IF EXISTS referring_table CASCADE;") + cur.execute("DROP TABLE IF EXISTS not_valid_referring_table CASCADE;") + cur.execute("DROP TABLE IF EXISTS other_referring_table CASCADE;") + # Insert a row with a datetime field for the last day of June to + # guarantee partitions from Jan to Jun (plus extras) will be created. + create_row_sql = dedent( + """ + INSERT INTO + to_repack + (datetime_field) + VALUES + ('2025-06-30 23:59:59'); + """ + ) + cur.execute(create_row_sql) + repack = Psycopack( + table="to_repack", + batch_size=1, + conn=connection, + cur=cur, + partition_config=_partition.PartitionConfig( + column="datetime_field", + num_of_extra_partitions_ahead=3, + strategy=_partition.DateRangeStrategy( + partition_by=_partition.PartitionInterval.MONTH + ), + ), + ) + repack.full() + + # Verify the table is now partitioned + cur.execute(""" + SELECT + pt.relname AS partition_name + FROM + pg_inherits + JOIN + pg_class pt + ON pt.oid = inhrelid + JOIN + pg_class p + ON p.oid = inhparent + WHERE + p.relname = 'to_repack' + ORDER BY + partition_name; + """) + partitions = cur.fetchall() + assert len(partitions) > 0, "Table should have partitions" + + # Verify partition names follow the date-based format + # Example: to_repack_p202501 (table_name_pYYYYMM for MONTH) + # We should have partitions from MIN(datetime_field) to MAX(datetime_field) + 3 extra months + # Since we inserted a row with 2025-06-30, we know June will be included + partition_names = [p[0] for p in partitions] + assert partition_names == [ + "to_repack_p202501", + "to_repack_p202502", + "to_repack_p202503", + "to_repack_p202504", + "to_repack_p202505", + "to_repack_p202506", + "to_repack_p202507", + "to_repack_p202508", + "to_repack_p202509", + ] + + # Verify the primary key includes both id and datetime_field (partition column) + cur.execute( + """ + SELECT + a.attname AS column_name + FROM + pg_index i + JOIN + pg_attribute a + ON a.attrelid = i.indrelid + AND a.attnum = ANY(i.indkey) + WHERE + i.indrelid = 'to_repack'::regclass + AND i.indisprimary + ORDER BY + array_position(i.indkey, a.attnum); + """ + ) + pk_columns = [row[0] for row in cur.fetchall()] + assert pk_columns == ["id", "datetime_field"] + + # Verify data integrity + cur.execute("SELECT COUNT(*) FROM to_repack;") + count = cur.fetchone() + assert count is not None + assert count[0] == 101, "All rows should be present (100 + 1 inserted)" + + +def test_partition_calling_stages_individually( + connection: _psycopg.Connection, +) -> None: + """ + Test that partitioning works when calling stages individually (e.g., swap() + directly) rather than using full(). This simulates resuming a repack process. + """ + with _cur.get_cursor(connection) as cur: + factories.create_table_for_repacking( + connection=connection, + cur=cur, + table_name="to_repack", + rows=100, + ) + # Clean up any referring tables from other tests + cur.execute("DROP TABLE IF EXISTS referring_table CASCADE;") + cur.execute("DROP TABLE IF EXISTS not_valid_referring_table CASCADE;") + cur.execute("DROP TABLE IF EXISTS other_referring_table CASCADE;") + + # Create the Psycopack instance with partition config + repack = Psycopack( + table="to_repack", + batch_size=1, + conn=connection, + cur=cur, + partition_config=_partition.PartitionConfig( + column="datetime_field", + num_of_extra_partitions_ahead=3, + strategy=_partition.DateRangeStrategy( + partition_by=_partition.PartitionInterval.MONTH + ), + ), + ) + + # Call stages individually + repack.pre_validate() + repack.setup_repacking() + repack.backfill() + repack.sync_schemas() + repack.post_sync_update() + + # Now instantiate a new Psycopack and call swap() directly + # (simulating resuming from a previous run) + repack2 = Psycopack( + table="to_repack", + batch_size=1, + conn=connection, + cur=cur, + partition_config=_partition.PartitionConfig( + column="datetime_field", + num_of_extra_partitions_ahead=3, + strategy=_partition.DateRangeStrategy( + partition_by=_partition.PartitionInterval.MONTH + ), + ), + ) + repack2.swap() + repack2.clean_up() + + # Verify the table is now partitioned + cur.execute(""" + SELECT + pt.relname AS partition_name + FROM + pg_inherits + JOIN + pg_class pt + ON pt.oid = inhrelid + JOIN + pg_class p + ON p.oid = inhparent + WHERE + p.relname = 'to_repack' + ORDER BY + partition_name; + """) + partitions = cur.fetchall() + assert len(partitions) > 0, "Table should have partitions"