Add environment check and build image check for more Breeze commands (#…

…23687)

Several commands of Breeze depends on docker, docker compose
being available as well as breeze image. They will work
fine if you "just" built the image but they might benefit
from the image being rebuilt (to make sure all latest
dependencies are installed in the image). The common checks
done in "shell" command for that are now extracted to common
utils and run as first thing in those commands that need it.

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

@@ -69,9 +69,7 @@
 from airflow_breeze.utils.console import get_console
 from airflow_breeze.utils.docker_command_utils import (
     build_cache,
-    check_docker_compose_version,
-    check_docker_resources,
-    check_docker_version,
+    perform_environment_checks,
     prepare_docker_build_command,
     prepare_empty_docker_build_command,
 )
@@ -236,6 +234,7 @@ def run_build(ci_image_params: BuildCiParams) -> None:
             get_console().print(f"[error]Error when building image! {info}")
             sys.exit(return_code)
 
+    perform_environment_checks(verbose=verbose)
     parameters_passed = filter_out_none(**kwargs)
     if build_multiple_images:
         python_version_list = get_python_version_list(python_versions)
@@ -279,6 +278,7 @@ def pull_image(
     extra_pytest_args: Tuple,
 ):
     """Pull and optionally verify CI images - possibly in parallel for all Python versions."""
+    perform_environment_checks(verbose=verbose)
     if run_in_parallel:
         python_version_list = get_python_version_list(python_versions)
         ci_image_params_list = [
@@ -342,6 +342,7 @@ def verify_image(
     extra_pytest_args: Tuple,
 ):
     """Verify CI image."""
+    perform_environment_checks(verbose=verbose)
     if image_name is None:
         build_params = BuildCiParams(python=python, image_tag=image_tag, github_repository=github_repository)
         image_name = build_params.airflow_image_name_with_tag
@@ -517,9 +518,6 @@ def rebuild_ci_image_if_needed(
     :param dry_run: whether it's a dry_run
     :param verbose: should we print verbose messages
     """
-    check_docker_version(verbose=verbose)
-    check_docker_compose_version(verbose=verbose)
-    check_docker_resources(build_params.airflow_image_name, verbose=verbose, dry_run=dry_run)
     build_ci_image_check_cache = Path(
         BUILD_CACHE_DIR, build_params.airflow_branch, f".built_{build_params.python}"
     )

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

@@ -27,6 +27,7 @@
 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
@@ -48,6 +49,7 @@
     check_docker_resources,
     get_env_variables_for_docker_commands,
     get_extra_docker_flags,
+    perform_environment_checks,
 )
 from airflow_breeze.utils.path_utils import (
     AIRFLOW_SOURCES_ROOT,
@@ -406,6 +408,7 @@ def free_space(verbose: bool, dry_run: bool, answer: str):
 @option_verbose
 @option_dry_run
 def resource_check(verbose: bool, dry_run: bool):
+    perform_environment_checks(verbose=verbose)
     shell_params = ShellParams(verbose=verbose, python=DEFAULT_PYTHON_MAJOR_MINOR_VERSION)
     check_docker_resources(shell_params.airflow_image_name, verbose=verbose, dry_run=dry_run)
 
@@ -437,12 +440,14 @@ def command_hash_export(verbose: bool, output: IO):
 @option_verbose
 @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)
     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/developer_commands.py

@@ -61,12 +61,10 @@
 from airflow_breeze.utils.console import get_console
 from airflow_breeze.utils.custom_param_types import BetterChoice, NotVerifiedBetterChoice
 from airflow_breeze.utils.docker_command_utils import (
-    check_docker_compose_version,
-    check_docker_is_running,
     check_docker_resources,
-    check_docker_version,
     get_env_variables_for_docker_commands,
     get_extra_docker_flags,
+    perform_environment_checks,
 )
 from airflow_breeze.utils.path_utils import AIRFLOW_SOURCES_ROOT
 from airflow_breeze.utils.run_utils import assert_pre_commit_installed, filter_out_none, run_command
@@ -393,6 +391,7 @@ def build_docs(
     package_filter: Tuple[str],
 ):
     """Build documentation in the container."""
+    perform_environment_checks(verbose=verbose)
     params = BuildCiParams(github_repository=github_repository, python=DEFAULT_PYTHON_MAJOR_MINOR_VERSION)
     rebuild_ci_image_if_needed(build_params=params, dry_run=dry_run, verbose=verbose)
     ci_image_name = params.airflow_image_name
@@ -470,6 +469,7 @@ def static_checks(
     precommit_args: Tuple,
 ):
     assert_pre_commit_installed(verbose=verbose)
+    perform_environment_checks(verbose=verbose)
     command_to_execute = [sys.executable, "-m", "pre_commit", 'run']
     if last_commit and commit_ref:
         get_console().print("\n[error]You cannot specify both --last-commit and --commit-ref[/]\n")
@@ -529,6 +529,7 @@ def stop(verbose: bool, dry_run: bool, preserve_volumes: bool):
 @option_dry_run
 @click.argument('exec_args', nargs=-1, type=click.UNPROCESSED)
 def exec(verbose: bool, dry_run: bool, exec_args: Tuple):
+    perform_environment_checks(verbose=verbose)
     container_running = find_airflow_container(verbose, dry_run)
     if container_running:
         cmd_to_run = [
@@ -567,26 +568,21 @@ def enter_shell(**kwargs) -> Union[subprocess.CompletedProcess, subprocess.Calle
     """
     verbose = kwargs['verbose']
     dry_run = kwargs['dry_run']
-    check_docker_is_running(verbose)
-    check_docker_version(verbose)
-    check_docker_compose_version(verbose)
+    perform_environment_checks(verbose=verbose)
     if read_from_cache_file('suppress_asciiart') is None:
         get_console().print(ASCIIART, style=ASCIIART_STYLE)
     if read_from_cache_file('suppress_cheatsheet') is None:
         get_console().print(CHEATSHEET, style=CHEATSHEET_STYLE)
     enter_shell_params = ShellParams(**filter_out_none(**kwargs))
-    return run_shell_with_build_image_checks(verbose, dry_run, enter_shell_params)
+    rebuild_ci_image_if_needed(build_params=enter_shell_params, dry_run=dry_run, verbose=verbose)
+    return run_shell(verbose, dry_run, enter_shell_params)
 
 
-def run_shell_with_build_image_checks(
+def run_shell(
     verbose: bool, dry_run: bool, shell_params: ShellParams
 ) -> Union[subprocess.CompletedProcess, subprocess.CalledProcessError]:
     """
-    Executes a shell command built from params passed, checking if build is not needed.
-    * checks if there are enough resources to run shell
-    * checks if image was built at least once (if not - forces the build)
-    * if not forces, checks if build is needed and asks the user if so
-    * builds the image if needed
+    Executes a shell command built from params passed.
     * prints information about the build
     * constructs docker compose command to enter shell
     * executes it
@@ -595,7 +591,6 @@ def run_shell_with_build_image_checks(
     :param dry_run: do not execute "write" commands - just print what would happen
     :param shell_params: parameters of the execution
     """
-    rebuild_ci_image_if_needed(build_params=shell_params, dry_run=dry_run, verbose=verbose)
     shell_params.print_badge_info()
     cmd = ['docker-compose', 'run', '--service-ports', "-e", "BREEZE", '--rm', 'airflow']
     cmd_added = shell_params.command_passed

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

@@ -67,6 +67,7 @@
 from airflow_breeze.utils.custom_param_types import BetterChoice
 from airflow_breeze.utils.docker_command_utils import (
     build_cache,
+    perform_environment_checks,
     prepare_docker_build_command,
     prepare_empty_docker_build_command,
 )
@@ -278,6 +279,7 @@ def run_build(prod_image_params: BuildProdParams) -> None:
             get_console().print(f"[error]Error when building image! {info}")
             sys.exit(return_code)
 
+    perform_environment_checks(verbose=verbose)
     parameters_passed = filter_out_none(**kwargs)
     if build_multiple_images:
         python_version_list = get_python_version_list(python_versions)
@@ -321,6 +323,7 @@ def pull_prod_image(
     extra_pytest_args: Tuple,
 ):
     """Pull and optionally verify Production images - possibly in parallel for all Python versions."""
+    perform_environment_checks(verbose=verbose)
     if run_in_parallel:
         python_version_list = get_python_version_list(python_versions)
         prod_image_params_list = [
@@ -384,6 +387,7 @@ def verify_prod_image(
     extra_pytest_args: Tuple,
 ):
     """Verify Production image."""
+    perform_environment_checks(verbose=verbose)
     if image_name is None:
         build_params = BuildProdParams(
             python=python, image_tag=image_tag, github_repository=github_repository

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

@@ -25,6 +25,7 @@
 
 import click
 
+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 (
     ALLOWED_PLATFORMS,
@@ -64,6 +65,7 @@
 from airflow_breeze.utils.docker_command_utils import (
     get_env_variables_for_docker_commands,
     get_extra_docker_flags,
+    perform_environment_checks,
 )
 from airflow_breeze.utils.find_newer_dependencies import find_newer_dependencies
 from airflow_breeze.utils.parallel import check_async_run_results
@@ -239,6 +241,7 @@ def prepare_airflow_packages(
     version_suffix_for_pypi: str,
     debug: bool,
 ):
+    perform_environment_checks(verbose=verbose)
     shell_params = ShellParams(
         verbose=verbose,
         github_repository=github_repository,
@@ -249,6 +252,7 @@ def prepare_airflow_packages(
         install_providers_from_sources=False,
         mount_sources=MOUNT_ALL,
     )
+    rebuild_ci_image_if_needed(build_params=shell_params, dry_run=dry_run, verbose=verbose)
     result_command = run_with_debug(
         params=shell_params,
         command=["/opt/airflow/scripts/in_container/run_prepare_airflow_packages.sh"],
@@ -277,6 +281,7 @@ def prepare_provider_documentation(
     debug: bool,
     packages: List[str],
 ):
+    perform_environment_checks(verbose=verbose)
     shell_params = ShellParams(
         verbose=verbose,
         mount_sources=MOUNT_ALL,
@@ -285,6 +290,7 @@ def prepare_provider_documentation(
         answer=answer,
         skip_environment_initialization=True,
     )
+    rebuild_ci_image_if_needed(build_params=shell_params, dry_run=dry_run, verbose=verbose)
     cmd_to_run = ["/opt/airflow/scripts/in_container/run_prepare_provider_documentation.sh", *packages]
     result_command = run_with_debug(
         params=shell_params,
@@ -323,6 +329,7 @@ def prepare_provider_packages(
     debug: bool,
     packages: Tuple[str, ...],
 ):
+    perform_environment_checks(verbose=verbose)
     packages_list = list(packages)
     if package_list_file:
         packages_list.extend([package.strip() for package in package_list_file.readlines()])
@@ -335,6 +342,7 @@ def prepare_provider_packages(
         skip_environment_initialization=True,
         version_suffix_for_pypi=version_suffix_for_pypi,
     )
+    rebuild_ci_image_if_needed(build_params=shell_params, dry_run=dry_run, verbose=verbose)
     cmd_to_run = ["/opt/airflow/scripts/in_container/run_prepare_provider_packages.sh", *packages_list]
     result_command = run_with_debug(
         params=shell_params,
@@ -417,6 +425,7 @@ def generate_constraints(
     debug: bool,
     airflow_constraints_mode: str,
 ):
+    perform_environment_checks(verbose=verbose)
     if debug and run_in_parallel:
         get_console().print("\n[error]Cannot run --debug and --run-in-parallel at the same time[/]\n")
         sys.exit(1)
@@ -508,6 +517,7 @@ def verify_provider_packages(
     package_format: str,
     github_repository: str,
 ):
+    perform_environment_checks(verbose=verbose)
     shell_params = ShellParams(
         verbose=verbose,
         mount_sources=MOUNT_SELECTED,
@@ -519,6 +529,7 @@ def verify_provider_packages(
         use_packages_from_dist=use_packages_from_dist,
         package_format=package_format,
     )
+    rebuild_ci_image_if_needed(build_params=shell_params, dry_run=dry_run, verbose=verbose)
     cmd_to_run = [
         "-c",
         "python /opt/airflow/scripts/in_container/verify_providers.py",
@@ -597,6 +608,12 @@ def release_prod_images(
     verbose: bool,
     dry_run: bool,
 ):
+    perform_environment_checks(verbose=verbose)
+    rebuild_ci_image_if_needed(
+        build_params=ShellParams(verbose=verbose, python=DEFAULT_PYTHON_MAJOR_MINOR_VERSION),
+        dry_run=dry_run,
+        verbose=verbose,
+    )
     if not match(r"^\d*\.\d*\.\d*$", airflow_version):
         get_console().print(
             f"[warning]Skipping latest image tagging as this is a pre-release version: {airflow_version}"

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

@@ -21,7 +21,6 @@
 
 import click
 
-from airflow_breeze.commands.developer_commands import check_docker_is_running, check_docker_resources
 from airflow_breeze.commands.main_command import main
 from airflow_breeze.global_constants import ALLOWED_TEST_TYPES
 from airflow_breeze.params.build_prod_params import BuildProdParams
@@ -38,7 +37,10 @@
 )
 from airflow_breeze.utils.console import get_console
 from airflow_breeze.utils.custom_param_types import BetterChoice
-from airflow_breeze.utils.docker_command_utils import get_env_variables_for_docker_commands
+from airflow_breeze.utils.docker_command_utils import (
+    get_env_variables_for_docker_commands,
+    perform_environment_checks,
+)
 from airflow_breeze.utils.run_tests import run_docker_compose_tests
 from airflow_breeze.utils.run_utils import run_command
 
@@ -150,9 +152,7 @@ def tests(
 
     exec_shell_params = ShellParams(verbose=verbose, dry_run=dry_run)
     env_variables = get_env_variables_for_docker_commands(exec_shell_params)
-    check_docker_is_running(verbose)
-    check_docker_resources(exec_shell_params.airflow_image_name, verbose=verbose, dry_run=dry_run)
-
+    perform_environment_checks(verbose=verbose)
     cmd = ['docker-compose', 'run', '--service-ports', '--rm', 'airflow']
     cmd.extend(list(extra_pytest_args))
     run_command(

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

@@ -573,3 +573,9 @@ def get_env_variables_for_docker_commands(params: Union[ShellParams, BuildCiPara
             env_variables[variable] = str(constant_param_value)
     update_expected_environment_variables(env_variables)
     return env_variables
+
+
+def perform_environment_checks(verbose: bool):
+    check_docker_is_running(verbose=verbose)
+    check_docker_version(verbose=verbose)
+    check_docker_compose_version(verbose=verbose)