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