From 38036ec861c62270c5f9e5543273519699503f1f Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Sat, 16 Dec 2023 09:16:49 +0100 Subject: [PATCH] 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. (cherry picked from commit a68b4194fe7201bba0544856b60c7d6724da60b3) --- .../src/airflow_breeze/params/build_ci_params.py | 3 +-- .../src/airflow_breeze/params/build_prod_params.py | 3 +-- .../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 e5322309c576e..954f405f9c584 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 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 8a534cc425881..ed0cec9a8cb66 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 @@ 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 @@ def airflow_version(self) -> str: 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 d9bd3d3a1fbf6..e1c77990e9f78 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 @@ 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 @@ def _to_build_args(self): 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