Add tagging image as latest for CI image wait

.github/workflows/ci.yml

@@ -579,7 +579,7 @@ jobs:
         run: breeze free-space
       - name: Wait for CI images ${{ env.PYTHON_VERSIONS }}:${{ env.IMAGE_TAG_FOR_THE_BUILD }}
         id: wait-for-images
-        run: breeze pull-image --run-in-parallel --verify-image --wait-for-image
+        run: breeze pull-image --run-in-parallel --verify-image --wait-for-image --tag-as-latest
         env:
           PYTHON_VERSIONS: ${{ needs.build-info.outputs.pythonVersionsListAsString }}
           IMAGE_TAG: ${{ env.IMAGE_TAG_FOR_THE_BUILD }}

dev/breeze/src/airflow_breeze/commands/ci_image_commands.py

@@ -73,11 +73,11 @@
     prepare_docker_build_command,
     prepare_empty_docker_build_command,
 )
+from airflow_breeze.utils.image import run_pull_image, run_pull_in_parallel
 from airflow_breeze.utils.mark_image_as_refreshed import mark_image_as_refreshed
 from airflow_breeze.utils.md5_build_check import md5sum_check_if_build_is_needed
 from airflow_breeze.utils.parallel import check_async_run_results
 from airflow_breeze.utils.path_utils import AIRFLOW_SOURCES_ROOT, BUILD_CACHE_DIR
-from airflow_breeze.utils.pulll_image import run_pull_image, run_pull_in_parallel
 from airflow_breeze.utils.python_versions import get_python_version_list
 from airflow_breeze.utils.registry import login_to_github_docker_registry
 from airflow_breeze.utils.run_tests import verify_an_image

dev/breeze/src/airflow_breeze/commands/configuration_and_maintenance_commands.py

@@ -27,7 +27,6 @@
 from click import Context
 
 from airflow_breeze import NAME, VERSION
-from airflow_breeze.commands.ci_image_commands import rebuild_ci_image_if_needed
 from airflow_breeze.commands.main_command import main
 from airflow_breeze.global_constants import DEFAULT_PYTHON_MAJOR_MINOR_VERSION, MOUNT_ALL
 from airflow_breeze.params.shell_params import ShellParams
@@ -51,6 +50,7 @@
     get_extra_docker_flags,
     perform_environment_checks,
 )
+from airflow_breeze.utils.image import find_available_ci_image
 from airflow_breeze.utils.path_utils import (
     AIRFLOW_SOURCES_ROOT,
     BUILD_CACHE_DIR,
@@ -441,13 +441,7 @@ def command_hash_export(verbose: bool, output: IO):
 @option_dry_run
 def fix_ownership(verbose: bool, dry_run: bool):
     perform_environment_checks(verbose=verbose)
-    shell_params = ShellParams(
-        verbose=verbose,
-        mount_sources=MOUNT_ALL,
-        python=DEFAULT_PYTHON_MAJOR_MINOR_VERSION,
-        skip_environment_initialization=True,
-    )
-    rebuild_ci_image_if_needed(build_params=shell_params, dry_run=dry_run, verbose=verbose)
+    shell_params = find_available_ci_image(dry_run, verbose)
     extra_docker_flags = get_extra_docker_flags(MOUNT_ALL)
     env = get_env_variables_for_docker_commands(shell_params)
     cmd = [

dev/breeze/src/airflow_breeze/commands/production_image_commands.py

@@ -71,8 +71,8 @@
     prepare_docker_build_command,
     prepare_empty_docker_build_command,
 )
+from airflow_breeze.utils.image import run_pull_image, run_pull_in_parallel
 from airflow_breeze.utils.path_utils import AIRFLOW_SOURCES_ROOT, DOCKER_CONTEXT_DIR
-from airflow_breeze.utils.pulll_image import run_pull_image, run_pull_in_parallel
 from airflow_breeze.utils.python_versions import get_python_version_list
 from airflow_breeze.utils.registry import login_to_github_docker_registry
 from airflow_breeze.utils.run_tests import verify_an_image

dev/breeze/src/airflow_breeze/configure_rich_click.py

@@ -64,5 +64,9 @@
             RELEASE_MANAGEMENT_COMMANDS,
         ]
     }
-except ImportError:
-    import click  # type: ignore[no-redef]
+except ImportError as e:
+    if "No module named 'rich_click'" in e.msg:
+        # just ignore the import error when rich_click is missing
+        import click  # type: ignore[no-redef]
+    else:
+        raise

dev/breeze/src/airflow_breeze/utils/docker_command_utils.py

@@ -528,6 +528,7 @@ def update_expected_environment_variables(env: Dict[str, str]) -> None:
     "PYTHON_MAJOR_MINOR_VERSION": "python",
     "SQLITE_URL": "sqlite_url",
     "START_AIRFLOW": "start_airflow",
+    "SKIP_ENVIRONMENT_INITIALIZATION": "skip_environment_initialization",
     "USE_AIRFLOW_VERSION": "use_airflow_version",
     "USE_PACKAGES_FROM_DIST": "use_packages_from_dist",
     "VERSION_SUFFIX_FOR_PYPI": "version_suffix_for_pypi",

dev/breeze/src/airflow_breeze/utils/pulll_image.py → dev/breeze/src/airflow_breeze/utils/image.py

@@ -16,12 +16,19 @@
 # under the License.
 
 import multiprocessing as mp
+import subprocess
 import time
 from typing import List, Tuple, Union
 
+from airflow_breeze.global_constants import (
+    ALLOWED_PYTHON_MAJOR_MINOR_VERSIONS,
+    DEFAULT_PYTHON_MAJOR_MINOR_VERSION,
+    MOUNT_ALL,
+)
 from airflow_breeze.params._common_build_params import _CommonBuildParams
 from airflow_breeze.params.build_ci_params import BuildCiParams
 from airflow_breeze.params.build_prod_params import BuildProdParams
+from airflow_breeze.params.shell_params import ShellParams
 from airflow_breeze.utils.console import get_console
 from airflow_breeze.utils.mark_image_as_refreshed import mark_image_as_refreshed
 from airflow_breeze.utils.parallel import check_async_run_results
@@ -183,3 +190,56 @@ def run_pull_and_verify_image(
         verbose=verbose,
         extra_pytest_args=extra_pytest_args,
     )
+
+
+def just_pull_ci_image(
+    python_version: str, dry_run: bool, verbose: bool
+) -> Tuple[ShellParams, Union[subprocess.CompletedProcess, subprocess.CalledProcessError]]:
+    shell_params = ShellParams(
+        verbose=verbose,
+        mount_sources=MOUNT_ALL,
+        python=python_version,
+        skip_environment_initialization=True,
+    )
+    get_console().print(f"[info]Pulling {shell_params.airflow_image_name_with_tag}.[/]")
+    pull_command_result = run_command(
+        ["docker", "pull", shell_params.airflow_image_name_with_tag],
+        verbose=verbose,
+        dry_run=dry_run,
+        check=True,
+    )
+    return shell_params, pull_command_result
+
+
+def check_if_ci_image_available(
+    python_version: str, dry_run: bool, verbose: bool
+) -> Tuple[ShellParams, Union[subprocess.CompletedProcess, subprocess.CalledProcessError]]:
+    shell_params = ShellParams(
+        verbose=verbose,
+        mount_sources=MOUNT_ALL,
+        python=python_version,
+        skip_environment_initialization=True,
+    )
+    inspect_command_result = run_command(
+        ["docker", "inspect", shell_params.airflow_image_name_with_tag],
+        stdout=subprocess.DEVNULL,
+        verbose=verbose,
+        dry_run=dry_run,
+        check=False,
+    )
+    return (
+        shell_params,
+        inspect_command_result,
+    )
+
+
+def find_available_ci_image(dry_run: bool, verbose: bool) -> ShellParams:
+    for python_version in ALLOWED_PYTHON_MAJOR_MINOR_VERSIONS:
+        shell_params, inspect_command_result = check_if_ci_image_available(python_version, dry_run, verbose)
+        if inspect_command_result.returncode == 0:
+            get_console().print(
+                "[info]Running fix_ownership " f"with {shell_params.airflow_image_name_with_tag}.[/]"
+            )
+            return shell_params
+    shell_params, _ = just_pull_ci_image(DEFAULT_PYTHON_MAJOR_MINOR_VERSION, dry_run, verbose)
+    return shell_params