diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index 615ecb66c6799..35c2d0559c652 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -1019,7 +1019,7 @@ def synchronize_log_template(*, session: Session = NEW_SESSION) -> None: from airflow.models.tasklog import LogTemplate metadata = reflect_tables([LogTemplate], session) - log_template_table: Table | None = metadata.tables.get(LogTemplate.__tablename__) + log_template_table: Table | None = metadata.tables.get(LogTemplate.__table__.key) if log_template_table is None: log.info("Log template table does not exist (added in 2.3.0); skipping log template sync.") @@ -1104,11 +1104,20 @@ def reflect_tables(tables: list[MappedClassProtocol | str] | None, session, sche else: for tbl in tables: try: - table_name = tbl if isinstance(tbl, str) else tbl.__tablename__ tbl_schema: str | None - tbl_schema, sep, name = table_name.partition(".") - if not sep: - tbl_schema, name = None, table_name + if isinstance(tbl, str): + tbl_schema, sep, name = tbl.partition(".") + if not sep: + tbl_schema, name = None, tbl + else: + # A mapped class already carries its configured schema (e.g. via + # ``sql_alchemy_schema``) on its table; a bare ``__tablename__`` would + # discard it and fall back to the connection's default schema. + # ``__table__`` isn't declared on ``MappedClassProtocol`` (SQLAlchemy sets + # it dynamically, invisible to mypy's static checks here), so read it via + # ``getattr`` instead of a direct attribute access. + name = tbl.__tablename__ + tbl_schema = getattr(tbl, "__table__").schema metadata.reflect( bind=bind, schema=tbl_schema, only=[name], extend_existing=True, resolve_fks=False ) diff --git a/airflow-core/tests/unit/utils/test_db.py b/airflow-core/tests/unit/utils/test_db.py index c81859edbdcfa..6464943faf56b 100644 --- a/airflow-core/tests/unit/utils/test_db.py +++ b/airflow-core/tests/unit/utils/test_db.py @@ -31,7 +31,7 @@ from alembic.migration import MigrationContext from alembic.runtime.environment import EnvironmentContext from alembic.script import ScriptDirectory -from sqlalchemy import Column, Integer, MetaData, Table, select +from sqlalchemy import Column, Integer, MetaData, String, Table, select from airflow import settings from airflow.models import Base as airflow_base @@ -46,7 +46,9 @@ create_default_connections, downgrade, initdb, + reflect_tables, resetdb, + synchronize_log_template, upgradedb, ) from airflow.utils.db_manager import RunDBManager @@ -402,6 +404,49 @@ def scalar(self, stmt): assert bool(lss) is False +class TestReflectTablesSchemaQualified: + def test_reflect_tables_uses_mapped_class_configured_schema(self, mocker): + metadata = MetaData() + Table("scoped_model", metadata, Column("id", Integer, primary_key=True), schema="custom_schema") + + class _ScopedModel: + __tablename__ = "scoped_model" + __table__ = metadata.tables["custom_schema.scoped_model"] + + mock_reflect = mocker.patch.object(MetaData, "reflect") + session = mocker.MagicMock() + + reflect_tables([_ScopedModel], session) + + mock_reflect.assert_called_once_with( + bind=session.bind, + schema="custom_schema", + only=["scoped_model"], + extend_existing=True, + resolve_fks=False, + ) + + def test_synchronize_log_template_finds_table_in_configured_schema(self, mocker): + from airflow.models.tasklog import LogTemplate + + metadata = MetaData() + Table( + LogTemplate.__tablename__, + metadata, + Column("id", Integer, primary_key=True), + Column("filename", String(2000)), + Column("elasticsearch_id", String(2000)), + schema="custom_schema", + ) + mocker.patch.object(LogTemplate.__table__, "schema", "custom_schema") + mocker.patch("airflow.utils.db.reflect_tables", return_value=metadata) + session = mocker.MagicMock() + + synchronize_log_template(session=session) + + session.execute.assert_called() + + class TestAutocommitEngineForMySQL: """Test the AutocommitEngineForMySQL context manager."""