Refactor schema version checks to share target-branch and worktree logic

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"

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<YYYY>_<MM>_<DD>.py`` file under:\n"

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:

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"})