Let providers contribute tables to 'db clean'

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,

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):

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",

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",

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.

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",

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()

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

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.