(airflow) branch main updated: Add version suffix if it is set during image builds (#36253)

[ephraimanierobi]

This is an automated email from the ASF dual-hosted git repository.

ephraimanierobi 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 a68b4194fe Add version suffix if it is set during image builds (#36253)
a68b4194fe is described below

commit a68b4194fe7201bba0544856b60c7d6724da60b3
Author: Jarek Potiuk <ja...@potiuk.com>
AuthorDate: Sat Dec 16 09:16:49 2023 +0100

    Add version suffix if it is set during image builds (#36253)
    
    When we are building CI/PROD images, AIRFLOW_VERSION arg is retrieved
    from the current Airflow version set in version.py. This is find for
    main build when we are running image build tests, because there, the
    version contains `.dev0` and when we extend the image we can use
    `pip install airflow-version=${AIRFLOW_VERSION}`. However in release
    builds, the version in `version.py` does not contain version suffix,
    and this version of Airflow is not released yet.
    
    This PR fixes it by checking if version_suffix_for_pypi is set,
    and it case it is and airflow version does not contain it, we
    will add the suffix automatically.
---
 dev/breeze/src/airflow_breeze/params/build_ci_params.py    |  3 +--
 dev/breeze/src/airflow_breeze/params/build_prod_params.py  |  3 +--
 .../src/airflow_breeze/params/common_build_params.py       | 14 ++++++++++++++
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/dev/breeze/src/airflow_breeze/params/build_ci_params.py 
b/dev/breeze/src/airflow_breeze/params/build_ci_params.py
index e5322309c5..954f405f9c 100644
--- a/dev/breeze/src/airflow_breeze/params/build_ci_params.py
+++ b/dev/breeze/src/airflow_breeze/params/build_ci_params.py
@@ -21,7 +21,6 @@ from dataclasses import dataclass
 from pathlib import Path
 
 from airflow_breeze.branch_defaults import DEFAULT_AIRFLOW_CONSTRAINTS_BRANCH
-from airflow_breeze.global_constants import get_airflow_version
 from airflow_breeze.params.common_build_params import CommonBuildParams
 from airflow_breeze.utils.path_utils import BUILD_CACHE_DIR
 
@@ -46,7 +45,7 @@ class BuildCiParams(CommonBuildParams):
 
     @property
     def airflow_version(self):
-        return get_airflow_version()
+        return self._get_version_with_suffix()
 
     @property
     def image_type(self) -> str:
diff --git a/dev/breeze/src/airflow_breeze/params/build_prod_params.py 
b/dev/breeze/src/airflow_breeze/params/build_prod_params.py
index 8a534cc425..ed0cec9a8c 100644
--- a/dev/breeze/src/airflow_breeze/params/build_prod_params.py
+++ b/dev/breeze/src/airflow_breeze/params/build_prod_params.py
@@ -26,7 +26,6 @@ from airflow_breeze.global_constants import (
     AIRFLOW_SOURCES_FROM,
     AIRFLOW_SOURCES_TO,
     get_airflow_extras,
-    get_airflow_version,
 )
 from airflow_breeze.params.common_build_params import CommonBuildParams
 from airflow_breeze.utils.console import get_console
@@ -62,7 +61,7 @@ class BuildProdParams(CommonBuildParams):
         if self.install_airflow_version:
             return self.install_airflow_version
         else:
-            return get_airflow_version()
+            return self._get_version_with_suffix()
 
     @property
     def image_type(self) -> str:
diff --git a/dev/breeze/src/airflow_breeze/params/common_build_params.py 
b/dev/breeze/src/airflow_breeze/params/common_build_params.py
index d9bd3d3a1f..e1c77990e9 100644
--- a/dev/breeze/src/airflow_breeze/params/common_build_params.py
+++ b/dev/breeze/src/airflow_breeze/params/common_build_params.py
@@ -28,6 +28,7 @@ from airflow_breeze.global_constants import (
     ALLOWED_INSTALL_MYSQL_CLIENT_TYPES,
     APACHE_AIRFLOW_GITHUB_REPOSITORY,
     DOCKER_DEFAULT_PLATFORM,
+    get_airflow_version,
 )
 from airflow_breeze.utils.console import get_console
 from airflow_breeze.utils.platforms import get_real_platform
@@ -196,3 +197,16 @@ class CommonBuildParams:
         for arg in self.build_arg_values:
             build_args.extend(["--build-arg", arg])
         return build_args
+
+    def _get_version_with_suffix(self) -> str:
+        from packaging.version import Version
+
+        airflow_version = get_airflow_version()
+        try:
+            if self.version_suffix_for_pypi and self.version_suffix_for_pypi 
not in airflow_version:
+                version = Version(airflow_version)
+                return version.base_version + 
f".{self.version_suffix_for_pypi}"
+        except Exception:
+            # in case of any failure just fall back to the original version set
+            pass
+        return airflow_version