This is an automated email from the ASF dual-hosted git repository.
uranusjr 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 30506688167 Support schema-qualified table names in db clean (#70147)
30506688167 is described below
commit 305066881677d4d3fef84bf03b1b4f741af6dae1
Author: Sean Muth <[email protected]>
AuthorDate: Tue Jul 28 12:31:31 2026 -0500
Support schema-qualified table names in db clean (#70147)
Co-authored-by: TP <[email protected]>
---
airflow-core/src/airflow/utils/db.py | 19 +++-
airflow-core/src/airflow/utils/db_cleanup.py | 90 +++++++++++++------
airflow-core/tests/unit/utils/test_db_cleanup.py | 110 +++++++++++++++++++++++
3 files changed, 189 insertions(+), 30 deletions(-)
diff --git a/airflow-core/src/airflow/utils/db.py
b/airflow-core/src/airflow/utils/db.py
index 1076715ee71..615ecb66c67 100644
--- a/airflow-core/src/airflow/utils/db.py
+++ b/airflow-core/src/airflow/utils/db.py
@@ -1080,12 +1080,19 @@ def synchronize_log_template(*, session: Session =
NEW_SESSION) -> None:
session.add(LogTemplate(filename=filename,
elasticsearch_id=elasticsearch_id))
-def reflect_tables(tables: list[MappedClassProtocol | str] | None, session):
+def reflect_tables(tables: list[MappedClassProtocol | str] | None, session,
schema: str | None = None):
"""
When running checks prior to upgrades, we use reflection to determine
current state of the database.
This function gets the current state of each table in the set of models
provided and returns a SqlAlchemy metadata object containing them.
+
+ A string entry in ``tables`` may be schema-qualified using dot notation
(e.g.
+ ``"celery.celery_taskmeta"``) to reflect a table from a non-default
schema; the returned
+ metadata's ``tables`` mapping is then keyed by that same ``schema.table``
string. When
+ ``tables`` is ``None``, ``schema`` selects which schema to reflect in full
(default: the
+ connection's default schema); SQLAlchemy keys every reflected table by
``schema.table`` in
+ that case too, whenever a non-``None`` schema was given.
"""
import sqlalchemy.schema
@@ -1093,12 +1100,18 @@ def reflect_tables(tables: list[MappedClassProtocol |
str] | None, session):
metadata = sqlalchemy.schema.MetaData()
if tables is None:
- metadata.reflect(bind=bind, resolve_fks=False)
+ metadata.reflect(bind=bind, schema=schema, resolve_fks=False)
else:
for tbl in tables:
try:
table_name = tbl if isinstance(tbl, str) else tbl.__tablename__
- metadata.reflect(bind=bind, only=[table_name],
extend_existing=True, resolve_fks=False)
+ tbl_schema: str | None
+ tbl_schema, sep, name = table_name.partition(".")
+ if not sep:
+ tbl_schema, name = None, table_name
+ metadata.reflect(
+ bind=bind, schema=tbl_schema, only=[name],
extend_existing=True, resolve_fks=False
+ )
except exc.InvalidRequestError:
continue
return metadata
diff --git a/airflow-core/src/airflow/utils/db_cleanup.py
b/airflow-core/src/airflow/utils/db_cleanup.py
index 6deb4e78f68..1a79dbc9e30 100644
--- a/airflow-core/src/airflow/utils/db_cleanup.py
+++ b/airflow-core/src/airflow/utils/db_cleanup.py
@@ -24,11 +24,11 @@ See:
from __future__ import annotations
import csv
+import dataclasses
import logging
import os
from collections.abc import Generator
from contextlib import contextmanager, suppress
-from dataclasses import dataclass
from types import SimpleNamespace
from typing import TYPE_CHECKING, Any
@@ -63,12 +63,27 @@ ARCHIVED_TABLES_FROM_DB_MIGRATIONS = [
]
-@dataclass
+def _split_schema_table(table_name: str) -> tuple[str | None, str]:
+ """Split a possibly schema-qualified table name (``schema.table``) into
``(schema, table)``."""
+ schema, sep, name = table_name.partition(".")
+ if sep:
+ return schema, name
+ return None, table_name
+
+
+def _format_table_name(schema: str | None, table: str) -> str:
+ """Format a fully qualified table name from schema and table names."""
+ if schema:
+ return f"{schema}.{table}"
+ return table
+
+
[email protected]
class _TableConfig:
"""
Config class for performing cleanup on a table.
- :param table_name: the table
+ :param table_name: the table name; may be schema-qualified with dot
notation
:param extra_columns: any columns besides recency_column_name that we'll
need in queries
:param recency_column_name: date column to filter by
:param keep_last: whether the last record should be kept even if it's
older than clean_before_timestamp
@@ -101,20 +116,29 @@ class _TableConfig:
skip_if_referenced: list[tuple[str, str]] | None = None
referenced_pk_column: str = "id"
+ # Calculated from table_name and populated in __post_init__.
+ schema_name: str = dataclasses.field(init=False)
+ bare_table_name: str = dataclasses.field(init=False)
+
def __post_init__(self):
+ self.schema_name, self.bare_table_name =
_split_schema_table(self.table_name)
self.recency_column = column(self.recency_column_name)
if self.dag_id_column_name is None:
self.dag_id_column = None
self.orm_model: Base = table(
- self.table_name, *[column(x) for x in self.extra_columns or
[]], self.recency_column
+ self.bare_table_name,
+ *[column(x) for x in self.extra_columns or []],
+ self.recency_column,
+ schema=self.schema_name,
)
else:
self.dag_id_column = column(self.dag_id_column_name)
self.orm_model: Base = table(
- self.table_name,
+ self.bare_table_name,
*[column(x) for x in self.extra_columns or []],
self.dag_id_column,
self.recency_column,
+ schema=self.schema_name,
)
# skip_if_referenced filters on referenced_pk_column, which must be a
column of orm_model
@@ -133,7 +157,7 @@ class _TableConfig:
@property
def readable_config(self):
return {
- "table": self.orm_model.name,
+ "table": self.table_name,
"recency_column": str(self.recency_column),
"dag_id_column": str(self.dag_id_column),
"keep_last": self.keep_last,
@@ -223,7 +247,7 @@ if (
):
config_list.append(_TableConfig(table_name="session",
recency_column_name="expiry"))
-config_dict: dict[str, _TableConfig] = {x.orm_model.name: x for x in
sorted(config_list)}
+config_dict: dict[str, _TableConfig] = {x.table_name: x for x in
sorted(config_list)}
def _check_for_rows(*, session: Session, query: Select, print_rows: bool =
False) -> int:
@@ -264,6 +288,7 @@ def _do_delete(
bind = session.get_bind()
dialect_name = bind.dialect.name
batch_counter = itertools.count(1)
+ source_table_name = _format_table_name(orm_model.schema, orm_model.name)
while True:
limited_query = query.limit(batch_size) if batch_size else query
@@ -283,7 +308,10 @@ def _do_delete(
# using bulk delete
# create a new table and copy the rows there
timestamp_str = re.sub(r"[^\d]", "",
timezone.utcnow().isoformat())[:14]
- target_table_name =
f"{ARCHIVE_TABLE_PREFIX}{orm_model.name}__{timestamp_str}{suffix}"
+ target_table_name = _format_table_name(
+ orm_model.schema,
+ f"{ARCHIVE_TABLE_PREFIX}{orm_model.name}__{timestamp_str}{suffix}",
+ )
print(f"Moving data to table {target_table_name}")
target_table = None
# Lets the ``finally`` cleanup below tell the failure path (don't let a
@@ -295,7 +323,7 @@ def _do_delete(
if dialect_name == "mysql":
# MySQL with replication needs this split into two queries, so
just do it for all MySQL
# ERROR 1786 (HY000): Statement violates GTID consistency:
CREATE TABLE ... SELECT.
- session.execute(text(f"CREATE TABLE {target_table_name} LIKE
{orm_model.name}"))
+ session.execute(text(f"CREATE TABLE {target_table_name} LIKE
{source_table_name}"))
metadata = reflect_tables([target_table_name], session)
target_table = metadata.tables[target_table_name]
insert_stm = target_table.insert().from_select(target_table.c,
limited_query)
@@ -308,8 +336,8 @@ def _do_delete(
session.commit()
# delete the rows from the old table
- metadata = reflect_tables([orm_model.name, target_table_name],
session)
- source_table = metadata.tables[orm_model.name]
+ metadata = reflect_tables([source_table_name, target_table_name],
session)
+ source_table = metadata.tables[source_table_name]
target_table = metadata.tables[target_table_name]
logger.debug("rows moved; purging from %s", source_table.name)
if dialect_name == "sqlite":
@@ -635,21 +663,26 @@ def _effective_table_names(*, table_names: list[str] |
None) -> tuple[list[str],
def _get_archived_table_names(table_names: list[str] | None, session: Session)
-> list[str]:
inspector = inspect(session.bind)
- db_table_names = [
- x
- for x in (inspector.get_table_names() if inspector else [])
- if x.startswith(ARCHIVE_TABLE_PREFIX) or x in
ARCHIVED_TABLES_FROM_DB_MIGRATIONS
- ]
- effective_table_names, _ = _effective_table_names(table_names=table_names)
- # Filter out tables that don't start with the archive prefix
- archived_table_names = [
- table_name
- for table_name in db_table_names
- if (
- any("__" + x + "__" in table_name for x in effective_table_names)
- or table_name in ARCHIVED_TABLES_FROM_DB_MIGRATIONS
+ _, effective_config_dict = _effective_table_names(table_names=table_names)
+ schemas = {config.schema_name for config in effective_config_dict.values()}
+
+ archived_table_names: list[str] = []
+ for schema in schemas:
+ db_table_names = [
+ name
+ for name in (inspector.get_table_names(schema=schema) if inspector
else [])
+ if name.startswith(ARCHIVE_TABLE_PREFIX)
+ or (schema is None and name in ARCHIVED_TABLES_FROM_DB_MIGRATIONS)
+ ]
+ # Further filter to tables belonging to one of the effective configs
+ archived_table_names.extend(
+ _format_table_name(schema, name)
+ for name in db_table_names
+ if (
+ any(f"__{config.bare_table_name}__" in name for config in
effective_config_dict.values())
+ or (schema is None and name in
ARCHIVED_TABLES_FROM_DB_MIGRATIONS)
+ )
)
- ]
return archived_table_names
@@ -712,9 +745,12 @@ def run_cleanup(
dag_ids=dag_ids,
exclude_dag_ids=exclude_dag_ids,
)
- existing_tables = reflect_tables(tables=None, session=session).tables
+ existing_tables: set[str] = {
+ table
+ for schema in {config.schema_name for config in
effective_config_dict.values()}
+ for table in reflect_tables(tables=None, session=session,
schema=schema).tables
+ }
failed_tables: list[str] = []
-
for table_name, table_config in effective_config_dict.items():
if table_name in existing_tables:
with _suppress_with_logging(table_name, session) as ctx:
diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py
b/airflow-core/tests/unit/utils/test_db_cleanup.py
index 8271d38968b..abe11aa7c58 100644
--- a/airflow-core/tests/unit/utils/test_db_cleanup.py
+++ b/airflow-core/tests/unit/utils/test_db_cleanup.py
@@ -1427,3 +1427,113 @@ def _dag_version_config_without_row_exclusion():
if key in config:
config[key] = None
return config
+
+
+class TestSchemaQualifiedTableConfig:
+ """
+ ``_TableConfig`` / ``reflect_tables`` support for schema-qualified table
names.
+
+ A table config may be schema-qualified with ``schema.table`` dot notation
so that
+ ``airflow db clean`` can reach a table living in a schema other than the
metadata
+ connection's default (for example a result backend provisioned into its
own schema),
+ instead of silently skipping it.
+ """
+
+ def test_table_config_parses_schema_qualified_table_name(self):
+ config = _TableConfig(table_name="my_schema.some_table",
recency_column_name="date_done")
+ assert config.schema_name == "my_schema"
+ assert config.bare_table_name == "some_table"
+ assert config.table_name == "my_schema.some_table"
+ assert config.orm_model.schema == "my_schema"
+ assert config.orm_model.name == "some_table"
+
+ def test_table_config_without_schema_prefix(self):
+ config = _TableConfig(table_name="some_table",
recency_column_name="date_done")
+ assert config.schema_name is None
+ assert config.bare_table_name == "some_table"
+ assert config.table_name == "some_table"
+ assert config.orm_model.schema is None
+
+
[email protected]("postgres")
+class TestSchemaQualifiedTableCleanupIntegration:
+ """
+ End-to-end db clean + archive + export/drop against a schema-qualified
table config --
+ the scenario a table provisioned into a non-default schema (e.g. a
schema-split result
+ backend) hits.
+ """
+
+ SCHEMA = "test_schema_qualified_cleanup"
+ TABLE = "schema_qualified_cleanup_test"
+
+ def setup_method(self):
+ with create_session() as session:
+ session.execute(text(f"DROP SCHEMA IF EXISTS {self.SCHEMA}
CASCADE"))
+ session.execute(text(f"CREATE SCHEMA {self.SCHEMA}"))
+ session.execute(
+ text(
+ f"CREATE TABLE {self.SCHEMA}.{self.TABLE} "
+ "(id serial primary key, task_id varchar(155), date_done
timestamp)"
+ )
+ )
+ session.commit()
+
+ def teardown_method(self):
+ with create_session() as session:
+ session.execute(text(f"DROP SCHEMA IF EXISTS {self.SCHEMA}
CASCADE"))
+ session.commit()
+
+ def test_clean_archive_export_and_drop_schema_qualified_table(self,
tmp_path):
+ old_date = pendulum.now("UTC").subtract(days=400)
+ new_date = pendulum.now("UTC")
+ with create_session() as session:
+ session.execute(
+ text(
+ f"INSERT INTO {self.SCHEMA}.{self.TABLE} (task_id,
date_done) "
+ "VALUES ('old-task', :old_date), ('new-task', :new_date)"
+ ),
+ {"old_date": old_date, "new_date": new_date},
+ )
+ session.commit()
+
+ qualified_name = f"{self.SCHEMA}.{self.TABLE}"
+ # Register a schema-qualified table config for the duration of the
test rather than
+ # relying on any provider-specific configuration -- this exercises the
generic
+ # schema-qualified support directly.
+ test_config = _TableConfig(table_name=qualified_name,
recency_column_name="date_done")
+ with patch.dict(config_dict, {qualified_name: test_config}):
+ assert qualified_name in config_dict
+
+ with create_session() as session:
+ run_cleanup(
+
clean_before_timestamp=pendulum.now("UTC").subtract(days=300),
+ table_names=[qualified_name],
+ dry_run=False,
+ confirm=False,
+ session=session,
+ )
+
+ remaining = session.execute(text(f"SELECT task_id FROM
{qualified_name}")).scalars().all()
+ assert remaining == ["new-task"]
+
+ archived = _get_archived_table_names([qualified_name], session)
+ assert len(archived) == 1
+ assert archived[0].startswith(f"{self.SCHEMA}.")
+
+ archived_rows = session.execute(text(f"SELECT task_id FROM
{archived[0]}")).scalars().all()
+ assert archived_rows == ["old-task"]
+
+ export_archived_records(
+ export_format="csv",
+ output_path=str(tmp_path),
+ table_names=[qualified_name],
+ drop_archives=True,
+ needs_confirm=False,
+ session=session,
+ )
+ # export_archived_records takes a caller-owned session and
does not commit it;
+ # commit here so the reflection-based existence check below,
which may use a
+ # separate connection, observes the DROP TABLE.
+ session.commit()
+
+ assert _get_archived_table_names([qualified_name], session) ==
[]
