[v3-1-test] Add pre-commit hook to verify version consistency (#59509)

.pre-commit-config.yaml

@@ -180,6 +180,18 @@ repos:
         entry: ./scripts/ci/prek/check_min_python_version.py
         language: python
         require_serial: true
+      - id: check-version-consistency
+        name: Check version consistency
+        entry: ./scripts/ci/prek/check_version_consistency.py
+        language: python
+        files: >
+          (?x)
+          ^airflow-core/src/airflow/__init__\.py$|
+          ^airflow-core/pyproject\.toml$|
+          ^task-sdk/src/airflow/sdk/__init__\.py$|
+          ^pyproject\.toml$
+        pass_filenames: false
+        require_serial: true
       - id: upgrade-important-versions
         name: Upgrade important versions (manual)
         entry: ./scripts/ci/prek/upgrade_important_versions.py

scripts/ci/prek/check_version_consistency.py

@@ -0,0 +1,302 @@
+#!/usr/bin/env python
+# 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.
+# /// script
+# requires-python = ">=3.10,<3.11"
+# dependencies = [
+#   "rich>=13.6.0",
+#   "packaging>=25.0",
+#   "tomli>=2.0.1",
+# ]
+# ///
+from __future__ import annotations
+
+import ast
+import re
+import sys
+from pathlib import Path
+
+try:
+    import tomllib
+except ImportError:
+    import tomli as tomllib
+
+from packaging.specifiers import SpecifierSet
+from packaging.version import Version
+
+sys.path.insert(0, str(Path(__file__).parent.resolve()))
+
+from common_prek_utils import (
+    AIRFLOW_CORE_SOURCES_PATH,
+    AIRFLOW_ROOT_PATH,
+    AIRFLOW_TASK_SDK_SOURCES_PATH,
+    console,
+)
+
+
+def read_airflow_version() -> str:
+    """Read Airflow version from airflow-core/src/airflow/__init__.py"""
+    ast_obj = ast.parse((AIRFLOW_CORE_SOURCES_PATH / "airflow" / "__init__.py").read_text())
+    for node in ast_obj.body:
+        if isinstance(node, ast.Assign):
+            if node.targets[0].id == "__version__":  # type: ignore[attr-defined]
+                return ast.literal_eval(node.value)
+
+    raise RuntimeError("Couldn't find __version__ in airflow-core/src/airflow/__init__.py")
+
+
+def read_task_sdk_version() -> str:
+    """Read Task SDK version from task-sdk/src/airflow/sdk/__init__.py"""
+    ast_obj = ast.parse((AIRFLOW_TASK_SDK_SOURCES_PATH / "airflow" / "sdk" / "__init__.py").read_text())
+    for node in ast_obj.body:
+        if isinstance(node, ast.Assign):
+            if node.targets[0].id == "__version__":  # type: ignore[attr-defined]
+                return ast.literal_eval(node.value)
+
+    raise RuntimeError("Couldn't find __version__ in task-sdk/src/airflow/sdk/__init__.py")
+
+
+def read_airflow_version_from_pyproject() -> str:
+    """Read Airflow version from airflow-core/pyproject.toml"""
+    pyproject_path = AIRFLOW_ROOT_PATH / "airflow-core" / "pyproject.toml"
+    with pyproject_path.open("rb") as f:
+        data = tomllib.load(f)
+    version = data.get("project", {}).get("version")
+    if not version:
+        raise RuntimeError("Couldn't find version in airflow-core/pyproject.toml")
+    return str(version)
+
+
+def read_task_sdk_dependency_constraint() -> str:
+    """Read Task SDK dependency constraint from airflow-core/pyproject.toml"""
+    pyproject_path = AIRFLOW_ROOT_PATH / "airflow-core" / "pyproject.toml"
+    with pyproject_path.open("rb") as f:
+        data = tomllib.load(f)
+    dependencies = data.get("project", {}).get("dependencies", [])
+    for dep in dependencies:
+        # Extract package name (everything before >=, >, <, ==, etc.)
+        package_name = re.split(r"[<>=!]", dep)[0].strip().strip("\"'")
+        if package_name == "apache-airflow-task-sdk":
+            # Extract the version constraint part
+            constraint_match = re.search(r"apache-airflow-task-sdk\s*(.*)", dep)
+            if constraint_match:
+                return constraint_match.group(1).strip().strip("\"'")
+            return dep
+    raise RuntimeError("Couldn't find apache-airflow-task-sdk dependency in airflow-core/pyproject.toml")
+
+
+def read_root_pyproject_version() -> str:
+    """Read version from root pyproject.toml"""
+    pyproject_path = AIRFLOW_ROOT_PATH / "pyproject.toml"
+    with pyproject_path.open("rb") as f:
+        data = tomllib.load(f)
+    version = data.get("project", {}).get("version")
+    if not version:
+        raise RuntimeError("Couldn't find version in pyproject.toml")
+    return str(version)
+
+
+def read_root_airflow_core_dependency() -> str:
+    """Read apache-airflow-core dependency from root pyproject.toml"""
+    pyproject_path = AIRFLOW_ROOT_PATH / "pyproject.toml"
+    with pyproject_path.open("rb") as f:
+        data = tomllib.load(f)
+    dependencies = data.get("project", {}).get("dependencies", [])
+    for dep in dependencies:
+        # Extract package name (everything before >=, >, <, ==, etc.)
+        package_name = re.split(r"[<>=!]", dep)[0].strip().strip("\"'")
+        if package_name == "apache-airflow-core":
+            # Extract the version constraint part
+            constraint_match = re.search(r"apache-airflow-core\s*(.*)", dep)
+            if constraint_match:
+                return constraint_match.group(1).strip().strip("\"'")
+            return dep
+    raise RuntimeError("Couldn't find apache-airflow-core dependency in pyproject.toml")
+
+
+def read_root_task_sdk_dependency_constraint() -> str:
+    """Read Task SDK dependency constraint from root pyproject.toml"""
+    pyproject_path = AIRFLOW_ROOT_PATH / "pyproject.toml"
+    with pyproject_path.open("rb") as f:
+        data = tomllib.load(f)
+    dependencies = data.get("project", {}).get("dependencies", [])
+    for dep in dependencies:
+        # Extract package name (everything before >=, >, <, ==, etc.)
+        package_name = re.split(r"[<>=!]", dep)[0].strip().strip("\"'")
+        if package_name == "apache-airflow-task-sdk":
+            # Extract the version constraint part
+            constraint_match = re.search(r"apache-airflow-task-sdk\s*(.*)", dep)
+            if constraint_match:
+                return constraint_match.group(1).strip().strip("\"'")
+            return dep
+    raise RuntimeError("Couldn't find apache-airflow-task-sdk dependency in pyproject.toml")
+
+
+def check_version_in_constraint(version: str, constraint: str) -> bool:
+    """Check if version satisfies the constraint"""
+    try:
+        spec = SpecifierSet(constraint)
+        return Version(version) in spec
+    except Exception as e:
+        console.print(f"[red]Error parsing constraint '{constraint}': {e}[/red]")
+        return False
+
+
+def get_exact_version_from_constraint(constraint: str) -> str | None:
+    """
+    Extract the exact version from a constraint string (== operator).
+    Returns the version if found, or None if not found.
+    """
+    try:
+        spec = SpecifierSet(constraint)
+        exact_version = None
+
+        for specifier in spec:
+            if specifier.operator == "==":
+                exact_version = specifier.version
+                break
+
+        return exact_version
+    except Exception:
+        return None
+
+
+def check_constraint_matches_version(version: str, constraint: str) -> tuple[bool, str | None]:
+    """
+    Check if the constraint has an exact match (==) that matches the actual version.
+    Returns (is_match, exact_version_from_constraint)
+    """
+    exact_version = get_exact_version_from_constraint(constraint)
+    if exact_version is None:
+        return (False, None)
+
+    # Check if the exact version in the constraint matches the actual version
+    return (Version(exact_version) == Version(version), exact_version)
+
+
+def main():
+    errors: list[str] = []
+
+    # Read versions
+    try:
+        airflow_version_init = read_airflow_version()
+        airflow_version_pyproject = read_airflow_version_from_pyproject()
+        root_pyproject_version = read_root_pyproject_version()
+        root_airflow_core_dep = read_root_airflow_core_dependency()
+        task_sdk_version_init = read_task_sdk_version()
+        task_sdk_constraint = read_task_sdk_dependency_constraint()
+        root_task_sdk_constraint = read_root_task_sdk_dependency_constraint()
+    except Exception as e:
+        console.print(f"[red]Error reading versions: {e}[/red]")
+        sys.exit(1)
+
+    # Check Airflow version consistency
+    if airflow_version_init != airflow_version_pyproject:
+        errors.append(
+            f"Airflow version mismatch:\n"
+            f"  airflow-core/src/airflow/__init__.py: {airflow_version_init}\n"
+            f"  airflow-core/pyproject.toml: {airflow_version_pyproject}"
+        )
+
+    # Check root pyproject.toml version matches Airflow version
+    if airflow_version_init != root_pyproject_version:
+        errors.append(
+            f"Root pyproject.toml version mismatch:\n"
+            f"  airflow-core/src/airflow/__init__.py: {airflow_version_init}\n"
+            f"  pyproject.toml: {root_pyproject_version}"
+        )
+
+    # Check root pyproject.toml apache-airflow-core dependency matches Airflow version exactly
+    expected_core_dep = f"=={airflow_version_init}"
+    if root_airflow_core_dep != expected_core_dep:
+        errors.append(
+            f"Root pyproject.toml apache-airflow-core dependency mismatch:\n"
+            f"  Expected: apache-airflow-core=={airflow_version_init}\n"
+            f"  Found: apache-airflow-core{root_airflow_core_dep}"
+        )
+
+    # Check Task SDK version is within constraint in airflow-core/pyproject.toml
+    if not check_version_in_constraint(task_sdk_version_init, task_sdk_constraint):
+        errors.append(
+            f"Task SDK version does not satisfy constraint in airflow-core/pyproject.toml:\n"
+            f"  task-sdk/src/airflow/sdk/__init__.py: {task_sdk_version_init}\n"
+            f"  airflow-core/pyproject.toml constraint: apache-airflow-task-sdk{task_sdk_constraint}"
+        )
+
+    # Check Task SDK constraint exact version matches actual version in airflow-core/pyproject.toml
+    constraint_matches, exact_version = check_constraint_matches_version(
+        task_sdk_version_init, task_sdk_constraint
+    )
+    if not constraint_matches:
+        errors.append(
+            f"Task SDK constraint exact version does not match actual version in airflow-core/pyproject.toml:\n"
+            f"  task-sdk/src/airflow/sdk/__init__.py: {task_sdk_version_init}\n"
+            f"  airflow-core/pyproject.toml constraint exact: {exact_version}\n"
+            f"  Expected constraint to have exact version: == {task_sdk_version_init}"
+        )
+
+    # Check Task SDK version is within constraint in root pyproject.toml
+    if not check_version_in_constraint(task_sdk_version_init, root_task_sdk_constraint):
+        errors.append(
+            f"Task SDK version does not satisfy constraint in pyproject.toml:\n"
+            f"  task-sdk/src/airflow/sdk/__init__.py: {task_sdk_version_init}\n"
+            f"  pyproject.toml constraint: apache-airflow-task-sdk{root_task_sdk_constraint}"
+        )
+
+    # Check Task SDK constraint exact version matches actual version in root pyproject.toml
+    root_constraint_matches, root_exact_version = check_constraint_matches_version(
+        task_sdk_version_init, root_task_sdk_constraint
+    )
+    if not root_constraint_matches:
+        errors.append(
+            f"Task SDK constraint exact version does not match actual version in pyproject.toml:\n"
+            f"  task-sdk/src/airflow/sdk/__init__.py: {task_sdk_version_init}\n"
+            f"  pyproject.toml constraint exact: {root_exact_version}\n"
+            f"  Expected constraint to have exact version: == {task_sdk_version_init}"
+        )
+
+    # Verify constraints match between airflow-core and root pyproject.toml
+    # Compare the exact versions extracted from constraints rather than raw strings
+    airflow_core_exact = get_exact_version_from_constraint(task_sdk_constraint)
+    root_exact = get_exact_version_from_constraint(root_task_sdk_constraint)
+    if airflow_core_exact != root_exact:
+        errors.append(
+            f"Task SDK constraint exact version mismatch between pyproject.toml files:\n"
+            f"  airflow-core/pyproject.toml: apache-airflow-task-sdk{task_sdk_constraint} (exact: {airflow_core_exact})\n"
+            f"  pyproject.toml: apache-airflow-task-sdk{root_task_sdk_constraint} (exact: {root_exact})"
+        )
+
+    # Report results
+    if errors:
+        console.print("[red]Version consistency check failed:[/red]\n")
+        for error in errors:
+            console.print(f"[red]{error}[/red]\n")
+        console.print(
+            "[yellow]Please ensure versions are consistent:\n"
+            "  1. Set the Airflow version in airflow-core/src/airflow/__init__.py\n"
+            "  2. Set the Airflow version in airflow-core/pyproject.toml\n"
+            "  3. Set the Airflow version in pyproject.toml\n"
+            "  4. Set apache-airflow-core==<version> in pyproject.toml dependencies\n"
+            "  5. Set the Task SDK version in task-sdk/src/airflow/sdk/__init__.py\n"
+            "  6. Update the Task SDK version constraint in airflow-core/pyproject.toml to include the Task SDK version\n"
+            "  7. Update the Task SDK version constraint in pyproject.toml to include the Task SDK version[/yellow]"
+        )
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()