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..b93815ae6d968 --- /dev/null +++ b/generated/known_sdk_imports_in_core.txt @@ -0,0 +1,44 @@ +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/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..aea9446305ba2 100755 --- a/scripts/ci/prek/check_sdk_imports_in_core.py +++ b/scripts/ci/prek/check_sdk_imports_in_core.py @@ -22,16 +22,64 @@ # "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. 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 +----- +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 [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 import argparse -import sys +import ast +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" +# 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).""" @@ -42,23 +90,153 @@ 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() +def _can_parse(path: Path) -> bool: + """Whether *path* is valid enough Python for its import count to be trustworthy. - if not args.files: - return + ``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 - 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, + +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 ( + 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: + 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)) + + 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 + 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: + 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\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." + ) + + +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=( + "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_for(files) if files else 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 files: + console.print( + "[yellow]No files provided. Pass filenames or use --all-files to scan the whole repo.[/yellow]" + ) + return 0 + + return manager.check(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..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 @@ -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,218 @@ 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_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", + """\ + 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_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", + "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} + + 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_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") + + 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}