-
-
Notifications
You must be signed in to change notification settings - Fork 325
Closed
Labels
Description
Describe the bug
I have a PG partitioned table with a foreign key to another partitioned table.
I used to ignore these foreign keys in include_object, which is probably not a good idea, but include_name does not provide enough information/context.
Since 1.18, alembic check raises an exception before include_object is called.
Expected behavior
No exception :)
As I was probably misusing include_object maybe this issue could be a feature request to have the referred_table in include_name for foreign keys
To Reproduce
-- initial schema
CREATE TABLE public.a (
id BIGINT GENERATED ALWAYS AS IDENTITY,
created timestamp with time zone DEFAULT now() NOT NULL
)
PARTITION BY RANGE (created);
CREATE TABLE public.b (
id BIGINT GENERATED ALWAYS AS IDENTITY,
a_id BIGINT NOT NULL,
a_created timestamp with time zone NOT NULL
)
PARTITION BY RANGE (a_created);
ALTER TABLE ONLY public.a
ADD CONSTRAINT pk_a PRIMARY KEY (id, created);
ALTER TABLE ONLY public.b
ADD CONSTRAINT pk_b PRIMARY KEY (id, a_created);
ALTER TABLE public.b
ADD CONSTRAINT fk_b_a_id_a FOREIGN KEY (a_id, a_created) REFERENCES public.a(id, created);
CREATE TABLE public.a_1 PARTITION OF public.a FOR VALUES FROM ('2025-01-01') TO ('2025-02-01');
CREATE TABLE public.a_2 PARTITION OF public.a FOR VALUES FROM ('2025-02-01') TO ('2025-03-01');
CREATE TABLE public.b_1 PARTITION OF public.b FOR VALUES FROM ('2025-01-01') TO ('2025-02-01');
CREATE TABLE public.b_2 PARTITION OF public.b FOR VALUES FROM ('2025-02-01') TO ('2025-03-01');# main.py
from datetime import datetime
from sqlalchemy import (
DateTime,
text,
PrimaryKeyConstraint,
BigInteger,
Identity,
ForeignKeyConstraint,
)
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
class Base(DeclarativeBase):
pass
class A(Base):
__tablename__ = "a"
__table_args__ = (
PrimaryKeyConstraint("id", "created"),
{"postgresql_partition_by": "RANGE (created)"},
)
id: Mapped[int] = mapped_column(BigInteger, Identity(always=True))
created: Mapped[datetime] = mapped_column(
DateTime(timezone=True), server_default=text("NOW()")
)
class B(Base):
__tablename__ = "b"
__table_args__ = (
PrimaryKeyConstraint("id", "a_created"),
ForeignKeyConstraint(
columns=["a_id", "a_created"],
refcolumns=["a.id", "a.created"],
),
{"postgresql_partition_by": "RANGE (a_created)"},
)
id: Mapped[int] = mapped_column(BigInteger, Identity(always=True))
a_id: Mapped[int] = mapped_column(BigInteger)
a_created: Mapped[datetime] = mapped_column(DateTime(timezone=True))# env.py
from logging.config import fileConfig
from sqlalchemy import engine_from_config
from sqlalchemy import pool
from sqlalchemy.sql.schema import SchemaItem
from alembic import context
from main import Base
config = context.config
if config.config_file_name is not None:
fileConfig(config.config_file_name)
target_metadata = Base.metadata
def include_object(
object: "SchemaItem",
name: str | None,
type_: str,
reflected: bool,
compare_to: "SchemaItem | None",
) -> bool:
print("include_object?", name, type_, reflected)
if type_ == "foreign_key_constraint" and object.referred_table.name.startswith(
("a_", "b_")
):
# ignore partitions
return False
return True
def include_name(name, type_, parent_names):
print("include_name?", name, type_, parent_names)
if type_ == "table":
assert isinstance(name, str)
if name.startswith(("a_", "b_")):
# ignore partitions
return False
return True
def run_migrations_online() -> None:
"""Run migrations in 'online' mode.
In this scenario we need to create an Engine
and associate a connection with the context.
"""
connectable = engine_from_config(
config.get_section(config.config_ini_section, {}),
prefix="sqlalchemy.",
poolclass=pool.NullPool,
)
with connectable.connect() as connection:
context.configure(
connection=connection,
target_metadata=target_metadata,
include_name=include_name,
include_object=include_object,
)
with context.begin_transaction():
context.run_migrations()
if context.is_offline_mode():
raise NotImplementedError
else:
run_migrations_online()
Error
❯ uv run alembic check
INFO [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO [alembic.runtime.migration] Will assume transactional DDL.
INFO [alembic.runtime.plugins] setting up autogenerate plugin alembic.autogenerate.schemas
INFO [alembic.runtime.plugins] setting up autogenerate plugin alembic.autogenerate.tables
INFO [alembic.runtime.plugins] setting up autogenerate plugin alembic.autogenerate.types
INFO [alembic.runtime.plugins] setting up autogenerate plugin alembic.autogenerate.constraints
INFO [alembic.runtime.plugins] setting up autogenerate plugin alembic.autogenerate.defaults
INFO [alembic.runtime.plugins] setting up autogenerate plugin alembic.autogenerate.comments
include_name? None schema {}
include_name? a_2 table {'schema_name': None, 'schema_qualified_table_name': 'a_2'}
include_name? a table {'schema_name': None, 'schema_qualified_table_name': 'a'}
include_name? b_2 table {'schema_name': None, 'schema_qualified_table_name': 'b_2'}
include_name? b_1 table {'schema_name': None, 'schema_qualified_table_name': 'b_1'}
include_name? b table {'schema_name': None, 'schema_qualified_table_name': 'b'}
include_name? a_1 table {'schema_name': None, 'schema_qualified_table_name': 'a_1'}
include_object? a table False
include_name? id column {'table_name': 'a', 'schema_name': None, 'schema_qualified_table_name': 'a'}
include_name? created column {'table_name': 'a', 'schema_name': None, 'schema_qualified_table_name': 'a'}
include_object? id column False
include_object? created column False
include_object? b table False
include_name? id column {'table_name': 'b', 'schema_name': None, 'schema_qualified_table_name': 'b'}
include_name? a_id column {'table_name': 'b', 'schema_name': None, 'schema_qualified_table_name': 'b'}
include_name? a_created column {'table_name': 'b', 'schema_name': None, 'schema_qualified_table_name': 'b'}
include_object? id column False
include_object? a_id column False
include_object? a_created column False
include_name? fk_b_a_id_a foreign_key_constraint {'table_name': 'b', 'schema_name': None, 'schema_qualified_table_name': 'b'}
include_name? fk_b_a_id_a_1 foreign_key_constraint {'table_name': 'b', 'schema_name': None, 'schema_qualified_table_name': 'b'}
include_name? fk_b_a_id_a_2 foreign_key_constraint {'table_name': 'b', 'schema_name': None, 'schema_qualified_table_name': 'b'}
Traceback (most recent call last):
File "/Users/kael/test-alembic-1.18/.venv/bin/alembic", line 10, in <module>
sys.exit(main())
^^^^^^
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/config.py", line 1047, in main
CommandLine(prog=prog).main(argv=argv)
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/config.py", line 1037, in main
self.run_cmd(cfg, options)
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/config.py", line 971, in run_cmd
fn(
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/command.py", line 363, in check
script_directory.run_env()
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/script/base.py", line 545, in run_env
util.load_python_file(self.dir, "env.py")
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/util/pyfiles.py", line 116, in load_python_file
module = load_module_py(module_id, path)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/util/pyfiles.py", line 136, in load_module_py
spec.loader.exec_module(module) # type: ignore
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "<frozen importlib._bootstrap_external>", line 999, in exec_module
File "<frozen importlib._bootstrap>", line 488, in _call_with_frames_removed
File "/Users/kael/test-alembic-1.18/alembic/env.py", line 73, in <module>
run_migrations_online()
File "/Users/kael/test-alembic-1.18/alembic/env.py", line 67, in run_migrations_online
context.run_migrations()
File "<string>", line 8, in run_migrations
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/runtime/environment.py", line 969, in run_migrations
self.get_context().run_migrations(**kw)
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/runtime/migration.py", line 614, in run_migrations
for step in self._migrations_fn(heads, self):
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/command.py", line 352, in retrieve_migrations
revision_context.run_autogenerate(rev, context)
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/autogenerate/api.py", line 587, in run_autogenerate
self._run_environment(rev, migration_context, True)
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/autogenerate/api.py", line 634, in _run_environment
compare._populate_migration_script(
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/autogenerate/compare/__init__.py", line 39, in _populate_migration_script
_produce_net_changes(autogen_context, upgrade_ops)
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/autogenerate/compare/__init__.py", line 48, in _produce_net_changes
autogen_context.comparators.dispatch(
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/util/langhelpers.py", line 420, in go
result = fn(*arg, **kw)
^^^^^^^^^^^^^^
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/autogenerate/compare/schema.py", line 51, in _produce_net_changes
autogen_context.comparators.dispatch(
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/util/langhelpers.py", line 420, in go
result = fn(*arg, **kw)
^^^^^^^^^^^^^^
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/autogenerate/compare/tables.py", line 75, in _autogen_for_tables
_compare_tables(
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/autogenerate/compare/tables.py", line 220, in _compare_tables
autogen_context.comparators.dispatch(
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/util/langhelpers.py", line 420, in go
result = fn(*arg, **kw)
^^^^^^^^^^^^^^
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/autogenerate/compare/constraints.py", line 637, in _compare_foreign_keys
impl._create_reflected_constraint_sig(fk) for fk in conn_fks
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/ddl/impl.py", line 768, in _create_reflected_constraint_sig
return _constraint_sig.from_constraint(False, self, constraint, **opts)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/ddl/_autogen.py", line 126, in from_constraint
sig = _clsreg[constraint.__visit_name__](is_metadata, impl, constraint)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/ddl/_autogen.py", line 280, in __init__
) = sqla_compat._fk_spec(const)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/alembic/util/sqla_compat.py", line 308, in _fk_spec
target_schema = constraint.elements[0].column.table.schema
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/sqlalchemy/util/langhelpers.py", line 1226, in __get__
obj.__dict__[self.__name__] = result = self.fget(obj)
^^^^^^^^^^^^^^
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/sqlalchemy/sql/schema.py", line 3199, in column
return self._resolve_column()
^^^^^^^^^^^^^^^^^^^^^^
File "/Users/kael/test-alembic-1.18/.venv/lib/python3.12/site-packages/sqlalchemy/sql/schema.py", line 3222, in _resolve_column
raise exc.NoReferencedTableError(
sqlalchemy.exc.NoReferencedTableError: Foreign key associated with column 'b.a_id' could not find table 'a_1' with which to generate a foreign key to target column 'id'
With alembic 1.17.2, the output ends up with
...
include_name? fk_b_a_id_a foreign_key_constraint {'table_name': 'b', 'schema_name': None, 'schema_qualified_table_name': 'b'}
include_name? fk_b_a_id_a_1 foreign_key_constraint {'table_name': 'b', 'schema_name': None, 'schema_qualified_table_name': 'b'}
include_name? fk_b_a_id_a_2 foreign_key_constraint {'table_name': 'b', 'schema_name': None, 'schema_qualified_table_name': 'b'}
include_object? fk_b_a_id_a_1 foreign_key_constraint True
include_object? fk_b_a_id_a_2 foreign_key_constraint True
No new upgrade operations detected.
Versions.
- OS: MacOS
- Python: 3.12
- Alembic: 1.18.2
- SQLAlchemy: 2.0.46
- Database: postgresql 18
- DBAPI:
Additional context
Have a nice day!
Reactions are currently unavailable