diff --git a/scripts/ci/prek/check_execution_api_versions.py b/scripts/ci/prek/check_execution_api_versions.py index 833ce32bcf753..6555cb4170f01 100755 --- a/scripts/ci/prek/check_execution_api_versions.py +++ b/scripts/ci/prek/check_execution_api_versions.py @@ -24,36 +24,21 @@ from __future__ import annotations import json -import os import subprocess import sys -import tempfile from pathlib import Path -from common_prek_utils import console, get_remote_for_main +from common_prek_utils import ( + console, + create_temporary_worktree, + fetch_target_branch, + get_changed_files, +) DATAMODELS_PREFIX = "airflow-core/src/airflow/api_fastapi/execution_api/datamodels/" VERSIONS_PREFIX = "airflow-core/src/airflow/api_fastapi/execution_api/versions/" -def get_target_branch() -> str: - """Branch to compare against. GITHUB_BASE_REF for PRs, DEFAULT_BRANCH in CI, else main.""" - return os.environ.get("GITHUB_BASE_REF") or os.environ.get("DEFAULT_BRANCH") or "main" - - -def get_changed_files(filenames: list[str]) -> list[str]: - """Get changed files. Uses filenames from prek when provided, else staged files for local runs.""" - if filenames: - return filenames - result = subprocess.run( - ["git", "diff", "--cached", "--name-only"], - capture_output=True, - text=True, - check=True, - ) - return [f for f in result.stdout.strip().splitlines() if f] - - def generate_schema(cwd: Path) -> dict: """Generate OpenAPI schema from repo at cwd.""" script_path = Path(__file__).parent / "generate_execution_api_schema.py" @@ -69,20 +54,10 @@ def generate_schema(cwd: Path) -> dict: return json.loads(result.stdout) -def generate_schema_from_main() -> dict: +def generate_schema_from_main(ref: str) -> dict: """Generate schema from target branch using worktree.""" - target_branch = get_target_branch() - remote = get_remote_for_main() - ref = f"{remote}/{target_branch}" - worktree_path = Path(tempfile.mkdtemp()) / "airflow-main" - subprocess.run(["git", "fetch", remote, target_branch], capture_output=True, check=False) - subprocess.run(["git", "worktree", "add", str(worktree_path), ref], capture_output=True, check=True) - try: + with create_temporary_worktree(ref) as worktree_path: return generate_schema(worktree_path) - finally: - subprocess.run( - ["git", "worktree", "remove", "--force", str(worktree_path)], capture_output=True, check=False - ) def normalize_schema(schema: dict) -> dict: @@ -109,8 +84,9 @@ def main() -> int: version_files = [f for f in changed_files if f.startswith(VERSIONS_PREFIX)] if datamodel_files and not version_files: + target_ref = fetch_target_branch() try: - main_schema = generate_schema_from_main() + main_schema = generate_schema_from_main(target_ref) except Exception as e: console.print(f"[yellow]WARNING: Could not generate schema from main: {e}[/]") console.print( @@ -145,10 +121,8 @@ def main() -> int: for f in datamodel_files: console.print(f" - [magenta]{f}[/]") console.print("") - remote = get_remote_for_main() - target_branch = get_target_branch() console.print( - f"Schema diff against [cyan]{remote}/{target_branch}[/] detected differences.\n" + f"Schema diff against [cyan]{target_ref}[/] detected differences.\n" "\n" "Please add or update a version file under:\n" f" [cyan]{VERSIONS_PREFIX}[/]\n" diff --git a/scripts/ci/prek/check_supervisor_schemas_versions.py b/scripts/ci/prek/check_supervisor_schemas_versions.py index 348382a3e6789..16b1c5985ef40 100755 --- a/scripts/ci/prek/check_supervisor_schemas_versions.py +++ b/scripts/ci/prek/check_supervisor_schemas_versions.py @@ -41,13 +41,16 @@ from __future__ import annotations -import os import subprocess import sys -import tempfile from pathlib import Path -from common_prek_utils import console, get_remote_for_main +from common_prek_utils import ( + console, + create_temporary_worktree, + fetch_target_branch, + get_changed_files, +) SUPERVISOR_SCHEMAS_PREFIX = "task-sdk/src/airflow/sdk/execution_time/schema/" VERSIONS_PREFIX = SUPERVISOR_SCHEMAS_PREFIX + "versions/" @@ -57,25 +60,6 @@ DUMP_SCRIPT = Path(__file__).parent / "dump_supervisor_schemas.py" -# TODO: We should consolidte the common logic with check_execution_api_versions.py into common_prek_utils -def get_target_branch() -> str: - """Branch to compare against. GITHUB_BASE_REF for PRs, DEFAULT_BRANCH in CI, else main.""" - return os.environ.get("GITHUB_BASE_REF") or os.environ.get("DEFAULT_BRANCH") or "main" - - -def get_changed_files(filenames: list[str]) -> list[str]: - """Get changed files. Uses filenames from prek when provided, else staged files for local runs.""" - if filenames: - return filenames - result = subprocess.run( - ["git", "diff", "--cached", "--name-only"], - capture_output=True, - text=True, - check=True, - ) - return [f for f in result.stdout.strip().splitlines() if f] - - def dump_snapshot(cwd: Path) -> str: """Run ``dump_supervisor_schemas.py`` in *cwd* and return its stdout.""" result = subprocess.run( @@ -101,12 +85,8 @@ def dump_snapshot(cwd: Path) -> str: return result.stdout -def _upstream_has_schema() -> bool: +def _upstream_has_schema(ref: str) -> bool: """Return True if the target branch carries the schema package.""" - target_branch = get_target_branch() - remote = get_remote_for_main() - ref = f"{remote}/{target_branch}" - subprocess.run(["git", "fetch", remote, target_branch], capture_output=True, check=False) # ``git cat-file -e`` exits zero iff the path exists at the ref. result = subprocess.run( ["git", "cat-file", "-e", f"{ref}:{VERSIONS_PREFIX}__init__.py"], @@ -116,22 +96,10 @@ def _upstream_has_schema() -> bool: return result.returncode == 0 -def dump_snapshot_from_main() -> str: +def dump_snapshot_from_main(ref: str) -> str: """Dump snapshot from target branch using a temporary worktree.""" - target_branch = get_target_branch() - remote = get_remote_for_main() - ref = f"{remote}/{target_branch}" - worktree_path = Path(tempfile.mkdtemp()) / "airflow-main" - subprocess.run(["git", "fetch", remote, target_branch], capture_output=True, check=False) - subprocess.run(["git", "worktree", "add", str(worktree_path), ref], capture_output=True, check=True) - try: + with create_temporary_worktree(ref) as worktree_path: return dump_snapshot(worktree_path) - finally: - subprocess.run( - ["git", "worktree", "remove", "--force", str(worktree_path)], - capture_output=True, - check=False, - ) def main() -> int: @@ -154,7 +122,8 @@ def main() -> int: # Contributor added a version-change entry: trust them. return 0 - if not _upstream_has_schema(): + target_ref = fetch_target_branch() + if not _upstream_has_schema(target_ref): # The package is being introduced in this PR -- nothing on the # target branch to compare against. The check will start firing # normally once the package is on the target branch. @@ -166,7 +135,7 @@ def main() -> int: return 0 try: - main_snapshot = dump_snapshot_from_main() + main_snapshot = dump_snapshot_from_main(target_ref) except Exception as e: console.print(f"[bold red]ERROR:[/] Failed to generate upstream snapshot for comparison: {e}") return 1 @@ -184,10 +153,8 @@ def main() -> int: for f in schema_source_files: console.print(f" - [magenta]{f}[/]") console.print("") - remote = get_remote_for_main() - target_branch = get_target_branch() console.print( - f"Snapshot diff against [cyan]{remote}/{target_branch}[/] detected differences.\n" + f"Snapshot diff against [cyan]{target_ref}[/] detected differences.\n" "\n" "Append a ``VersionChange`` subclass to the in-progress head " "``v__
.py`` file under:\n" diff --git a/scripts/ci/prek/common_prek_utils.py b/scripts/ci/prek/common_prek_utils.py index 62154dedad044..31959935943ba 100644 --- a/scripts/ci/prek/common_prek_utils.py +++ b/scripts/ci/prek/common_prek_utils.py @@ -30,7 +30,7 @@ from collections.abc import Callable, Generator, Iterable from contextlib import contextmanager from pathlib import Path -from tempfile import NamedTemporaryFile, _TemporaryFileWrapper +from tempfile import NamedTemporaryFile, TemporaryDirectory, _TemporaryFileWrapper from typing import Any AIRFLOW_ROOT_PATH = Path(__file__).parents[3].resolve() @@ -869,6 +869,59 @@ def get_remote_for_main() -> str: return apache_remote or origin_remote or "origin" +def get_target_branch() -> str: + """Return the target branch used for comparisons against the upstream repository.""" + return os.environ.get("GITHUB_BASE_REF") or os.environ.get("DEFAULT_BRANCH") or "main" + + +def get_changed_files(filenames: list[str]) -> list[str]: + """Return filenames passed by prek, or staged filenames when invoked directly.""" + if filenames: + return filenames + result = subprocess.run( + ["git", "diff", "--cached", "--name-only"], + capture_output=True, + text=True, + check=True, + ) + return result.stdout.strip().splitlines() + + +def fetch_target_branch() -> str: + """Fetch the target branch and return its remote ref.""" + remote = get_remote_for_main() + target_branch = get_target_branch() + subprocess.run( + ["git", "fetch", remote, target_branch], + capture_output=True, + text=True, + check=False, + ) + return f"{remote}/{target_branch}" + + +@contextmanager +def create_temporary_worktree(ref: str) -> Generator[Path, None, None]: + """Create a temporary worktree for *ref* and remove it on exit.""" + with TemporaryDirectory() as temp_dir: + worktree_path = Path(temp_dir) / "airflow" + subprocess.run( + ["git", "worktree", "add", str(worktree_path), ref], + capture_output=True, + text=True, + check=True, + ) + try: + yield worktree_path + finally: + subprocess.run( + ["git", "worktree", "remove", "--force", str(worktree_path)], + capture_output=True, + text=True, + check=False, + ) + + def env_without_github_tokens(env: dict[str, str] | None = None) -> dict[str, str]: cleaned_env = dict(os.environ if env is None else env) for token_env_var in GITHUB_TOKEN_ENV_VARS: diff --git a/scripts/tests/ci/prek/test_check_execution_api_versions.py b/scripts/tests/ci/prek/test_check_execution_api_versions.py new file mode 100644 index 0000000000000..d1fec692f8307 --- /dev/null +++ b/scripts/tests/ci/prek/test_check_execution_api_versions.py @@ -0,0 +1,57 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from contextlib import nullcontext +from unittest import mock + +import ci.prek.check_execution_api_versions as check_execution_api_versions + + +@mock.patch.object(check_execution_api_versions, "generate_schema", autospec=True) +@mock.patch.object(check_execution_api_versions, "create_temporary_worktree", autospec=True) +def test_generate_schema_from_main_uses_target_ref(create_temporary_worktree, generate_schema, tmp_path): + create_temporary_worktree.return_value = nullcontext(tmp_path) + generate_schema.return_value = {"schema": "main"} + + assert check_execution_api_versions.generate_schema_from_main("upstream/main") == {"schema": "main"} + create_temporary_worktree.assert_called_once_with("upstream/main") + generate_schema.assert_called_once_with(tmp_path) + + +@mock.patch.object(check_execution_api_versions, "schemas_equal", autospec=True, return_value=True) +@mock.patch.object(check_execution_api_versions, "generate_schema", autospec=True) +@mock.patch.object(check_execution_api_versions, "generate_schema_from_main", autospec=True) +@mock.patch.object( + check_execution_api_versions, "fetch_target_branch", autospec=True, return_value="upstream/main" +) +@mock.patch.object(check_execution_api_versions, "get_changed_files", autospec=True) +def test_main_fetches_target_once( + get_changed_files, + fetch_target_branch, + generate_schema_from_main, + generate_schema, + schemas_equal, +): + get_changed_files.return_value = [f"{check_execution_api_versions.DATAMODELS_PREFIX}taskinstance.py"] + generate_schema_from_main.return_value = {"schema": "main"} + generate_schema.return_value = {"schema": "current"} + + assert check_execution_api_versions.main() == 0 + fetch_target_branch.assert_called_once_with() + generate_schema_from_main.assert_called_once_with("upstream/main") + schemas_equal.assert_called_once_with({"schema": "current"}, {"schema": "main"}) diff --git a/scripts/tests/ci/prek/test_check_supervisor_schemas_versions.py b/scripts/tests/ci/prek/test_check_supervisor_schemas_versions.py new file mode 100644 index 0000000000000..21c25f6b89c20 --- /dev/null +++ b/scripts/tests/ci/prek/test_check_supervisor_schemas_versions.py @@ -0,0 +1,73 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import subprocess +from contextlib import nullcontext +from unittest import mock + +import ci.prek.check_supervisor_schemas_versions as check_supervisor_schemas_versions + + +@mock.patch.object(check_supervisor_schemas_versions, "dump_snapshot", autospec=True) +@mock.patch.object(check_supervisor_schemas_versions, "create_temporary_worktree", autospec=True) +def test_dump_snapshot_from_main_uses_target_ref(create_temporary_worktree, dump_snapshot, tmp_path): + create_temporary_worktree.return_value = nullcontext(tmp_path) + dump_snapshot.return_value = "main snapshot" + + assert check_supervisor_schemas_versions.dump_snapshot_from_main("upstream/main") == "main snapshot" + create_temporary_worktree.assert_called_once_with("upstream/main") + dump_snapshot.assert_called_once_with(tmp_path) + + +@mock.patch.object(check_supervisor_schemas_versions.subprocess, "run", autospec=True) +def test_upstream_has_schema_checks_target_ref(run): + run.return_value = subprocess.CompletedProcess(args=["git", "cat-file"], returncode=0) + + assert check_supervisor_schemas_versions._upstream_has_schema("upstream/main") is True + run.assert_called_once_with( + [ + "git", + "cat-file", + "-e", + f"upstream/main:{check_supervisor_schemas_versions.VERSIONS_PREFIX}__init__.py", + ], + capture_output=True, + check=False, + ) + + +@mock.patch.object(check_supervisor_schemas_versions, "dump_snapshot_from_main", autospec=True) +@mock.patch.object( + check_supervisor_schemas_versions, "_upstream_has_schema", autospec=True, return_value=False +) +@mock.patch.object( + check_supervisor_schemas_versions, "fetch_target_branch", autospec=True, return_value="upstream/main" +) +@mock.patch.object(check_supervisor_schemas_versions, "get_changed_files", autospec=True) +def test_main_fetches_target_once( + get_changed_files, + fetch_target_branch, + upstream_has_schema, + dump_snapshot_from_main, +): + get_changed_files.return_value = [check_supervisor_schemas_versions.TASK_SDK_COMMS_PATH] + + assert check_supervisor_schemas_versions.main() == 0 + fetch_target_branch.assert_called_once_with() + upstream_has_schema.assert_called_once_with("upstream/main") + dump_snapshot_from_main.assert_not_called() diff --git a/scripts/tests/ci/prek/test_common_prek_utils.py b/scripts/tests/ci/prek/test_common_prek_utils.py index b169b9b90dc05..27f58417f3563 100644 --- a/scripts/tests/ci/prek/test_common_prek_utils.py +++ b/scripts/tests/ci/prek/test_common_prek_utils.py @@ -18,6 +18,7 @@ import subprocess from pathlib import Path +from unittest import mock import ci.prek.common_prek_utils as common_prek_utils import pytest @@ -176,6 +177,85 @@ def test_no_default_branch_keeps_providers(self, monkeypatch): assert PROVIDERS_AMAZON_S3_PATH in result +class TestSchemaVersionCheckGitUtilities: + @pytest.mark.parametrize( + ("github_base_ref", "default_branch", "expected"), + [ + ("v3-1-test", "main", "v3-1-test"), + (None, "v3-0-test", "v3-0-test"), + (None, None, "main"), + ], + ) + def test_get_target_branch(self, monkeypatch, github_base_ref, default_branch, expected): + for name, value in (("GITHUB_BASE_REF", github_base_ref), ("DEFAULT_BRANCH", default_branch)): + if value is None: + monkeypatch.delenv(name, raising=False) + else: + monkeypatch.setenv(name, value) + + assert common_prek_utils.get_target_branch() == expected + + @mock.patch.object(common_prek_utils.subprocess, "run", autospec=True) + def test_get_changed_files_uses_filenames_from_prek(self, run): + filenames = ["airflow-core/src/airflow/example.py"] + + assert common_prek_utils.get_changed_files(filenames) == filenames + run.assert_not_called() + + @mock.patch.object(common_prek_utils.subprocess, "run", autospec=True) + def test_get_changed_files_uses_staged_files_when_invoked_directly(self, run): + run.return_value = subprocess.CompletedProcess( + args=["git", "diff"], + returncode=0, + stdout="first.py\nsecond.py\n", + ) + + assert common_prek_utils.get_changed_files([]) == ["first.py", "second.py"] + run.assert_called_once_with( + ["git", "diff", "--cached", "--name-only"], + capture_output=True, + text=True, + check=True, + ) + + @mock.patch.object(common_prek_utils, "get_target_branch", autospec=True, return_value="main") + @mock.patch.object(common_prek_utils, "get_remote_for_main", autospec=True, return_value="upstream") + @mock.patch.object(common_prek_utils.subprocess, "run", autospec=True) + def test_fetch_target_branch(self, run, get_remote_for_main, get_target_branch): + assert common_prek_utils.fetch_target_branch() == "upstream/main" + get_remote_for_main.assert_called_once_with() + get_target_branch.assert_called_once_with() + run.assert_called_once_with( + ["git", "fetch", "upstream", "main"], + capture_output=True, + text=True, + check=False, + ) + + @mock.patch.object(common_prek_utils.subprocess, "run", autospec=True) + def test_create_temporary_worktree_removes_worktree_after_error(self, run): + with pytest.raises(RuntimeError, match="schema generation failed"): + with common_prek_utils.create_temporary_worktree("upstream/main") as worktree_path: + temp_dir = worktree_path.parent + raise RuntimeError("schema generation failed") + + assert not temp_dir.exists() + assert run.call_args_list == [ + mock.call( + ["git", "worktree", "add", str(worktree_path), "upstream/main"], + capture_output=True, + text=True, + check=True, + ), + mock.call( + ["git", "worktree", "remove", "--force", str(worktree_path)], + capture_output=True, + text=True, + check=False, + ), + ] + + class TestIsHiddenWithinRoot: @pytest.mark.parametrize( ("relative_path", "expected"),