From f5c86edb7967102d26435018e367354735044f56 Mon Sep 17 00:00:00 2001
From: Jarek Potiuk <jarek@potiuk.com>
Date: Wed, 8 May 2024 00:09:32 +0200
Subject: [PATCH] Add hatch_build.py to k8s test venv cache calculation
 (#39473)

With dynamic project depdencies, the requirements for airlfow are
calculated dymamically from hatch_build.py, however cache invalidation
key for k8s tests only included pyproject.toml so when colorlog has
been bumped in @39453, main build continued to use cache from old
colorlog.

This PR adds also hatch_build.py to cache id calculation
---
 .github/workflows/k8s-tests.yml               |  2 +-
 .../airflow_breeze/utils/kubernetes_utils.py  | 30 ++++++++++++-------
 2 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/.github/workflows/k8s-tests.yml b/.github/workflows/k8s-tests.yml
index 55b54131c4c2a..c4b72a9afc924 100644
--- a/.github/workflows/k8s-tests.yml
+++ b/.github/workflows/k8s-tests.yml
@@ -95,7 +95,7 @@ jobs:
           path: ".build/.k8s-env"
           key: "\
             k8s-env-${{ steps.breeze.outputs.host-python-version }}-\
-            ${{ hashFiles('scripts/ci/kubernetes/k8s_requirements.txt','pyproject.toml') }}"
+            ${{ hashFiles('scripts/ci/kubernetes/k8s_requirements.txt','hatch_build.py') }}"
       - name: Run complete K8S tests ${{ inputs.kubernetes-combos-list-as-string }}
         run: breeze k8s run-complete-tests --run-in-parallel --upgrade --no-copy-local-sources
         env:
diff --git a/dev/breeze/src/airflow_breeze/utils/kubernetes_utils.py b/dev/breeze/src/airflow_breeze/utils/kubernetes_utils.py
index ef94500900024..34c9db0766ea2 100644
--- a/dev/breeze/src/airflow_breeze/utils/kubernetes_utils.py
+++ b/dev/breeze/src/airflow_breeze/utils/kubernetes_utils.py
@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+import hashlib
 import itertools
 import os
 import random
@@ -53,8 +54,9 @@
 HELM_BIN_PATH = K8S_BIN_BASE_PATH / "helm"
 PYTHON_BIN_PATH = K8S_BIN_BASE_PATH / "python"
 SCRIPTS_CI_KUBERNETES_PATH = AIRFLOW_SOURCES_ROOT / "scripts" / "ci" / "kubernetes"
-K8S_REQUIREMENTS = SCRIPTS_CI_KUBERNETES_PATH / "k8s_requirements.txt"
-CACHED_K8S_REQUIREMENTS = K8S_ENV_PATH / "k8s_requirements.txt"
+K8S_REQUIREMENTS_PATH = SCRIPTS_CI_KUBERNETES_PATH / "k8s_requirements.txt"
+HATCH_BUILD_PY_PATH = AIRFLOW_SOURCES_ROOT / "hatch_build.py"
+CACHED_K8S_DEPS_HASH_PATH = K8S_ENV_PATH / "k8s_deps_hash.txt"
 CHART_PATH = AIRFLOW_SOURCES_ROOT / "chart"
 
 # In case of parallel runs those ports will be quickly allocated by multiple threads and closed, which
@@ -273,15 +275,21 @@ def make_sure_kubernetes_tools_are_installed():
         )
 
 
+def _get_k8s_deps_hash():
+    md5_hash = hashlib.md5()
+    content = K8S_REQUIREMENTS_PATH.read_text() + HATCH_BUILD_PY_PATH.read_text()
+    md5_hash.update(content.encode("utf-8"))
+    k8s_deps_hash = md5_hash.hexdigest()
+    return k8s_deps_hash
+
+
 def _requirements_changed() -> bool:
-    if not CACHED_K8S_REQUIREMENTS.exists():
+    if not CACHED_K8S_DEPS_HASH_PATH.exists():
         get_console().print(
             f"\n[warning]The K8S venv in {K8S_ENV_PATH} has never been created. Installing it.\n"
         )
         return True
-    requirements_file_content = K8S_REQUIREMENTS.read_text()
-    cached_requirements_content = CACHED_K8S_REQUIREMENTS.read_text()
-    if cached_requirements_content != requirements_file_content:
+    if CACHED_K8S_DEPS_HASH_PATH.read_text() != _get_k8s_deps_hash():
         get_console().print(
             f"\n[warning]Requirements changed for the K8S venv in {K8S_ENV_PATH}. "
             f"Reinstalling the venv.\n"
@@ -297,7 +305,7 @@ def _install_packages_in_k8s_virtualenv():
         "pip",
         "install",
         "-r",
-        str(K8S_REQUIREMENTS.resolve()),
+        str(K8S_REQUIREMENTS_PATH.resolve()),
     ]
     env = os.environ.copy()
     capture_output = True
@@ -307,7 +315,9 @@ def _install_packages_in_k8s_virtualenv():
         install_command, check=False, capture_output=capture_output, text=True, env=env
     )
     if install_packages_result.returncode != 0:
-        get_console().print(f"[error]Error when installing packages from : {K8S_REQUIREMENTS.resolve()}[/]\n")
+        get_console().print(
+            f"[error]Error when installing packages from : {K8S_REQUIREMENTS_PATH.resolve()}[/]\n"
+        )
         if not get_verbose():
             get_console().print(install_packages_result.stdout)
             get_console().print(install_packages_result.stderr)
@@ -383,9 +393,9 @@ def create_virtualenv(force_venv_setup: bool) -> RunCommandResult:
     install_packages_result = _install_packages_in_k8s_virtualenv()
     if install_packages_result.returncode == 0:
         if get_dry_run():
-            get_console().print(f"[info]Dry run - would be saving {K8S_REQUIREMENTS} to cache")
+            get_console().print(f"[info]Dry run - would be saving {K8S_REQUIREMENTS_PATH} to cache")
         else:
-            CACHED_K8S_REQUIREMENTS.write_text(K8S_REQUIREMENTS.read_text())
+            CACHED_K8S_DEPS_HASH_PATH.write_text(_get_k8s_deps_hash())
     return install_packages_result