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 a68b419)

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:

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:

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