Skip to content

Support referring foreign keys from other schemas #58

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Bouallegui-Moudhaffer wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Support referring foreign keys from other schemas #58

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cfa0929
Support referring foreign keys from other schemas
Bouallegui-Moudhaffer Apr 4, 2026
aaef00f
Merge branch 'main' into feature/cross-schema-referring-fks
Bouallegui-Moudhaffer Apr 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 30 additions & 23 deletions src/psycopack/_commands.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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("""
Expand All @@ -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)
)
Expand Down Expand Up @@ -859,41 +862,37 @@ 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)
.format(
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("""
Expand All @@ -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)
)
Expand Down Expand Up @@ -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};"
Expand All @@ -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)
)
Expand Down
124 changes: 74 additions & 50 deletions src/psycopack/_repack.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand All @@ -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(
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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():
Expand Down
Loading