diff --git a/src/psycopack/_commands.py b/src/psycopack/_commands.py index 9db675e..8000473 100644 --- a/src/psycopack/_commands.py +++ b/src/psycopack/_commands.py @@ -24,7 +24,10 @@ def __init__( self.schema = schema self.partition_config = partition_config - def drop_constraint(self, *, table: str, constraint: str) -> None: + def drop_constraint( + self, *, table: str, constraint: str, schema: str | None = None + ) -> None: + schema = schema or self.schema self.cur.execute( psycopg.sql.SQL( dedent(""" @@ -35,7 +38,7 @@ def drop_constraint(self, *, table: str, constraint: str) -> None: .format( table=psycopg.sql.Identifier(table), constraint=psycopg.sql.Identifier(constraint), - schema=psycopg.sql.Identifier(self.schema), + schema=psycopg.sql.Identifier(schema), ) .as_string(self.conn) ) @@ -859,28 +862,21 @@ def create_unique_constraint_using_idx( ) def create_not_valid_constraint_from_def( - self, *, table: str, constraint: str, definition: str, is_validated: bool + self, + *, + table: str, + constraint: str, + definition: str, + is_validated: bool, + schema: str | None = None, ) -> 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", "") - + schema = schema or self.schema add_constraint_sql = dedent(""" ALTER TABLE {schema}.{table} ADD CONSTRAINT {constraint} {definition} """) - # 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. + if is_validated: add_constraint_sql += " NOT VALID" self.cur.execute( psycopg.sql.SQL(add_constraint_sql) @@ -888,12 +884,15 @@ def create_not_valid_constraint_from_def( table=psycopg.sql.Identifier(table), constraint=psycopg.sql.Identifier(constraint), definition=psycopg.sql.SQL(definition), - schema=psycopg.sql.Identifier(self.schema), + schema=psycopg.sql.Identifier(schema), ) .as_string(self.conn) ) - def validate_constraint(self, *, table: str, constraint: str) -> None: + def validate_constraint( + self, *, table: str, constraint: str, schema: str | None = None + ) -> None: + schema = schema or self.schema self.cur.execute( psycopg.sql.SQL( dedent(""" @@ -904,7 +903,7 @@ def validate_constraint(self, *, table: str, constraint: str) -> None: .format( table=psycopg.sql.Identifier(table), constraint=psycopg.sql.Identifier(constraint), - schema=psycopg.sql.Identifier(self.schema), + schema=psycopg.sql.Identifier(schema), ) .as_string(self.conn) ) @@ -953,7 +952,15 @@ def rename_sequence(self, *, seq_from: str, seq_to: str) -> None: .as_string(self.conn) ) - def rename_constraint(self, *, table: str, cons_from: str, cons_to: str) -> None: + def rename_constraint( + self, + *, + table: str, + cons_from: str, + cons_to: str, + schema: str | None = None, + ) -> None: + schema = schema or self.schema self.cur.execute( psycopg.sql.SQL( "ALTER TABLE {schema}.{table} RENAME CONSTRAINT {cons_from} TO {cons_to};" @@ -962,7 +969,7 @@ def rename_constraint(self, *, table: str, cons_from: str, cons_to: str) -> None table=psycopg.sql.Identifier(table), cons_from=psycopg.sql.Identifier(cons_from), cons_to=psycopg.sql.Identifier(cons_to), - schema=psycopg.sql.Identifier(self.schema), + schema=psycopg.sql.Identifier(schema), ) .as_string(self.conn) ) diff --git a/src/psycopack/_repack.py b/src/psycopack/_repack.py index a23c87f..3994de3 100644 --- a/src/psycopack/_repack.py +++ b/src/psycopack/_repack.py @@ -98,6 +98,14 @@ def __call__( ) -> None: ... # pragma: no cover +class _ForeignKeyCleanupPair(typing.TypedDict): + schema: str + referring_table: str + cons_from: str + cons_to: str + old_is_validated: bool + new_is_validated: bool + class Psycopack: """ Class for operating a full table copy-and-swap (Psycopack) process. @@ -358,6 +366,7 @@ def pre_validate(self) -> None: f"before proceeding with Psycopack: {invalid_indexes}." ) + """ fks_in_different_schema = [ f"{fk.schema}.{fk.name}" for fk in self.introspector.get_referring_fks(table=self.table) @@ -370,7 +379,7 @@ def pre_validate(self) -> None: f"has the following referring foreign keys in a different " f"schema: {fks_in_different_schema}." ) - + """ deferrable_unique_constraints = [ c.name for c in self.introspector.get_constraints( @@ -606,43 +615,47 @@ def clean_up(self) -> 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. - # Match FKs by their definition (columns), similar to how indexes are matched. - fk_definitions: dict[str, list[dict[str, str]]] = defaultdict(list) - - for fk in self.introspector.get_referring_fks(table=self.repacked_name): - # Normalize the FK definition by replacing the old table name with the new one - # This allows us to match FKs based on their structure. - fk_def = fk.definition.replace( - f"REFERENCES {self.schema}.{self.repacked_name}", - f"REFERENCES {self.schema}.{self.table}", - ).replace( - f"REFERENCES {self.repacked_name}", - f"REFERENCES {self.table}", + # Rename foreign keys from other tables by matching each original FK on the + # repacked table with the swapped-in replacement FK on the new table. + # + # The replacement FK is created in _create_referring_fks() with a deterministic + # temporary name: + # build_postgres_identifier([old_fk.name], "psycopack") + # + # Matching by name is more reliable than matching by pg_get_constraintdef() + # text in retry / reentrant scenarios. + fk_pairs: list[_ForeignKeyCleanupPair] = [] + + new_fks_by_location_and_name: dict[ + tuple[str, str, str], _introspect.ReferringForeignKey + ] = {} + + for fk in self.introspector.get_referring_fks(table=self.table): + new_fks_by_location_and_name[(fk.schema, fk.referring_table, fk.name)] = fk + + for old_fk in self.introspector.get_referring_fks(table=self.repacked_name): + expected_new_name = _identifiers.build_postgres_identifier( + [old_fk.name], "psycopack" + ) + new_fk = new_fks_by_location_and_name.get( + (old_fk.schema, old_fk.referring_table, expected_new_name) ) - fk_definitions[fk_def].append( + assert new_fk is not None, ( + f"Could not find swapped-in FK {expected_new_name} on " + f"{old_fk.schema}.{old_fk.referring_table} for original FK " + f"{old_fk.name}." + ) + fk_pairs.append( { - "cons_from": fk.name, - "referring_table": fk.referring_table, + "schema": old_fk.schema, + "referring_table": old_fk.referring_table, + "cons_from": old_fk.name, + "cons_to": new_fk.name, + "old_is_validated": old_fk.is_validated, + "new_is_validated": new_fk.is_validated, } ) - for fk in self.introspector.get_referring_fks(table=self.table): - # Use the FK definition as the key to match with the old FK. - fk_def = fk.definition - if fk_def in fk_definitions: - # Find the first unmatched FK with this definition. - fk_pair = next( - ( - pair - for pair in fk_definitions[fk_def] - if "cons_to" not in pair - ), - ) - if fk_pair: - fk_pair["cons_to"] = fk.name - with ( self.command.db_transaction(), self.command.lock_timeout(self.lock_timeout), @@ -673,19 +686,23 @@ def clean_up(self) -> None: idx_to=index_data["idx_from"], ) - 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( + for fk_data in fk_pairs: + self.command.drop_constraint( + table=fk_data["referring_table"], + constraint=fk_data["cons_from"], + schema=fk_data["schema"], + ) + self.command.rename_constraint( + table=fk_data["referring_table"], + cons_from=fk_data["cons_to"], + cons_to=fk_data["cons_from"], + schema=fk_data["schema"], + ) + if fk_data["old_is_validated"] and not fk_data["new_is_validated"]: + self.command.validate_constraint( table=fk_data["referring_table"], constraint=fk_data["cons_from"], - ) - self.command.rename_constraint( - table=fk_data["referring_table"], - cons_from=fk_data["cons_to"], - cons_to=fk_data["cons_from"], + schema=fk_data["schema"], ) # For partitioned tables, drop referring FKs before dropping the table @@ -727,7 +744,9 @@ def reset(self) -> None: fks = self.introspector.get_referring_fks(table=self.copy_table) for fk in fks: self.command.drop_constraint( - table=fk.referring_table, constraint=fk.name + table=fk.referring_table, + constraint=fk.name, + schema=fk.schema, ) self.command.drop_trigger_if_exists(table=self.table, trigger=self.trigger) self.command.drop_function_if_exists(function=self.function) @@ -755,7 +774,9 @@ def _create_copy_table(self) -> None: if self.introspector.get_table_oid(table=self.copy_table) is not None: for fk in self.introspector.get_referring_fks(table=self.copy_table): self.command.drop_constraint( - table=fk.referring_table, constraint=fk.name + table=fk.referring_table, + constraint=fk.name, + schema=fk.schema, ) self.command.drop_table_if_exists(table=self.copy_table) @@ -1094,11 +1115,14 @@ def _create_referring_fks(self) -> None: constraint=constraint_name, definition=definition, is_validated=fk.is_validated, + schema=fk.schema, ) - if fk.is_validated: - self.command.validate_constraint( - table=fk.referring_table, constraint=constraint_name - ) + if fk.is_validated: + self.command.validate_constraint( + table=fk.referring_table, + constraint=constraint_name, + schema=fk.schema, + ) def _check_user_permissions(self) -> None: if not self.introspector.has_create_and_usage_privilege_on_schema(): diff --git a/tests/test_repack.py b/tests/test_repack.py index 1daf0ae..b76855d 100644 --- a/tests/test_repack.py +++ b/tests/test_repack.py @@ -21,7 +21,6 @@ PartitioningForTableWithReferringFKs, PrimaryKeyNotFound, Psycopack, - ReferringForeignKeyInDifferentSchema, SyncStrategy, TableDoesNotExist, TableHasTriggers, @@ -1991,7 +1990,14 @@ def test_with_non_default_schema(connection: _psycopg.Connection) -> None: ) -def test_with_fks_from_another_schema(connection: _psycopg.Connection) -> None: +@pytest.mark.parametrize( + "sync_strategy", + [SyncStrategy.DIRECT_TRIGGER, SyncStrategy.CHANGE_LOG], +) +def test_repack_with_referring_fks_from_another_schema( + connection: _psycopg.Connection, + sync_strategy: SyncStrategy, +) -> None: with _cur.get_cursor(connection, logged=True) as cur: factories.create_table_for_repacking( connection=connection, @@ -2000,25 +2006,186 @@ def test_with_fks_from_another_schema(connection: _psycopg.Connection) -> None: rows=10, pk_type="integer", ) + + cur.execute("CREATE SCHEMA sweet_schema;") + + cur.execute( + dedent(""" + CREATE TABLE sweet_schema.referring_table ( + id SERIAL PRIMARY KEY, + to_repack_id INTEGER + ); + """) + ) + cur.execute( + dedent(""" + ALTER TABLE sweet_schema.referring_table + ADD CONSTRAINT cross_schema_referring + FOREIGN KEY (to_repack_id) + REFERENCES public.to_repack(id); + """) + ) + cur.execute( + dedent(""" + INSERT INTO sweet_schema.referring_table (to_repack_id) + SELECT generate_series(1, 10); + """) + ) + + cur.execute( + dedent(""" + CREATE TABLE sweet_schema.not_valid_referring_table ( + id SERIAL PRIMARY KEY, + to_repack_id INTEGER + ); + """) + ) + cur.execute( + dedent(""" + ALTER TABLE sweet_schema.not_valid_referring_table + ADD CONSTRAINT cross_schema_not_valid_referring + FOREIGN KEY (to_repack_id) + REFERENCES public.to_repack(id) + NOT VALID; + """) + ) + cur.execute( + dedent(""" + INSERT INTO sweet_schema.not_valid_referring_table (to_repack_id) + SELECT generate_series(1, 10); + """) + ) + + table_before = _collect_table_info(table="to_repack", connection=connection) + repack = Psycopack( table="to_repack", batch_size=1, conn=connection, cur=cur, + sync_strategy=sync_strategy, + change_log_batch_size=10 + if sync_strategy == SyncStrategy.CHANGE_LOG + else None, ) - schema = "sweet_schema" - cur.execute(f"CREATE SCHEMA {schema};") + 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, + ) + + introspector = _introspect.Introspector( + conn=connection, + cur=cur, + schema="public", + ) + fks_after = introspector.get_referring_fks(table="to_repack") + + assert { + (fk.schema, fk.referring_table, fk.name, fk.is_validated) + for fk in fks_after + } == { + ("public", "referring_table", "referring_table_to_repack_id_fkey", True), + ("public", "not_valid_referring_table", "not_valid_referring", False), + ("sweet_schema", "referring_table", "cross_schema_referring", True), + ( + "sweet_schema", + "not_valid_referring_table", + "cross_schema_not_valid_referring", + False, + ), + } + cur.execute( - dedent(f""" - CREATE TABLE {schema}.referring_table ( - to_repack_id INTEGER REFERENCES public.to_repack(id) - ); + dedent(""" + SELECT + referring_ns.nspname AS referring_schema, + referring_cls.relname AS referring_table, + con.conname AS constraint_name, + con.convalidated AS is_validated, + referred_ns.nspname AS referred_schema, + referred_cls.relname AS referred_table + FROM + pg_constraint con + INNER JOIN + pg_class referring_cls + ON con.conrelid = referring_cls.oid + INNER JOIN + pg_namespace referring_ns + ON referring_cls.relnamespace = referring_ns.oid + INNER JOIN + pg_class referred_cls + ON con.confrelid = referred_cls.oid + INNER JOIN + pg_namespace referred_ns + ON referred_cls.relnamespace = referred_ns.oid + WHERE + con.contype = 'f' + AND referred_ns.nspname = 'public' + AND referred_cls.relname = 'to_repack' + ORDER BY + referring_ns.nspname, + referring_cls.relname, + con.conname; """) ) - with pytest.raises( - ReferringForeignKeyInDifferentSchema, match=f"{schema}.referring_table" - ): - repack.full() + constraints = cur.fetchall() + assert constraints == [ + ( + "public", + "not_valid_referring_table", + "not_valid_referring", + False, + "public", + "to_repack", + ), + ( + "public", + "referring_table", + "referring_table_to_repack_id_fkey", + True, + "public", + "to_repack", + ), + ( + "sweet_schema", + "not_valid_referring_table", + "cross_schema_not_valid_referring", + False, + "public", + "to_repack", + ), + ( + "sweet_schema", + "referring_table", + "cross_schema_referring", + True, + "public", + "to_repack", + ), + ] + + cur.execute( + dedent(""" + SELECT + p.id, + t.id + FROM + sweet_schema.referring_table p + INNER JOIN + public.to_repack t + ON p.to_repack_id = t.id + ORDER BY + p.id + LIMIT 1; + """) + ) + row = cur.fetchone() + assert row is not None def test_with_deferrable_unique_constraint(connection: _psycopg.Connection) -> None: