This is an automated email from the ASF dual-hosted git repository.
amoghrajesh pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new d8cfc951f7e Convert check-sdk-imports to capture occurences than files
(#71212)
d8cfc951f7e is described below
commit d8cfc951f7edcb15a0473b25a26b9495ff53d3c1
Author: Amogh Desai <[email protected]>
AuthorDate: Tue Aug 11 09:27:06 2026 +0530
Convert check-sdk-imports to capture occurences than files (#71212)
---
airflow-core/.pre-commit-config.yaml | 103 +--------
generated/known_sdk_imports_in_core.txt | 44 ++++
scripts/ci/prek/check_sdk_imports_in_core.py | 210 +++++++++++++++++--
.../ci/prek/test_check_sdk_imports_in_core.py | 232 ++++++++++++++++++++-
4 files changed, 470 insertions(+), 119 deletions(-)
diff --git a/airflow-core/.pre-commit-config.yaml
b/airflow-core/.pre-commit-config.yaml
index 08382f3934f..c59444e234d 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 00000000000..b93815ae6d9
--- /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 03192728930..aea9446305b 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 ed1bcbe057f..3a430c3c2c1 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 @@ from __future__ import annotations
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
+
+
[email protected]
+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 @@ class TestNocheckMarker:
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}
