This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 6aee277d4a9 Fix log template sync failing when a custom SQL Alchemy
schema is configured (#70793)
6aee277d4a9 is described below
commit 6aee277d4a9bbf518d9c6359f0d0ab468b4da898
Author: Aaryan Mahajan <[email protected]>
AuthorDate: Tue Oct 20 13:18:49 2026 +0530
Fix log template sync failing when a custom SQL Alchemy schema is
configured (#70793)
* Fix log template sync failing when a custom SQL Alchemy schema is
configured
When sql_alchemy_schema is set, reflect_tables() derived the schema to
reflect a mapped class's table under by splitting its bare __tablename__
on a dot, which discards the schema entirely for classes like
LogTemplate whose table name has no dot. synchronize_log_template()
then looked up the reflected table by that same bare name, which also
does not match how SQLAlchemy keys schema-qualified tables in the
reflected metadata. Together these left the log_template table
unrecognized under a custom schema, causing triggering a Dag to fail
with a TypeError.
Derive the schema from the mapped class's own configured table instead
of parsing its name, and look the reflected table up by its schema
-qualified key.
* Add newsfragment for log template custom schema fix
* Remove newsfragment for log template custom schema fix
Not needed for this change per review feedback.
* Use inspect().local_table instead of getattr for the mapped class's schema
getattr(tbl, "__table__") worked around __table__ not being declared on
MappedClassProtocol, but SQLAlchemy's inspect() already exposes the same
table through a properly typed Mapper.local_table, so the getattr and its
explanatory comment are no longer needed.
---
airflow-core/src/airflow/utils/db.py | 14 +++++----
airflow-core/tests/unit/utils/test_db.py | 50 +++++++++++++++++++++++++++++++-
2 files changed, 58 insertions(+), 6 deletions(-)
diff --git a/airflow-core/src/airflow/utils/db.py
b/airflow-core/src/airflow/utils/db.py
index 335284a2016..6e078f3b0e9 100644
--- a/airflow-core/src/airflow/utils/db.py
+++ b/airflow-core/src/airflow/utils/db.py
@@ -1020,7 +1020,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.")
@@ -1105,11 +1105,15 @@ def reflect_tables(tables: Sequence[MappedClassProtocol
| str] | None, session,
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:
+ # Preserve the mapped class's own configured schema; see
PR #70793.
+ name = tbl.__tablename__
+ tbl_schema = inspect(tbl, raiseerr=True).local_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 83a1e653045..6af67dec69a 100644
--- a/airflow-core/tests/unit/utils/test_db.py
+++ b/airflow-core/tests/unit/utils/test_db.py
@@ -30,7 +30,8 @@ from alembic.config import Config
from alembic.migration import MigrationContext
from alembic.runtime.environment import EnvironmentContext
from alembic.script import ScriptDirectory
-from sqlalchemy import Column, Integer, MetaData, Table, delete, select
+from sqlalchemy import Column, Integer, MetaData, String, Table, delete, select
+from sqlalchemy.orm import declarative_base
from airflow import settings
from airflow.models import Base as airflow_base
@@ -48,7 +49,9 @@ from airflow.utils.db import (
downgrade,
get_default_connections,
initdb,
+ reflect_tables,
resetdb,
+ synchronize_log_template,
upgradedb,
)
from airflow.utils.db_manager import RunDBManager
@@ -453,6 +456,51 @@ class TestDb:
assert bool(lss) is False
+class TestReflectTablesSchemaQualified:
+ def test_reflect_tables_uses_mapped_class_configured_schema(self, mocker):
+ metadata = MetaData()
+ table = Table(
+ "scoped_model", metadata, Column("id", Integer, primary_key=True),
schema="custom_schema"
+ )
+
+ class _ScopedModel(declarative_base(metadata=metadata)):
+ __tablename__ = "scoped_model"
+ __table__ = table
+
+ 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."""
