Skip to content

Ignore foreign key on partitioned table (alembic check raises an exception in 1.18) #1787

@mgu

Description

@mgu

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!

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions