From 8975d7abf7302a9ccf00d68e57504b5805be9627 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Thu, 6 Aug 2026 10:51:24 +0530 Subject: [PATCH 1/3] Convert check-sdk-imports to capture occurences than files --- airflow-core/.pre-commit-config.yaml | 103 +------------- generated/known_sdk_imports_in_core.txt | 91 ++++++++++++ scripts/ci/prek/check_sdk_imports_in_core.py | 120 +++++++++++++--- .../ci/prek/test_check_sdk_imports_in_core.py | 133 +++++++++++++++++- 4 files changed, 328 insertions(+), 119 deletions(-) create mode 100644 generated/known_sdk_imports_in_core.txt diff --git a/airflow-core/.pre-commit-config.yaml b/airflow-core/.pre-commit-config.yaml index 08382f3934f43..c59444e234d5d 100644 --- a/airflow-core/.pre-commit-config.yaml +++ b/airflow-core/.pre-commit-config.yaml @@ -308,113 +308,12 @@ repos: ^tests/unit/core/test_configuration\.py$| ^tests/unit/models/test_renderedtifields\.py$| ^tests/unit/models/test_variable\.py$ - - id: check-sdk-imports + - id: check-sdk-imports-in-core name: Check for SDK imports in core files entry: ../scripts/ci/prek/check_sdk_imports_in_core.py language: python types: [python] files: ^src/airflow/ - exclude: | - (?x) - # Allow SDK imports in these legitimate locations - ^src/airflow/example_dags/.*\.py$| - - # TODO: These files need to be refactored to remove SDK coupling - ^src/airflow/__init__\.py$| - ^src/airflow/api/common/mark_tasks\.py$| - ^src/airflow/api_fastapi/core_api/datamodels/assets\.py$| - ^src/airflow/api_fastapi/core_api/datamodels/connections\.py$| - ^src/airflow/api_fastapi/core_api/datamodels/hitl\.py$| - ^src/airflow/api_fastapi/core_api/datamodels/variables\.py$| - ^src/airflow/api_fastapi/core_api/routes/ui/grid\.py$| - ^src/airflow/api_fastapi/core_api/routes/ui/structure\.py$| - ^src/airflow/api_fastapi/core_api/services/public/connections\.py$| - ^src/airflow/api_fastapi/core_api/services/ui/connections\.py$| - ^src/airflow/api_fastapi/core_api/services/ui/grid\.py$| - ^src/airflow/api_fastapi/core_api/services/ui/task_group.py$| - ^src/airflow/api_fastapi/execution_api/routes/hitl\.py$| - ^src/airflow/api_fastapi/execution_api/routes/task_instances\.py$| - ^src/airflow/api_fastapi/execution_api/versions/v2026_04_06\.py$| - ^src/airflow/api_fastapi/logging/decorators\.py$| - ^src/airflow/assets/evaluation\.py$| - ^src/airflow/assets/manager\.py$| - ^src/airflow/cli/commands/connection_command\.py$| - ^src/airflow/cli/commands/task_command\.py$| - ^src/airflow/cli/commands/triggerer_command.py$| - ^src/airflow/configuration\.py$| - ^src/airflow/dag_processing/collection\.py$| - ^src/airflow/dag_processing/dagbag\.py$| - ^src/airflow/dag_processing/importers/.*\.py$| - ^src/airflow/dag_processing/manager\.py$| - ^src/airflow/dag_processing/processor\.py$| - ^src/airflow/datasets/metadata\.py$| - ^src/airflow/exceptions\.py$| - ^src/airflow/executors/base_executor\.py$| - ^src/airflow/jobs/triggerer_job_runner\.py$| - ^src/airflow/lineage/hook\.py$| - ^src/airflow/listeners/spec/asset\.py$| - ^src/airflow/listeners/spec/taskinstance\.py$| - ^src/airflow/logging/remote\.py$| - ^src/airflow/models/__init__\.py$| - ^src/airflow/models/asset\.py$| - ^src/airflow/models/baseoperator\.py$| - ^src/airflow/models/callback\.py$| - ^src/airflow/models/connection\.py$| - ^src/airflow/models/dag\.py$| - ^src/airflow/models/dagrun\.py$| - ^src/airflow/models/deadline\.py$| - ^src/airflow/models/expandinput\.py$| - ^src/airflow/models/mappedoperator\.py$| - ^src/airflow/models/operator\.py$| - ^src/airflow/models/param\.py$| - ^src/airflow/models/renderedtifields\.py$| - ^src/airflow/models/serialized_dag\.py$| - ^src/airflow/models/taskinstance\.py$| - ^src/airflow/models/taskinstancekey\.py$| - ^src/airflow/models/taskmap\.py$| - ^src/airflow/models/taskmixin\.py$| - ^src/airflow/models/taskreschedule\.py$| - ^src/airflow/models/trigger\.py$| - ^src/airflow/stats\.py$| - ^src/airflow/models/variable\.py$| - ^src/airflow/models/xcom\.py$| - ^src/airflow/models/xcom_arg\.py$| - ^src/airflow/operators/subdag\.py$| - ^src/airflow/plugins_manager\.py$| - ^src/airflow/providers_manager\.py$| - ^src/airflow/secrets/__init__.py$| - ^src/airflow/serialization/decoders\.py$| - ^src/airflow/serialization/definitions/[_/a-z]+\.py$| - ^src/airflow/serialization/encoders\.py$| - ^src/airflow/serialization/enums\.py$| - ^src/airflow/serialization/helpers\.py$| - ^src/airflow/serialization/serialized_objects\.py$| - ^src/airflow/settings\.py$| - ^src/airflow/task/task_runner/bash_task_runner\.py$| - ^src/airflow/task/task_runner/standard_task_runner\.py$| - ^src/airflow/ti_deps/deps/mapped_task_upstream_dep\.py$| - ^src/airflow/ti_deps/deps/prev_dagrun_dep\.py$| - ^src/airflow/ti_deps/deps/trigger_rule_dep\.py$| - ^src/airflow/timetables/assets\.py$| - ^src/airflow/timetables/base\.py$| - ^src/airflow/timetables/simple\.py$| - ^src/airflow/triggers/base\.py$| - ^src/airflow/utils/cli\.py$| - ^src/airflow/utils/context\.py$| - ^src/airflow/utils/dag_cycle_tester\.py$| - ^src/airflow/utils/dag_edges\.py$| - ^src/airflow/utils/dag_parsing_context\.py$| - ^src/airflow/utils/decorators\.py$| - ^src/airflow/utils/dot_renderer\.py$| - ^src/airflow/utils/edgemodifier\.py$| - ^src/airflow/utils/email\.py$| - ^src/airflow/utils/helpers\.py$| - ^src/airflow/utils/operator_helpers\.py$| - ^src/airflow/utils/session\.py$| - ^src/airflow/utils/task_group\.py$| - ^src/airflow/utils/trigger_rule\.py$| - ^src/airflow/utils/json\.py$| - ^src/airflow/utils/types\.py$ ## ONLY ADD PREK HOOKS HERE THAT REQUIRE CI IMAGE - id: check-schema-defaults name: Check schema defaults match server-side defaults diff --git a/generated/known_sdk_imports_in_core.txt b/generated/known_sdk_imports_in_core.txt new file mode 100644 index 0000000000000..bf2e2b6dba457 --- /dev/null +++ b/generated/known_sdk_imports_in_core.txt @@ -0,0 +1,91 @@ +airflow-core/src/airflow/__init__.py::1 +airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py::1 +airflow-core/src/airflow/cli/commands/task_command.py::7 +airflow-core/src/airflow/cli/commands/triggerer_command.py::1 +airflow-core/src/airflow/configuration.py::1 +airflow-core/src/airflow/dag_processing/dagbag.py::1 +airflow-core/src/airflow/dag_processing/importers/base.py::1 +airflow-core/src/airflow/dag_processing/importers/python_importer.py::7 +airflow-core/src/airflow/dag_processing/manager.py::4 +airflow-core/src/airflow/dag_processing/processor.py::15 +airflow-core/src/airflow/example_dags/example_asset_alias.py::1 +airflow-core/src/airflow/example_dags/example_asset_alias_with_no_taskflow.py::1 +airflow-core/src/airflow/example_dags/example_asset_allow_teams.py::1 +airflow-core/src/airflow/example_dags/example_asset_decorator.py::1 +airflow-core/src/airflow/example_dags/example_asset_partition.py::1 +airflow-core/src/airflow/example_dags/example_asset_state_store.py::1 +airflow-core/src/airflow/example_dags/example_asset_with_watchers.py::1 +airflow-core/src/airflow/example_dags/example_assets.py::1 +airflow-core/src/airflow/example_dags/example_branch_labels.py::1 +airflow-core/src/airflow/example_dags/example_branch_python_dop_operator_3.py::1 +airflow-core/src/airflow/example_dags/example_complex.py::1 +airflow-core/src/airflow/example_dags/example_custom_weight.py::1 +airflow-core/src/airflow/example_dags/example_dag_decorator.py::2 +airflow-core/src/airflow/example_dags/example_display_name.py::1 +airflow-core/src/airflow/example_dags/example_dynamic_task_mapping.py::1 +airflow-core/src/airflow/example_dags/example_dynamic_task_mapping_with_no_taskflow_operators.py::1 +airflow-core/src/airflow/example_dags/example_failed_dag.py::1 +airflow-core/src/airflow/example_dags/example_inlet_event_extra.py::1 +airflow-core/src/airflow/example_dags/example_kubernetes_executor.py::1 +airflow-core/src/airflow/example_dags/example_latest_only_with_trigger.py::1 +airflow-core/src/airflow/example_dags/example_local_kubernetes_executor.py::1 +airflow-core/src/airflow/example_dags/example_nested_branch_dag.py::1 +airflow-core/src/airflow/example_dags/example_outlet_event_extra.py::1 +airflow-core/src/airflow/example_dags/example_params_trigger_ui.py::1 +airflow-core/src/airflow/example_dags/example_params_ui_tutorial.py::1 +airflow-core/src/airflow/example_dags/example_passing_params_via_test_command.py::1 +airflow-core/src/airflow/example_dags/example_retry_policy.py::1 +airflow-core/src/airflow/example_dags/example_setup_teardown.py::1 +airflow-core/src/airflow/example_dags/example_setup_teardown_taskflow.py::1 +airflow-core/src/airflow/example_dags/example_simplest_dag.py::1 +airflow-core/src/airflow/example_dags/example_skip_dag.py::3 +airflow-core/src/airflow/example_dags/example_task_group.py::1 +airflow-core/src/airflow/example_dags/example_task_group_decorator.py::1 +airflow-core/src/airflow/example_dags/example_task_state_store.py::2 +airflow-core/src/airflow/example_dags/example_task_state_store_mapped.py::1 +airflow-core/src/airflow/example_dags/example_time_delta_sensor_async.py::1 +airflow-core/src/airflow/example_dags/example_trigger_target_dag.py::1 +airflow-core/src/airflow/example_dags/example_workday_timetable.py::1 +airflow-core/src/airflow/example_dags/example_xcom.py::1 +airflow-core/src/airflow/example_dags/example_xcomargs.py::1 +airflow-core/src/airflow/example_dags/plugins/event_listener.py::1 +airflow-core/src/airflow/example_dags/tutorial.py::1 +airflow-core/src/airflow/example_dags/tutorial_dag.py::1 +airflow-core/src/airflow/example_dags/tutorial_objectstorage.py::1 +airflow-core/src/airflow/example_dags/tutorial_taskflow_api.py::1 +airflow-core/src/airflow/example_dags/tutorial_taskflow_api_virtualenv.py::1 +airflow-core/src/airflow/example_dags/tutorial_taskflow_templates.py::1 +airflow-core/src/airflow/exceptions.py::1 +airflow-core/src/airflow/executors/base_executor.py::3 +airflow-core/src/airflow/jobs/triggerer_job_runner.py::18 +airflow-core/src/airflow/models/__init__.py::5 +airflow-core/src/airflow/models/asset.py::2 +airflow-core/src/airflow/models/connection.py::5 +airflow-core/src/airflow/models/dag.py::1 +airflow-core/src/airflow/models/dagrun.py::2 +airflow-core/src/airflow/models/taskinstance.py::1 +airflow-core/src/airflow/models/trigger.py::4 +airflow-core/src/airflow/models/variable.py::8 +airflow-core/src/airflow/models/xcom_arg.py::1 +airflow-core/src/airflow/plugins_manager.py::1 +airflow-core/src/airflow/providers_manager.py::5 +airflow-core/src/airflow/secrets/__init__.py::1 +airflow-core/src/airflow/serialization/decoders.py::2 +airflow-core/src/airflow/serialization/definitions/baseoperator.py::1 +airflow-core/src/airflow/serialization/definitions/dag.py::2 +airflow-core/src/airflow/serialization/definitions/deadline.py::1 +airflow-core/src/airflow/serialization/definitions/mappedoperator.py::5 +airflow-core/src/airflow/serialization/encoders.py::11 +airflow-core/src/airflow/serialization/serialized_objects.py::16 +airflow-core/src/airflow/settings.py::1 +airflow-core/src/airflow/stats.py::1 +airflow-core/src/airflow/timetables/simple.py::1 +airflow-core/src/airflow/triggers/base.py::3 +airflow-core/src/airflow/utils/cli.py::2 +airflow-core/src/airflow/utils/context.py::2 +airflow-core/src/airflow/utils/dag_cycle_tester.py::1 +airflow-core/src/airflow/utils/dag_edges.py::1 +airflow-core/src/airflow/utils/dot_renderer.py::3 +airflow-core/src/airflow/utils/edgemodifier.py::1 +airflow-core/src/airflow/utils/json.py::4 +airflow-core/src/airflow/utils/operator_helpers.py::2 diff --git a/scripts/ci/prek/check_sdk_imports_in_core.py b/scripts/ci/prek/check_sdk_imports_in_core.py index 0319272893007..c07a025fd6249 100755 --- a/scripts/ci/prek/check_sdk_imports_in_core.py +++ b/scripts/ci/prek/check_sdk_imports_in_core.py @@ -22,13 +22,49 @@ # "rich>=13.6.0", # ] # /// +"""Check that no new ``airflow.sdk`` imports are introduced in ``airflow-core``. + +All *existing* imports are recorded in ``generated/known_sdk_imports_in_core.txt`` +as ``relative/path::N`` entries (one per file), where ``N`` is the maximum +number of ``airflow.sdk`` import statements allowed in that file. A file whose +current count exceeds the recorded limit is treated as a violation -- core +should not gain new runtime dependencies on the Task SDK; use ``# noqa: SDK001`` +for a one off exception, or regenerate the allowlist for a deliberate one. + +Modes +----- +Default (files passed by prek/pre-commit): + Check only the supplied files; fail if any file's count exceeds the limit. + When a file's count has *decreased*, the allowlist entry is tightened + automatically and the hook exits with a non-zero code so that pre-commit + reports the modified allowlist -- just stage + ``generated/known_sdk_imports_in_core.txt`` and re-run. + +``--all-files``: + Walk all of ``airflow-core/src/airflow`` and check every ``.py`` file. + +``--cleanup``: + Remove entries for files that no longer exist. Safe to run at any time; + does not add new entries or raise limits. + +``--generate``: + Scan ``airflow-core/src/airflow`` and *rebuild* the allowlist from scratch. + Intended for the initial setup or after a large-scale clean-up sprint. +""" + from __future__ import annotations import argparse -import sys +from collections.abc import Iterable from pathlib import Path -from common_prek_utils import find_import_violations, report_import_violations +from common_prek_utils import AIRFLOW_CORE_ROOT_PATH, AllowlistManager, find_import_violations +from rich.console import Console + +console = Console(color_system="standard", width=200) + +REPO_ROOT = AIRFLOW_CORE_ROOT_PATH.parent +CORE_SRC_ROOT = AIRFLOW_CORE_ROOT_PATH / "src" / "airflow" NOCHECK_CODE = "SDK001" @@ -42,23 +78,75 @@ def check_file_for_sdk_imports(file_path: Path) -> list[tuple[int, str]]: ) -def main(): - parser = argparse.ArgumentParser(description="Check for SDK imports in airflow-core files") - parser.add_argument("files", nargs="*", help="Files to check") - args = parser.parse_args() +class SdkImportsAllowlistManager(AllowlistManager): + def __init__(self, allowlist_file: Path) -> None: + super().__init__(allowlist_file, repo_root=REPO_ROOT) - if not args.files: - return + def iter_files(self) -> Iterable[Path]: + return CORE_SRC_ROOT.rglob("*.py") - report_import_violations( - args.files, - check_func=check_file_for_sdk_imports, - violation_label="SDK import(s) in core files", - nocheck_code=NOCHECK_CODE, - only_python_files=True, + def count_occurrences(self, path: Path) -> int: + return len(check_file_for_sdk_imports(path)) + + def violation_panel_text(self) -> str: + return ( + "New [bold]airflow.sdk[/bold] import detected in airflow-core.\n" + "Core (scheduler/API server) should not gain new runtime dependencies " + "on the Task SDK.\n" + "If this import is a genuine one-off, append `# noqa: SDK001` to the " + "import line.\n" + "If it's an intentional, broader exception, run:\n\n" + " [cyan]uv run ./scripts/ci/prek/check_sdk_imports_in_core.py --generate[/cyan]\n\n" + "to regenerate the allowlist, then commit the updated\n" + "[cyan]generated/known_sdk_imports_in_core.txt[/cyan]." + ) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description="Prevent new airflow.sdk imports in airflow-core.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + parser.add_argument("files", nargs="*", metavar="FILE", help="Files to check (provided by prek)") + parser.add_argument( + "--all-files", + action="store_true", + help="Check every Python file under airflow-core/src/airflow", + ) + parser.add_argument( + "--cleanup", + action="store_true", + help="Remove stale entries from the allowlist and exit", ) + parser.add_argument( + "--generate", + action="store_true", + help="Regenerate the allowlist from the current codebase and exit", + ) + args = parser.parse_args(argv) + + manager = SdkImportsAllowlistManager(REPO_ROOT / "generated" / "known_sdk_imports_in_core.txt") + + if args.generate: + return manager.generate() + + if args.cleanup: + return manager.cleanup() + + allowlist = manager.load() + + if args.all_files: + return manager.check(list(manager.iter_files()), allowlist) + + if not args.files: + console.print( + "[yellow]No files provided. Pass filenames or use --all-files to scan the whole repo.[/yellow]" + ) + return 0 + + return manager.check([Path(f).resolve() for f in args.files], allowlist) if __name__ == "__main__": - main() - sys.exit(0) + raise SystemExit(main()) diff --git a/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py b/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py index ed1bcbe057f46..c3c6a90a2cdb5 100644 --- a/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py +++ b/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py @@ -19,8 +19,23 @@ import textwrap from pathlib import Path +import check_sdk_imports_in_core as hook import pytest -from check_sdk_imports_in_core import check_file_for_sdk_imports +from check_sdk_imports_in_core import SdkImportsAllowlistManager, check_file_for_sdk_imports + + +@pytest.fixture +def create_fake_core_repo(tmp_path, monkeypatch): + monkeypatch.setattr(hook, "REPO_ROOT", tmp_path) + monkeypatch.setattr(hook, "CORE_SRC_ROOT", tmp_path / "airflow-core" / "src" / "airflow") + + def _write(rel: str, code: str) -> Path: + path = tmp_path / "airflow-core" / "src" / "airflow" / rel + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(textwrap.dedent(code)) + return path + + return _write class TestCheckFileForSdkImports: @@ -178,3 +193,119 @@ def test_nocheck_marker(self, tmp_path: Path, code: str, expected: list[tuple[in f = tmp_path / "example.py" f.write_text(code) assert check_file_for_sdk_imports(f) == expected + + +class TestSdkImportsAllowlistRatchet: + def test_no_violations_passes(self, create_fake_core_repo, tmp_path): + path = create_fake_core_repo( + "models/clean.py", + """\ + from airflow.models import DagRun + """, + ) + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + assert manager.check([path], {}) == 0 + + def test_new_violation_fails(self, create_fake_core_repo, tmp_path): + path = create_fake_core_repo( + "models/bad.py", + """\ + from airflow.sdk import DAG + """, + ) + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + assert manager.check([path], {}) == 1 + + def test_violation_within_allowlist_passes(self, create_fake_core_repo, tmp_path): + path = create_fake_core_repo( + "models/grandfathered.py", + """\ + from airflow.sdk import DAG + """, + ) + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + allowlist = {"airflow-core/src/airflow/models/grandfathered.py": 1} + assert manager.check([path], allowlist) == 0 + + def test_exceeding_allowlist_fails(self, create_fake_core_repo, tmp_path): + path = create_fake_core_repo( + "models/grew.py", + """\ + from airflow.sdk import DAG + from airflow.sdk.definitions.deadline import VariableInterval + """, + ) + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + allowlist = {"airflow-core/src/airflow/models/grew.py": 1} + assert manager.check([path], allowlist) == 1 + + def test_reducing_violations_tightens_allowlist(self, create_fake_core_repo, tmp_path): + path = create_fake_core_repo( + "models/improved.py", + """\ + from airflow.sdk import DAG + """, + ) + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + allowlist = {"airflow-core/src/airflow/models/improved.py": 2} + assert manager.check([path], allowlist) == 1 + assert manager.load() == {"airflow-core/src/airflow/models/improved.py": 1} + + def test_fixing_all_violations_removes_entry(self, create_fake_core_repo, tmp_path): + path = create_fake_core_repo( + "models/fixed.py", + """\ + from airflow.models import DagRun + """, + ) + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + allowlist = {"airflow-core/src/airflow/models/fixed.py": 1} + assert manager.check([path], allowlist) == 1 + assert manager.load() == {} + + def test_noqa_marker_avoids_ratchet_entirely(self, create_fake_core_repo, tmp_path): + path = create_fake_core_repo( + "models/one_off.py", + """\ + from airflow.sdk import DAG # noqa: SDK001 + """, + ) + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + assert manager.check([path], {}) == 0 + + def test_non_python_file_is_skipped(self, create_fake_core_repo, tmp_path): + path = create_fake_core_repo( + "models/not_python.txt", + "from airflow.sdk import DAG\n", + ) + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + assert manager.check([path], {}) == 0 + + +class TestSdkImportsAllowlistCleanup: + def test_cleanup_removes_stale_entries(self, create_fake_core_repo, tmp_path): + create_fake_core_repo("models/keeper.py", "from airflow.models import DagRun\n") + allowlist_path = tmp_path / "allowlist.txt" + manager = SdkImportsAllowlistManager(allowlist_path) + manager.save( + { + "airflow-core/src/airflow/models/keeper.py": 0, + "airflow-core/src/airflow/models/gone.py": 1, + } + ) + assert manager.cleanup() == 0 + assert manager.load() == {"airflow-core/src/airflow/models/keeper.py": 0} + + def test_cleanup_empty_allowlist(self, tmp_path): + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + assert manager.cleanup() == 0 + + +class TestSdkImportsAllowlistGenerate: + def test_generate_records_current_occurrences(self, create_fake_core_repo, tmp_path): + create_fake_core_repo("models/a.py", "from airflow.sdk import DAG\n") + create_fake_core_repo("models/b.py", "from airflow.models import DagRun\n") + allowlist_path = tmp_path / "allowlist.txt" + manager = SdkImportsAllowlistManager(allowlist_path) + assert manager.generate() == 0 + assert manager.load() == {"airflow-core/src/airflow/models/a.py": 1} From 766783d4bd0a42f80d5a7284c868abca69a822f0 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Fri, 7 Aug 2026 18:58:48 +0530 Subject: [PATCH 2/3] comments from jason and dennis --- generated/known_sdk_imports_in_core.txt | 47 ---------- scripts/ci/prek/check_sdk_imports_in_core.py | 92 +++++++++++++++---- .../ci/prek/test_check_sdk_imports_in_core.py | 65 +++++++++++++ 3 files changed, 140 insertions(+), 64 deletions(-) diff --git a/generated/known_sdk_imports_in_core.txt b/generated/known_sdk_imports_in_core.txt index bf2e2b6dba457..b93815ae6d968 100644 --- a/generated/known_sdk_imports_in_core.txt +++ b/generated/known_sdk_imports_in_core.txt @@ -8,53 +8,6 @@ airflow-core/src/airflow/dag_processing/importers/base.py::1 airflow-core/src/airflow/dag_processing/importers/python_importer.py::7 airflow-core/src/airflow/dag_processing/manager.py::4 airflow-core/src/airflow/dag_processing/processor.py::15 -airflow-core/src/airflow/example_dags/example_asset_alias.py::1 -airflow-core/src/airflow/example_dags/example_asset_alias_with_no_taskflow.py::1 -airflow-core/src/airflow/example_dags/example_asset_allow_teams.py::1 -airflow-core/src/airflow/example_dags/example_asset_decorator.py::1 -airflow-core/src/airflow/example_dags/example_asset_partition.py::1 -airflow-core/src/airflow/example_dags/example_asset_state_store.py::1 -airflow-core/src/airflow/example_dags/example_asset_with_watchers.py::1 -airflow-core/src/airflow/example_dags/example_assets.py::1 -airflow-core/src/airflow/example_dags/example_branch_labels.py::1 -airflow-core/src/airflow/example_dags/example_branch_python_dop_operator_3.py::1 -airflow-core/src/airflow/example_dags/example_complex.py::1 -airflow-core/src/airflow/example_dags/example_custom_weight.py::1 -airflow-core/src/airflow/example_dags/example_dag_decorator.py::2 -airflow-core/src/airflow/example_dags/example_display_name.py::1 -airflow-core/src/airflow/example_dags/example_dynamic_task_mapping.py::1 -airflow-core/src/airflow/example_dags/example_dynamic_task_mapping_with_no_taskflow_operators.py::1 -airflow-core/src/airflow/example_dags/example_failed_dag.py::1 -airflow-core/src/airflow/example_dags/example_inlet_event_extra.py::1 -airflow-core/src/airflow/example_dags/example_kubernetes_executor.py::1 -airflow-core/src/airflow/example_dags/example_latest_only_with_trigger.py::1 -airflow-core/src/airflow/example_dags/example_local_kubernetes_executor.py::1 -airflow-core/src/airflow/example_dags/example_nested_branch_dag.py::1 -airflow-core/src/airflow/example_dags/example_outlet_event_extra.py::1 -airflow-core/src/airflow/example_dags/example_params_trigger_ui.py::1 -airflow-core/src/airflow/example_dags/example_params_ui_tutorial.py::1 -airflow-core/src/airflow/example_dags/example_passing_params_via_test_command.py::1 -airflow-core/src/airflow/example_dags/example_retry_policy.py::1 -airflow-core/src/airflow/example_dags/example_setup_teardown.py::1 -airflow-core/src/airflow/example_dags/example_setup_teardown_taskflow.py::1 -airflow-core/src/airflow/example_dags/example_simplest_dag.py::1 -airflow-core/src/airflow/example_dags/example_skip_dag.py::3 -airflow-core/src/airflow/example_dags/example_task_group.py::1 -airflow-core/src/airflow/example_dags/example_task_group_decorator.py::1 -airflow-core/src/airflow/example_dags/example_task_state_store.py::2 -airflow-core/src/airflow/example_dags/example_task_state_store_mapped.py::1 -airflow-core/src/airflow/example_dags/example_time_delta_sensor_async.py::1 -airflow-core/src/airflow/example_dags/example_trigger_target_dag.py::1 -airflow-core/src/airflow/example_dags/example_workday_timetable.py::1 -airflow-core/src/airflow/example_dags/example_xcom.py::1 -airflow-core/src/airflow/example_dags/example_xcomargs.py::1 -airflow-core/src/airflow/example_dags/plugins/event_listener.py::1 -airflow-core/src/airflow/example_dags/tutorial.py::1 -airflow-core/src/airflow/example_dags/tutorial_dag.py::1 -airflow-core/src/airflow/example_dags/tutorial_objectstorage.py::1 -airflow-core/src/airflow/example_dags/tutorial_taskflow_api.py::1 -airflow-core/src/airflow/example_dags/tutorial_taskflow_api_virtualenv.py::1 -airflow-core/src/airflow/example_dags/tutorial_taskflow_templates.py::1 airflow-core/src/airflow/exceptions.py::1 airflow-core/src/airflow/executors/base_executor.py::3 airflow-core/src/airflow/jobs/triggerer_job_runner.py::18 diff --git a/scripts/ci/prek/check_sdk_imports_in_core.py b/scripts/ci/prek/check_sdk_imports_in_core.py index c07a025fd6249..7dca1b6be6c85 100755 --- a/scripts/ci/prek/check_sdk_imports_in_core.py +++ b/scripts/ci/prek/check_sdk_imports_in_core.py @@ -28,8 +28,13 @@ as ``relative/path::N`` entries (one per file), where ``N`` is the maximum number of ``airflow.sdk`` import statements allowed in that file. A file whose current count exceeds the recorded limit is treated as a violation -- core -should not gain new runtime dependencies on the Task SDK; use ``# noqa: SDK001`` -for a one off exception, or regenerate the allowlist for a deliberate one. +should not gain new runtime dependencies on the Task SDK. Two escape hatches +exist: ``# noqa: SDK001`` on a single import line, which keeps that line out of +the count and out of the allowlist entirely, or raising the file's recorded +count via ``--generate FILE`` when a file genuinely needs several. + +``airflow-core/src/airflow/example_dags`` is skipped entirely: those files are +Dags authored with the Task SDK, not core runtime code. Modes ----- @@ -47,9 +52,11 @@ Remove entries for files that no longer exist. Safe to run at any time; does not add new entries or raise limits. -``--generate``: - Scan ``airflow-core/src/airflow`` and *rebuild* the allowlist from scratch. - Intended for the initial setup or after a large-scale clean-up sprint. +``--generate [FILE ...]``: + With files, record the current counts for just those files and leave every + other entry alone -- the usual way to approve a deliberate new import. + With no files, *rebuild* the allowlist from scratch; that also approves any + unrelated drift, so it is meant for the initial setup or a clean-up sprint. """ from __future__ import annotations @@ -68,6 +75,10 @@ NOCHECK_CODE = "SDK001" +# Dags are authored *with* the task SDK, so example Dags shipped in core are not core code +# for the purpose of this check and never take part in the ratchet. +EXCLUDED_DIR_NAMES = frozenset({"example_dags"}) + def check_file_for_sdk_imports(file_path: Path) -> list[tuple[int, str]]: """Check file for airflow.sdk imports. Returns list of (line_num, import_statement).""" @@ -82,23 +93,66 @@ class SdkImportsAllowlistManager(AllowlistManager): def __init__(self, allowlist_file: Path) -> None: super().__init__(allowlist_file, repo_root=REPO_ROOT) + def is_excluded(self, path: Path) -> bool: + return not EXCLUDED_DIR_NAMES.isdisjoint(path.parts) + def iter_files(self) -> Iterable[Path]: - return CORE_SRC_ROOT.rglob("*.py") + return (path for path in CORE_SRC_ROOT.rglob("*.py") if not self.is_excluded(path)) + + def check(self, files: list[Path], allowlist: dict[str, int]) -> int: + return super().check([path for path in files if not self.is_excluded(path)], allowlist) def count_occurrences(self, path: Path) -> int: return len(check_file_for_sdk_imports(path)) + def generate_for(self, files: list[Path]) -> int: + """Rewrite the allowlist entries for *files* only, leaving every other entry untouched.""" + allowlist = self.load() + updated: list[tuple[str, int, int]] = [] + + for path in files: + if self.is_excluded(path) or not path.exists() or path.suffix != ".py": + continue + try: + rel = str(path.resolve().relative_to(self.repo_root.resolve())) + except ValueError: + console.print(f"[yellow]Skipping file outside the repository:[/yellow] {path}") + continue + before = allowlist.get(rel, 0) + after = self.count_occurrences(path) + if before == after: + continue + if after: + allowlist[rel] = after + else: + del allowlist[rel] + updated.append((rel, before, after)) + + if not updated: + console.print("[green]Allowlist already up to date for the given files.[/green]") + return 0 + + self.save(allowlist) + console.print( + f"[green]Updated[/green] [cyan]{self.allowlist_file.relative_to(self.repo_root)}[/cyan]:" + ) + for rel, before, after in updated: + console.print(f" [cyan]{rel}[/cyan] {before} -> {after}") + return 0 + def violation_panel_text(self) -> str: return ( "New [bold]airflow.sdk[/bold] import detected in airflow-core.\n" "Core (scheduler/API server) should not gain new runtime dependencies " - "on the Task SDK.\n" - "If this import is a genuine one-off, append `# noqa: SDK001` to the " - "import line.\n" - "If it's an intentional, broader exception, run:\n\n" - " [cyan]uv run ./scripts/ci/prek/check_sdk_imports_in_core.py --generate[/cyan]\n\n" - "to regenerate the allowlist, then commit the updated\n" - "[cyan]generated/known_sdk_imports_in_core.txt[/cyan]." + "on the Task SDK.\n\n" + "The preferred fix is to drop the import. If it has to stay, pick one:\n\n" + "[bold]1. A single deliberate import line[/bold] -- append `# noqa: SDK001` to it.\n" + " The line is then invisible to this check and never reaches the allowlist.\n\n" + "[bold]2. A file that genuinely needs several[/bold] -- record its new count:\n\n" + " [cyan]uv run ./scripts/ci/prek/check_sdk_imports_in_core.py --generate FILE ...[/cyan]\n\n" + " then commit the updated [cyan]generated/known_sdk_imports_in_core.txt[/cyan].\n" + " Passing no FILE rebuilds the whole allowlist and approves any unrelated\n" + " drift with it, so prefer the scoped form." ) @@ -122,14 +176,18 @@ def main(argv: list[str] | None = None) -> int: parser.add_argument( "--generate", action="store_true", - help="Regenerate the allowlist from the current codebase and exit", + help=( + "Record the current counts for the given FILEs and exit; " + "with no FILE, rebuild the whole allowlist from the current codebase" + ), ) args = parser.parse_args(argv) manager = SdkImportsAllowlistManager(REPO_ROOT / "generated" / "known_sdk_imports_in_core.txt") + files = [Path(f).resolve() for f in args.files] if args.generate: - return manager.generate() + return manager.generate_for(files) if files else manager.generate() if args.cleanup: return manager.cleanup() @@ -139,13 +197,13 @@ def main(argv: list[str] | None = None) -> int: if args.all_files: return manager.check(list(manager.iter_files()), allowlist) - if not args.files: + if not files: console.print( "[yellow]No files provided. Pass filenames or use --all-files to scan the whole repo.[/yellow]" ) return 0 - return manager.check([Path(f).resolve() for f in args.files], allowlist) + return manager.check(files, allowlist) if __name__ == "__main__": diff --git a/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py b/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py index c3c6a90a2cdb5..702744e1bd68a 100644 --- a/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py +++ b/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py @@ -273,6 +273,17 @@ def test_noqa_marker_avoids_ratchet_entirely(self, create_fake_core_repo, tmp_pa manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") assert manager.check([path], {}) == 0 + def test_example_dags_are_not_checked(self, create_fake_core_repo, tmp_path): + path = create_fake_core_repo( + "example_dags/example_simplest_dag.py", + """\ + from airflow.sdk import DAG + from airflow.sdk import task + """, + ) + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + assert manager.check([path], {}) == 0 + def test_non_python_file_is_skipped(self, create_fake_core_repo, tmp_path): path = create_fake_core_repo( "models/not_python.txt", @@ -309,3 +320,57 @@ def test_generate_records_current_occurrences(self, create_fake_core_repo, tmp_p manager = SdkImportsAllowlistManager(allowlist_path) assert manager.generate() == 0 assert manager.load() == {"airflow-core/src/airflow/models/a.py": 1} + + def test_main_generate_with_files_is_scoped(self, create_fake_core_repo, tmp_path): + path = create_fake_core_repo("models/a.py", "from airflow.sdk import DAG\n") + (tmp_path / "generated").mkdir() + allowlist_path = tmp_path / "generated" / "known_sdk_imports_in_core.txt" + allowlist_path.write_text("airflow-core/src/airflow/models/other.py::2\n") + + assert hook.main(["--generate", str(path)]) == 0 + assert allowlist_path.read_text().splitlines() == [ + "airflow-core/src/airflow/models/a.py::1", + "airflow-core/src/airflow/models/other.py::2", + ] + + def test_generate_for_leaves_other_entries_untouched(self, create_fake_core_repo, tmp_path): + target = create_fake_core_repo( + "models/target.py", + """\ + from airflow.sdk import DAG + from airflow.sdk import task + """, + ) + create_fake_core_repo("models/drifted.py", "from airflow.sdk import DAG\n") + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + manager.save({"airflow-core/src/airflow/models/target.py": 1}) + + assert manager.generate_for([target]) == 0 + assert manager.load() == {"airflow-core/src/airflow/models/target.py": 2} + + def test_generate_for_drops_entry_when_file_is_clean(self, create_fake_core_repo, tmp_path): + path = create_fake_core_repo("models/fixed.py", "from airflow.models import DagRun\n") + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + manager.save( + { + "airflow-core/src/airflow/models/fixed.py": 1, + "airflow-core/src/airflow/models/other.py": 3, + } + ) + + assert manager.generate_for([path]) == 0 + assert manager.load() == {"airflow-core/src/airflow/models/other.py": 3} + + def test_generate_for_skips_example_dags(self, create_fake_core_repo, tmp_path): + path = create_fake_core_repo("example_dags/tutorial.py", "from airflow.sdk import DAG\n") + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + + assert manager.generate_for([path]) == 0 + assert manager.load() == {} + + def test_generate_skips_example_dags(self, create_fake_core_repo, tmp_path): + create_fake_core_repo("models/a.py", "from airflow.sdk import DAG\n") + create_fake_core_repo("example_dags/tutorial.py", "from airflow.sdk import DAG\n") + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + assert manager.generate() == 0 + assert manager.load() == {"airflow-core/src/airflow/models/a.py": 1} From 578c487edaf68fd30859779975b90fc4c675df78 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Mon, 10 Aug 2026 13:04:35 +0530 Subject: [PATCH 3/3] kaxil review --- scripts/ci/prek/check_sdk_imports_in_core.py | 36 +++++++++++++++++-- .../ci/prek/test_check_sdk_imports_in_core.py | 34 ++++++++++++++++++ 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/scripts/ci/prek/check_sdk_imports_in_core.py b/scripts/ci/prek/check_sdk_imports_in_core.py index 7dca1b6be6c85..aea9446305ba2 100755 --- a/scripts/ci/prek/check_sdk_imports_in_core.py +++ b/scripts/ci/prek/check_sdk_imports_in_core.py @@ -62,6 +62,7 @@ from __future__ import annotations import argparse +import ast from collections.abc import Iterable from pathlib import Path @@ -89,6 +90,21 @@ def check_file_for_sdk_imports(file_path: Path) -> list[tuple[int, str]]: ) +def _can_parse(path: Path) -> bool: + """Whether *path* is valid enough Python for its import count to be trustworthy. + + ``find_import_violations`` swallows parse failures and reports zero violations for + them, which is indistinguishable from a file that genuinely has none. Callers must + filter those files out *before* counting, so a file mid-edit with a syntax error + never looks like a legitimate drop in its allowlist entry. + """ + try: + ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + except (OSError, UnicodeDecodeError, SyntaxError): + return False + return True + + class SdkImportsAllowlistManager(AllowlistManager): def __init__(self, allowlist_file: Path) -> None: super().__init__(allowlist_file, repo_root=REPO_ROOT) @@ -97,10 +113,23 @@ def is_excluded(self, path: Path) -> bool: return not EXCLUDED_DIR_NAMES.isdisjoint(path.parts) def iter_files(self) -> Iterable[Path]: - return (path for path in CORE_SRC_ROOT.rglob("*.py") if not self.is_excluded(path)) + return ( + path for path in CORE_SRC_ROOT.rglob("*.py") if not self.is_excluded(path) and _can_parse(path) + ) def check(self, files: list[Path], allowlist: dict[str, int]) -> int: - return super().check([path for path in files if not self.is_excluded(path)], allowlist) + checkable = [] + for path in files: + if self.is_excluded(path): + continue + if path.exists() and path.suffix == ".py" and not _can_parse(path): + console.print( + f"[yellow]Skipping unparsable file (not counted, allowlist entry left " + f"untouched):[/yellow] {path}" + ) + continue + checkable.append(path) + return super().check(checkable, allowlist) def count_occurrences(self, path: Path) -> int: return len(check_file_for_sdk_imports(path)) @@ -113,6 +142,9 @@ def generate_for(self, files: list[Path]) -> int: for path in files: if self.is_excluded(path) or not path.exists() or path.suffix != ".py": continue + if not _can_parse(path): + console.print(f"[yellow]Skipping unparsable file:[/yellow] {path}") + continue try: rel = str(path.resolve().relative_to(self.repo_root.resolve())) except ValueError: diff --git a/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py b/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py index 702744e1bd68a..3a430c3c2c116 100644 --- a/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py +++ b/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py @@ -251,6 +251,26 @@ def test_reducing_violations_tightens_allowlist(self, create_fake_core_repo, tmp assert manager.check([path], allowlist) == 1 assert manager.load() == {"airflow-core/src/airflow/models/improved.py": 1} + def test_unparseable_file_does_not_tighten_allowlist(self, create_fake_core_repo, tmp_path): + """A syntax error must not look like a legitimate drop to zero imports. + + ``find_import_violations`` reports zero violations for a file it cannot parse, + indistinguishable from a file that genuinely has none -- the allowlist entry must + be left alone rather than tightened or deleted. + """ + path = create_fake_core_repo( + "models/broken.py", + """\ + from airflow.sdk import DAG + def broken( + """, + ) + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + manager.save({"airflow-core/src/airflow/models/broken.py": 1}) + + assert manager.check([path], manager.load()) == 0 + assert manager.load() == {"airflow-core/src/airflow/models/broken.py": 1} + def test_fixing_all_violations_removes_entry(self, create_fake_core_repo, tmp_path): path = create_fake_core_repo( "models/fixed.py", @@ -361,6 +381,20 @@ def test_generate_for_drops_entry_when_file_is_clean(self, create_fake_core_repo assert manager.generate_for([path]) == 0 assert manager.load() == {"airflow-core/src/airflow/models/other.py": 3} + def test_generate_for_leaves_unparseable_file_untouched(self, create_fake_core_repo, tmp_path): + path = create_fake_core_repo( + "models/broken.py", + """\ + from airflow.sdk import DAG + def broken( + """, + ) + manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt") + manager.save({"airflow-core/src/airflow/models/broken.py": 1}) + + assert manager.generate_for([path]) == 0 + assert manager.load() == {"airflow-core/src/airflow/models/broken.py": 1} + def test_generate_for_skips_example_dags(self, create_fake_core_repo, tmp_path): path = create_fake_core_repo("example_dags/tutorial.py", "from airflow.sdk import DAG\n") manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")