From 1ffe313576bec821eddc16756f74403a4f0fed3b Mon Sep 17 00:00:00 2001 From: TP Date: Wed, 29 Jul 2026 03:11:28 +0800 Subject: [PATCH] Let providers contribute tables to 'db clean' This adds a 'db-cleanup-tables' key in provider.yaml to specify a callable returning cleanup configs, discovered via ProvidersManager. The 'db clean' command merges them into the list of tables to clean at runtime. Celery uses it to register 'celery_taskmeta' and 'celery_tasksetmeta' tables, using the '[celery] result_backend_schema' config that schema-qualifies them when the result backend is in its own schema. This allows them to always be correctly found by the 'db clean' command. Both are removed from Core's hardcoded list. --- airflow-core/src/airflow/cli/cli_config.py | 2 +- .../src/airflow/cli/commands/db_command.py | 6 +- .../src/airflow/provider.yaml.schema.json | 7 +++ .../src/airflow/provider_info.schema.json | 7 +++ airflow-core/src/airflow/providers_manager.py | 21 +++++++ airflow-core/src/airflow/utils/db_cleanup.py | 43 +++++++++++++-- .../unit/always/test_providers_manager.py | 26 +++++++++ .../tests/unit/utils/test_db_cleanup.py | 55 ++++++++++++++++--- providers/celery/provider.yaml | 14 +++++ .../src/airflow/providers/celery/cleanup.py | 43 +++++++++++++++ .../providers/celery/get_provider_info.py | 8 +++ .../celery/tests/unit/celery/test_cleanup.py | 36 ++++++++++++ 12 files changed, 251 insertions(+), 17 deletions(-) create mode 100644 providers/celery/src/airflow/providers/celery/cleanup.py create mode 100644 providers/celery/tests/unit/celery/test_cleanup.py diff --git a/airflow-core/src/airflow/cli/cli_config.py b/airflow-core/src/airflow/cli/cli_config.py index 46a0ebfa391f0..89ff4dac31c35 100644 --- a/airflow-core/src/airflow/cli/cli_config.py +++ b/airflow-core/src/airflow/cli/cli_config.py @@ -520,7 +520,7 @@ def string_lower_type(val): help=lazy_object_proxy.Proxy( lambda: ( f"Table names to perform maintenance on (use comma-separated list).\n" - f"Options: {import_string('airflow.cli.commands.db_command.all_tables')}" + f"Options: {import_string('airflow.utils.db_cleanup.get_all_table_names')()}" ) ), type=string_list_type, diff --git a/airflow-core/src/airflow/cli/commands/db_command.py b/airflow-core/src/airflow/cli/commands/db_command.py index 4c15ccc2d488d..1d745ac2c8e7b 100644 --- a/airflow-core/src/airflow/cli/commands/db_command.py +++ b/airflow-core/src/airflow/cli/commands/db_command.py @@ -30,7 +30,7 @@ from airflow.exceptions import AirflowException from airflow.utils import cli as cli_utils, db from airflow.utils.db import _REVISION_HEADS_MAP -from airflow.utils.db_cleanup import config_dict, drop_archived_tables, export_archived_records, run_cleanup +from airflow.utils.db_cleanup import drop_archived_tables, export_archived_records, run_cleanup from airflow.utils.db_manager import _callable_accepts_use_migration_files from airflow.utils.process_utils import execute_interactive from airflow.utils.providers_configuration_loader import providers_configuration_loaded @@ -347,10 +347,6 @@ def _warn_remaining_retries(retrystate: RetryCallState): db.check() -# lazily imported by CLI parser for `help` command -all_tables = sorted(config_dict) - - @cli_utils.action_cli(check_db=False) @providers_configuration_loaded def cleanup_tables(args): diff --git a/airflow-core/src/airflow/provider.yaml.schema.json b/airflow-core/src/airflow/provider.yaml.schema.json index 4f3b53f1077ec..23628fbb37665 100644 --- a/airflow-core/src/airflow/provider.yaml.schema.json +++ b/airflow-core/src/airflow/provider.yaml.schema.json @@ -609,6 +609,13 @@ "type": "string" } }, + "db-cleanup-tables": { + "type": "array", + "description": "Import paths of callables returning `airflow db clean` table configs contributed by the provider", + "items": { + "type": "string" + } + }, "cli": { "type": "array", "description": "CLI command functions exposed by the provider", diff --git a/airflow-core/src/airflow/provider_info.schema.json b/airflow-core/src/airflow/provider_info.schema.json index b3e9bf74b6284..dfd22d34ab488 100644 --- a/airflow-core/src/airflow/provider_info.schema.json +++ b/airflow-core/src/airflow/provider_info.schema.json @@ -385,6 +385,13 @@ "type": "string" } }, + "db-cleanup-tables": { + "type": "array", + "description": "Import paths of callables returning `airflow db clean` table configs contributed by the provider", + "items": { + "type": "string" + } + }, "cli": { "type": "array", "description": "CLI command functions exposed by the provider", diff --git a/airflow-core/src/airflow/providers_manager.py b/airflow-core/src/airflow/providers_manager.py index 8e270209a8b0a..2d177d5250a57 100644 --- a/airflow-core/src/airflow/providers_manager.py +++ b/airflow-core/src/airflow/providers_manager.py @@ -449,6 +449,7 @@ def __init__(self): self._executor_without_check_set: set[tuple[str, str]] = set() self._queue_class_name_set: set[str] = set() self._db_manager_class_name_set: set[str] = set() + self._db_cleanup_table_provider_set: set[str] = set() self._provider_configs: dict[str, dict[str, Any]] = {} self._trigger_info_set: set[TriggerInfo] = set() self._notification_info_set: set[NotificationInfo] = set() @@ -617,6 +618,12 @@ def initialize_providers_db_managers(self): self.initialize_providers_list() self._discover_db_managers() + @provider_info_cache("db_cleanup_tables") + def initialize_providers_db_cleanup_tables(self): + """Lazy initialization of providers db cleanup table information.""" + self.initialize_providers_list() + self._discover_db_cleanup_tables() + @provider_info_cache("notifications") def initialize_providers_notifications(self): """Lazy initialization of providers notifications information.""" @@ -1315,6 +1322,13 @@ def _discover_db_managers(self) -> None: if _correctness_check(provider_package, db_manager_class_name, provider): self._db_manager_class_name_set.add(db_manager_class_name) + def _discover_db_cleanup_tables(self) -> None: + """Retrieve all ``airflow db clean`` table contributors defined in the providers.""" + for provider_package, provider in self._provider_dict.items(): + for callable_name in provider.data.get("db-cleanup-tables", []): + if _correctness_check(provider_package, callable_name, provider): + self._db_cleanup_table_provider_set.add(callable_name) + def _discover_config(self) -> None: """Retrieve all configs defined in the providers.""" for provider_package, provider in self._provider_dict.items(): @@ -1529,6 +1543,12 @@ def db_managers(self) -> list[str]: self.initialize_providers_db_managers() return sorted(self._db_manager_class_name_set) + @property + def db_cleanup_table_providers(self) -> list[str]: + """Import paths of provider callables that contribute ``airflow db clean`` table configs.""" + self.initialize_providers_db_cleanup_tables() + return sorted(self._db_cleanup_table_provider_set) + @property def filesystem_module_names(self) -> list[str]: self.initialize_providers_filesystems() @@ -1592,6 +1612,7 @@ def _cleanup(self): self._executor_class_name_set.clear() self._executor_without_check_set.clear() self._queue_class_name_set.clear() + self._db_cleanup_table_provider_set.clear() self._provider_configs.clear() # Imported lazily to avoid a configuration/providers_manager import cycle during cleanup. diff --git a/airflow-core/src/airflow/utils/db_cleanup.py b/airflow-core/src/airflow/utils/db_cleanup.py index 1a79dbc9e30ce..9836de6bcdc78 100644 --- a/airflow-core/src/airflow/utils/db_cleanup.py +++ b/airflow-core/src/airflow/utils/db_cleanup.py @@ -25,6 +25,8 @@ import csv import dataclasses +import functools +import itertools import logging import os from collections.abc import Generator @@ -38,6 +40,7 @@ from sqlalchemy.orm import aliased from sqlalchemy.sql.expression import ClauseElement, Executable, tuple_ +from airflow._shared.module_loading import import_string from airflow._shared.timezones import timezone from airflow.cli.simple_table import AirflowConsole from airflow.configuration import conf @@ -206,8 +209,6 @@ def readable_config(self): _TableConfig(table_name="xcom", recency_column_name="timestamp", dag_id_column_name="dag_id"), _TableConfig(table_name="_xcom_archive", recency_column_name="timestamp", dag_id_column_name="dag_id"), _TableConfig(table_name="callback_request", recency_column_name="created_at"), - _TableConfig(table_name="celery_taskmeta", recency_column_name="date_done"), - _TableConfig(table_name="celery_tasksetmeta", recency_column_name="date_done"), _TableConfig( table_name="trigger", recency_column_name="created_date", @@ -247,7 +248,40 @@ def readable_config(self): ): config_list.append(_TableConfig(table_name="session", recency_column_name="expiry")) -config_dict: dict[str, _TableConfig] = {x.table_name: x for x in sorted(config_list)} + +def _provider_table_configs() -> Generator[tuple[str, _TableConfig], None, None]: + """ + Collect ``_TableConfig`` data contributed by installed providers. + + Providers declare ``db-cleanup-tables`` in their ``provider.yaml``, which is + a list of import paths to callables that each return a list of dicts of + ``_TableConfig`` kwargs. + """ + from airflow.providers_manager import ProvidersManager + + for import_path in ProvidersManager().db_cleanup_table_providers: + try: + for spec in import_string(import_path)(): + config = _TableConfig(**spec) + yield config.table_name, config + except Exception: + logger.warning( + "Failed to load db clean table configs from %r; skipping.", + import_path, + exc_info=True, + ) + + +@functools.cache +def _all_table_configs() -> dict[str, _TableConfig]: + """Core built-in table configs merged with provider-contributed ones.""" + return dict(itertools.chain(((x.table_name, x) for x in config_list), _provider_table_configs())) + + +@functools.cache +def get_all_table_names() -> list[str]: + """Sorted names of every table ``airflow db clean`` can operate on (core + providers).""" + return sorted(_all_table_configs()) def _check_for_rows(*, session: Session, query: Select, print_rows: bool = False) -> int: @@ -628,9 +662,10 @@ def _suppress_with_logging(table: str, session: Session) -> Generator[SimpleName def _effective_table_names(*, table_names: list[str] | None) -> tuple[list[str], dict[str, _TableConfig]]: + config_dict = _all_table_configs() desired_table_names = set(table_names or config_dict) - outliers = desired_table_names - set(config_dict.keys()) + outliers = desired_table_names.difference(config_dict) if outliers: logger.warning( "The following table(s) are not valid choices and will be skipped: %s", diff --git a/airflow-core/tests/unit/always/test_providers_manager.py b/airflow-core/tests/unit/always/test_providers_manager.py index 25e3774d4f7e7..c6ee7d6b7b03e 100644 --- a/airflow-core/tests/unit/always/test_providers_manager.py +++ b/airflow-core/tests/unit/always/test_providers_manager.py @@ -58,6 +58,10 @@ def from_config(cls): return cls() +def fake_db_cleanup_table_configs(): + return [{"table_name": "fake_schema.fake_table", "recency_column_name": "created_at"}] + + def test_cleanup_providers_manager(cleanup_providers_manager): """Check the cleanup provider manager functionality.""" provider_manager = ProvidersManager() @@ -246,6 +250,28 @@ def test_providers_manager_remote_logging_bad_class_filtered(self): assert "bad" not in providers_manager._remote_logging_by_scheme assert providers_manager._remote_logging_info_list == [] + def test_providers_manager_register_db_cleanup_tables(self): + providers_manager = ProvidersManager() + providers_manager._provider_dict = LazyDictWithCache() + providers_manager._provider_dict["fake.db.cleanup"] = ProviderInfo( + version="0.0.1", + data={"db-cleanup-tables": [f"{__name__}.fake_db_cleanup_table_configs"]}, + ) + providers_manager._discover_db_cleanup_tables() + assert providers_manager._db_cleanup_table_provider_set == { + f"{__name__}.fake_db_cleanup_table_configs" + } + + def test_providers_manager_register_db_cleanup_tables_bad_callable_filtered(self): + providers_manager = ProvidersManager() + providers_manager._provider_dict = LazyDictWithCache() + providers_manager._provider_dict["fake.db.cleanup"] = ProviderInfo( + version="0.0.1", + data={"db-cleanup-tables": ["fake.module.does.not.exist.callable"]}, + ) + providers_manager._discover_db_cleanup_tables() + assert providers_manager._db_cleanup_table_provider_set == set() + def test_connection_form_widgets(self, yaml_ui_metadata_counts): yaml_widgets, _ = yaml_ui_metadata_counts provider_manager = ProvidersManager() diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py b/airflow-core/tests/unit/utils/test_db_cleanup.py index abe11aa7c589d..71bacfd8b23e8 100644 --- a/airflow-core/tests/unit/utils/test_db_cleanup.py +++ b/airflow-core/tests/unit/utils/test_db_cleanup.py @@ -53,7 +53,7 @@ _dump_table_to_file, _get_archived_table_names, _TableConfig, - config_dict, + config_list, drop_archived_tables, export_archived_records, run_cleanup, @@ -72,6 +72,8 @@ pytestmark = pytest.mark.db_test +config_dict = {x.table_name: x for x in config_list} + @pytest.fixture(autouse=True) def clean_database(): @@ -190,7 +192,7 @@ def test_run_cleanup_tables(self, clean_table_mock, table_names): verbose=None, ) run_cleanup(**base_kwargs, table_names=table_names) - assert clean_table_mock.call_count == len(table_names) if table_names else len(config_dict) + assert clean_table_mock.call_count == len(table_names) if table_names else len(config_list) @patch("airflow.utils.db_cleanup._cleanup_table") @patch("airflow.utils.db_cleanup._confirm_delete") @@ -871,8 +873,6 @@ def test_no_models_missing(self): "dag_bundle", # leave alone - not appropriate for cleanup } - from airflow.utils.db_cleanup import config_dict - print(f"all_models={set(all_models)}") print(f"excl+conf={exclusion_list.union(config_dict)}") assert set(all_models) - exclusion_list.union(config_dict) == set() @@ -1297,7 +1297,6 @@ def test_extra_filters_keep_in_flight_rows(self): import uuid6 from airflow.models.connection_test import ConnectionTestRequest, ConnectionTestState - from airflow.utils.db_cleanup import config_dict from airflow.utils.session import create_session cfg = config_dict["connection_test_request"] @@ -1496,13 +1495,15 @@ def test_clean_archive_export_and_drop_schema_qualified_table(self, tmp_path): ) session.commit() + from airflow.utils import db_cleanup + 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 patch.object(db_cleanup, "_all_table_configs", return_value={qualified_name: test_config}): + assert qualified_name in db_cleanup._all_table_configs() with create_session() as session: run_cleanup( @@ -1537,3 +1538,43 @@ def test_clean_archive_export_and_drop_schema_qualified_table(self, tmp_path): session.commit() assert _get_archived_table_names([qualified_name], session) == [] + + +def _fake_provider_cleanup_configs(): + return [{"table_name": "fake_schema.fake_cleanup_table", "recency_column_name": "created_at"}] + + +class TestProviderContributedTableConfigs: + """Provider-contributed ``db-cleanup-tables`` configs are merged into the effective config.""" + + def teardown_method(self): + from airflow.utils import db_cleanup + + db_cleanup._all_table_configs.cache_clear() + db_cleanup.get_all_table_names.cache_clear() + + def test_provider_configs_merged_into_all_tables(self): + from airflow.utils import db_cleanup + + db_cleanup._all_table_configs.cache_clear() + db_cleanup.get_all_table_names.cache_clear() + mock_pm = MagicMock() + mock_pm.db_cleanup_table_providers = [f"{__name__}._fake_provider_cleanup_configs"] + with patch("airflow.providers_manager.ProvidersManager", return_value=mock_pm): + assert "fake_schema.fake_cleanup_table" in db_cleanup.get_all_table_names() + _, effective = db_cleanup._effective_table_names(table_names=["fake_schema.fake_cleanup_table"]) + assert "fake_schema.fake_cleanup_table" in effective + assert effective["fake_schema.fake_cleanup_table"].schema_name == "fake_schema" + + def test_bad_provider_callable_is_skipped(self): + from airflow.utils import db_cleanup + + db_cleanup._all_table_configs.cache_clear() + db_cleanup.get_all_table_names.cache_clear() + mock_pm = MagicMock() + mock_pm.db_cleanup_table_providers = ["airflow.does.not.exist.callable"] + with patch("airflow.providers_manager.ProvidersManager", return_value=mock_pm): + names = db_cleanup.get_all_table_names() + # core tables still present; the broken provider is skipped, not fatal + assert "job" in names + assert "airflow.does.not.exist.callable" not in names diff --git a/providers/celery/provider.yaml b/providers/celery/provider.yaml index 2f94171331bd7..221129753d122 100644 --- a/providers/celery/provider.yaml +++ b/providers/celery/provider.yaml @@ -113,6 +113,9 @@ executors: cli: - airflow.providers.celery.cli.definition.get_celery_cli_commands +db-cleanup-tables: + - airflow.providers.celery.cleanup.get_db_cleanup_table_configs + config: celery_kubernetes_executor: description: | @@ -251,6 +254,17 @@ config: sensitive: true example: "db+postgresql+psycopg://postgres:airflow@postgres/airflow" default: ~ + result_backend_schema: + description: | + The database schema that the Celery result backend tables (``celery_taskmeta``, + ``celery_tasksetmeta``) live in, when they are provisioned into a schema other than + the one used by ``sql_alchemy_conn``. When set, ``airflow db clean`` (and the + archived-table export/drop commands) operate on these tables in this schema instead + of silently skipping them. + version_added: ~ + type: string + example: "celery" + default: ~ result_backend_sqlalchemy_engine_options: description: | Optional configuration dictionary to pass to the Celery result backend SQLAlchemy engine. diff --git a/providers/celery/src/airflow/providers/celery/cleanup.py b/providers/celery/src/airflow/providers/celery/cleanup.py new file mode 100644 index 0000000000000..a84492b259c3c --- /dev/null +++ b/providers/celery/src/airflow/providers/celery/cleanup.py @@ -0,0 +1,43 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""``airflow db clean`` table contributions for the Celery result backend.""" + +from __future__ import annotations + +from typing import Any + +from airflow.providers.common.compat.sdk import conf + + +def get_db_cleanup_table_configs() -> list[dict[str, Any]]: + """ + Specify Celery result backend table configs for ``airflow db clean`` to clean. + + Registered via the ``db-cleanup-tables`` provider extension point. + + This needs to be defined in the provider since the returned table names are + schema-qualified when ``[celery] result_backend_schema`` is set. + """ + schema = conf.get("celery", "result_backend_schema", fallback=None) + + def qualified(table_name: str) -> str: + return f"{schema}.{table_name}" if schema else table_name + + return [ + {"table_name": qualified("celery_taskmeta"), "recency_column_name": "date_done"}, + {"table_name": qualified("celery_tasksetmeta"), "recency_column_name": "date_done"}, + ] diff --git a/providers/celery/src/airflow/providers/celery/get_provider_info.py b/providers/celery/src/airflow/providers/celery/get_provider_info.py index 4a79cc58821d5..af7bcce3324a6 100644 --- a/providers/celery/src/airflow/providers/celery/get_provider_info.py +++ b/providers/celery/src/airflow/providers/celery/get_provider_info.py @@ -45,6 +45,7 @@ def get_provider_info(): "airflow.providers.celery.executors.celery_kubernetes_executor.CeleryKubernetesExecutor", ], "cli": ["airflow.providers.celery.cli.definition.get_celery_cli_commands"], + "db-cleanup-tables": ["airflow.providers.celery.cleanup.get_db_cleanup_table_configs"], "config": { "celery_kubernetes_executor": { "description": "This section only applies if you are using the ``CeleryKubernetesExecutor`` in\n``[core]`` section above\n", @@ -133,6 +134,13 @@ def get_provider_info(): "example": "db+postgresql+psycopg://postgres:airflow@postgres/airflow", "default": None, }, + "result_backend_schema": { + "description": "The database schema that the Celery result backend tables (``celery_taskmeta``,\n``celery_tasksetmeta``) live in, when they are provisioned into a schema other than\nthe one used by ``sql_alchemy_conn``. When set, ``airflow db clean`` (and the\narchived-table export/drop commands) operate on these tables in this schema instead\nof silently skipping them.\n", + "version_added": None, + "type": "string", + "example": "celery", + "default": None, + }, "result_backend_sqlalchemy_engine_options": { "description": "Optional configuration dictionary to pass to the Celery result backend SQLAlchemy engine.\n", "version_added": None, diff --git a/providers/celery/tests/unit/celery/test_cleanup.py b/providers/celery/tests/unit/celery/test_cleanup.py new file mode 100644 index 0000000000000..4a2a503299e72 --- /dev/null +++ b/providers/celery/tests/unit/celery/test_cleanup.py @@ -0,0 +1,36 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from airflow.providers.celery.cleanup import get_db_cleanup_table_configs + +from tests_common.test_utils.config import conf_vars + + +class TestCeleryDbCleanupTableConfigs: + def test_default_uses_bare_table_names(self): + configs = get_db_cleanup_table_configs() + assert [c["table_name"] for c in configs] == ["celery_taskmeta", "celery_tasksetmeta"] + assert all(c["recency_column_name"] == "date_done" for c in configs) + + def test_result_backend_schema_qualifies_table_names(self): + with conf_vars({("celery", "result_backend_schema"): "celery"}): + configs = get_db_cleanup_table_configs() + assert [c["table_name"] for c in configs] == [ + "celery.celery_taskmeta", + "celery.celery_tasksetmeta", + ]