Restore capability of reproducing CI failures using new Breeze (#24402)

The old breeze-legacy used to have a possibility of very easy
reproduction of CI failures by executing the right breeze command
that contained the commit hash of the PR being tested. This has
been broken for some time after we migrated to the new breeze,
but finally it was the time when it was needed again.

This PR brings back the capability by:

* addding --image-tag parameters for tests, shell and start-airflow
  commands
* if --image-tag is specified, then rather than building the
  image, it is pulled using the specified hash
* if --image-tag is specified, the local sources are not mounted
  to breeze when started, but the sources already embedded in the
  image are used ("skipped" set for --mount-sources).
* new "removed" command value is added to --mount-sources, it
  causes breeze command to remove the sources from the image (it
  is used when installing airflow during the tests for specified
  version (it's automatically used when --use-airflow-version
  is used).

.github/workflows/ci.yml

@@ -923,6 +923,7 @@ ${{ hashFiles('.pre-commit-config.yaml') }}"
         run: ./scripts/ci/testing/ci_run_airflow_testing.sh
         env:
           PR_LABELS: "${{ needs.build-info.outputs.pullRequestLabels }}"
+          IMAGE_TAG: ${{ env.IMAGE_TAG_FOR_THE_BUILD }}
       - name: "Upload airflow logs"
         uses: actions/upload-artifact@v2
         if: failure()
@@ -996,6 +997,7 @@ ${{ hashFiles('.pre-commit-config.yaml') }}"
         run: ./scripts/ci/testing/ci_run_airflow_testing.sh
         env:
           PR_LABELS: "${{ needs.build-info.outputs.pullRequestLabels }}"
+          IMAGE_TAG: ${{ env.IMAGE_TAG_FOR_THE_BUILD }}
       - name: "Upload airflow logs"
         uses: actions/upload-artifact@v2
         if: failure()
@@ -1068,6 +1070,7 @@ ${{ hashFiles('.pre-commit-config.yaml') }}"
         run: ./scripts/ci/testing/ci_run_airflow_testing.sh
         env:
           PR_LABELS: "${{ needs.build-info.outputs.pullRequestLabels }}"
+          IMAGE_TAG: ${{ env.IMAGE_TAG_FOR_THE_BUILD }}
       - name: "Upload airflow logs"
         uses: actions/upload-artifact@v2
         if: failure()
@@ -1138,6 +1141,7 @@ ${{ hashFiles('.pre-commit-config.yaml') }}"
         run: ./scripts/ci/testing/ci_run_airflow_testing.sh
         env:
           PR_LABELS: "${{ needs.build-info.outputs.pullRequestLabels }}"
+          IMAGE_TAG: ${{ env.IMAGE_TAG_FOR_THE_BUILD }}
       - name: "Upload airflow logs"
         uses: actions/upload-artifact@v2
         if: failure()
@@ -1206,6 +1210,7 @@ ${{ hashFiles('.pre-commit-config.yaml') }}"
         run: ./scripts/ci/testing/ci_run_airflow_testing.sh
         env:
           PR_LABELS: "${{ needs.build-info.outputs.pullRequestLabels }}"
+          IMAGE_TAG: ${{ env.IMAGE_TAG_FOR_THE_BUILD }}
       - name: "Upload airflow logs"
         uses: actions/upload-artifact@v2
         if: failure()

CI.rst

@@ -96,7 +96,7 @@ You can read more about Breeze in `BREEZE.rst <BREEZE.rst>`_ but in essence it i
 you to re-create CI environment in your local development instance and interact with it. In its basic
 form, when you do development you can run all the same tests that will be run in CI - but locally,
 before you submit them as PR. Another use case where Breeze is useful is when tests fail on CI. You can
-take the full ``COMMIT_SHA`` of the failed build pass it as ``--github-image-id`` parameter of Breeze and it will
+take the full ``COMMIT_SHA`` of the failed build pass it as ``--image-tag`` parameter of Breeze and it will
 download the very same version of image that was used in CI and run it locally. This way, you can very
 easily reproduce any failed test that happens in CI - even if you do not check out the sources
 connected with the run.
@@ -275,7 +275,7 @@ You can use those variables when you try to reproduce the build locally.
 |                                         |             |              |            | should set it to false, especially              |
 |                                         |             |              |            | in case our local sources are not the           |
 |                                         |             |              |            | ones we intend to use (for example              |
-|                                         |             |              |            | when ``--github-image-id`` is used              |
+|                                         |             |              |            | when ``--image-tag`` is used              |
 |                                         |             |              |            | in Breeze.                                      |
 |                                         |             |              |            |                                                 |
 |                                         |             |              |            | In CI jobs it is set to true                    |
@@ -668,12 +668,11 @@ For example knowing that the CI job was for commit ``cd27124534b46c9688a1d89e75f
 
 But you usually need to pass more variables and complex setup if you want to connect to a database or
 enable some integrations. Therefore it is easiest to use `Breeze <BREEZE.rst>`_ for that. For example if
-you need to reproduce a MySQL environment with kerberos integration enabled for commit
-cd27124534b46c9688a1d89e75fcd137ab5137e3, in python 3.8 environment you can run:
+you need to reproduce a MySQL environment in python 3.8 environment you can run:
 
 .. code-block:: bash
 
-  ./breeze-legacy --github-image-id cd27124534b46c9688a1d89e75fcd137ab5137e3 --python 3.8
+  breeze --image-tag cd27124534b46c9688a1d89e75fcd137ab5137e3 --python 3.8 --backend mysql
 
 You will be dropped into a shell with the exact version that was used during the CI run and you will
 be able to run pytest tests manually, easily reproducing the environment that was used in CI. Note that in

IMAGES.rst

@@ -281,16 +281,15 @@ to refresh them.
 
 Every developer can also pull and run images being result of a specific CI run in GitHub Actions.
 This is a powerful tool that allows to reproduce CI failures locally, enter the images and fix them much
-faster. It is enough to pass ``--github-image-id`` and the registry and Breeze will download and execute
+faster. It is enough to pass ``--image-tag`` and the registry and Breeze will download and execute
 commands using the same image that was used during the CI tests.
 
 For example this command will run the same Python 3.8 image as was used in build identified with
 9a621eaa394c0a0a336f8e1b31b35eff4e4ee86e commit SHA  with enabled rabbitmq integration.
 
 .. code-block:: bash
 
-  ./breeze-legacy --github-image-id 9a621eaa394c0a0a336f8e1b31b35eff4e4ee86e \
-    --python 3.8 --integration rabbitmq
+  breeze --image-tag 9a621eaa394c0a0a336f8e1b31b35eff4e4ee86e --python 3.8 --integration rabbitmq
 
 You can see more details and examples in `Breeze <BREEZE.rst>`_
 

breeze-complete

@@ -173,9 +173,6 @@ function breeze_complete::get_known_values_breeze() {
     -g | --github-repository)
         _breeze_known_values="${_breeze_default_github_repository}"
         ;;
-    -s | --github-image-id)
-        _breeze_known_values="${_breeze_default_github_image_id}"
-        ;;
     kind-cluster)
         _breeze_known_values="${_breeze_allowed_kind_operations}"
         ;;

breeze-legacy

@@ -970,21 +970,6 @@ function breeze::parse_arguments() {
             export GITHUB_REPOSITORY="${2}"
             shift 2
             ;;
-        -s | --github-image-id)
-            echo
-            echo "GitHub image id: ${2}"
-            echo
-            echo "Force pulling the image, using GitHub registry and skip mounting local sources."
-            echo "This is in order to get the exact same version as used in CI environment for SHA!."
-            echo "You can specify --skip-mounting-local-sources to not mount local sources to get exact. "
-            echo "behaviour as in the CI environment."
-            echo
-            export GITHUB_REGISTRY_PULL_IMAGE_TAG="${2}"
-            export CHECK_IMAGE_FOR_REBUILD="false"
-            export SKIP_BUILDING_PROD_IMAGE="true"
-            export SKIP_CHECK_REMOTE_IMAGE="true"
-            shift 2
-            ;;
         --init-script)
             export INIT_SCRIPT_FILE="${2}"
             echo "The initialization file is in ${INIT_SCRIPT_FILE}"
@@ -1370,15 +1355,6 @@ ${CMDNAME} shell [FLAGS] [-- <EXTRA_ARGS>]
       '${CMDNAME} shell -- -c \"ls -la\"'
       '${CMDNAME} -- -c \"ls -la\"'
 
-      For GitHub repository, the --github-repository flag can be used to specify the repository
-      to pull and push images. You can also use --github-image-id <COMMIT_SHA> in case
-      you want to pull the image with specific COMMIT_SHA tag.
-
-      '${CMDNAME} shell \\
-            --github-image-id 9a621eaa394c0a0a336f8e1b31b35eff4e4ee86e' - pull/use image with SHA
-      '${CMDNAME} \\
-            --github-image-id 9a621eaa394c0a0a336f8e1b31b35eff4e4ee86e' - pull/use image with SHA
-
 "
     readonly DETAILED_USAGE_SHELL
     export DETAILED_USAGE_EXEC="

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

@@ -319,7 +319,6 @@ def pull_image(
             verbose=verbose,
             wait_for_image=wait_for_image,
             tag_as_latest=tag_as_latest,
-            poll_time=10.0,
         )
         if return_code != 0:
             get_console().print(f"[error]There was an error when pulling CI image: {info}[/]")
@@ -528,26 +527,41 @@ def build_ci_image_in_parallel(
     pool.close()
 
 
-def rebuild_ci_image_if_needed(
-    build_params: Union[ShellParams, BuildCiParams], dry_run: bool, verbose: bool
+def rebuild_or_pull_ci_image_if_needed(
+    command_params: Union[ShellParams, BuildCiParams], dry_run: bool, verbose: bool
 ) -> None:
     """
     Rebuilds CI image if needed and user confirms it.
 
-    :param build_params: parameters of the shell
+    :param command_params: parameters of the command to execute
     :param dry_run: whether it's a dry_run
     :param verbose: should we print verbose messages
     """
     build_ci_image_check_cache = Path(
-        BUILD_CACHE_DIR, build_params.airflow_branch, f".built_{build_params.python}"
+        BUILD_CACHE_DIR, command_params.airflow_branch, f".built_{command_params.python}"
     )
-    ci_image_params = BuildCiParams(python=build_params.python, upgrade_to_newer_dependencies=False)
+    ci_image_params = BuildCiParams(
+        python=command_params.python, upgrade_to_newer_dependencies=False, image_tag=command_params.image_tag
+    )
+    if command_params.image_tag is not None:
+        return_code, message = run_pull_image(
+            image_params=ci_image_params,
+            dry_run=dry_run,
+            verbose=verbose,
+            parallel=False,
+            wait_for_image=True,
+            tag_as_latest=False,
+        )
+        if return_code != 0:
+            get_console().print(f"[error]Pulling image with {command_params.image_tag} failed! {message}[/]")
+            sys.exit(return_code)
+        return
     if build_ci_image_check_cache.exists():
         if verbose:
-            get_console().print(f'[info]{build_params.image_type} image already built locally.[/]')
+            get_console().print(f'[info]{command_params.image_type} image already built locally.[/]')
     else:
         get_console().print(
-            f'[warning]{build_params.image_type} image was never built locally or deleted. '
+            f'[warning]{command_params.image_type} image was never built locally or deleted. '
             'Forcing build.[/]'
         )
         ci_image_params.force_build = True

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

@@ -21,7 +21,7 @@
 
 import rich_click as click
 
-from airflow_breeze.commands.ci_image_commands import rebuild_ci_image_if_needed
+from airflow_breeze.commands.ci_image_commands import rebuild_or_pull_ci_image_if_needed
 from airflow_breeze.commands.main_command import main
 from airflow_breeze.global_constants import (
     DEFAULT_PYTHON_MAJOR_MINOR_VERSION,
@@ -45,6 +45,7 @@
     option_force_build,
     option_forward_credentials,
     option_github_repository,
+    option_image_tag,
     option_installation_package_format,
     option_integration,
     option_load_default_connection,
@@ -113,6 +114,7 @@
                 "--force-build",
                 "--mount-sources",
                 "--debian-version",
+                "--image-tag",
             ],
         },
     ],
@@ -141,6 +143,7 @@
                 "--force-build",
                 "--mount-sources",
                 "--debian-version",
+                "--image-tag",
             ],
         },
     ],
@@ -170,6 +173,7 @@
                 "--package-format",
                 "--force-build",
                 "--mount-sources",
+                "--image-tag",
             ],
         },
     ],
@@ -238,6 +242,7 @@
 @option_mount_sources
 @option_integration
 @option_db_reset
+@option_image_tag
 @option_answer
 @click.argument('extra-args', nargs=-1, type=click.UNPROCESSED)
 def shell(
@@ -261,6 +266,7 @@ def shell(
     force_build: bool,
     db_reset: bool,
     answer: Optional[str],
+    image_tag: Optional[str],
     extra_args: Tuple,
 ):
     """Enter breeze.py environment. this is the default command use when no other is selected."""
@@ -289,6 +295,7 @@ def shell(
         extra_args=extra_args,
         answer=answer,
         debian_version=debian_version,
+        image_tag=image_tag,
     )
 
 
@@ -312,6 +319,7 @@ def shell(
 @option_installation_package_format
 @option_mount_sources
 @option_integration
+@option_image_tag
 @option_db_reset
 @option_answer
 @click.argument('extra-args', nargs=-1, type=click.UNPROCESSED)
@@ -335,6 +343,7 @@ def start_airflow(
     use_packages_from_dist: bool,
     package_format: str,
     force_build: bool,
+    image_tag: Optional[str],
     db_reset: bool,
     answer: Optional[str],
     extra_args: Tuple,
@@ -362,6 +371,7 @@ def start_airflow(
         force_build=force_build,
         db_reset=db_reset,
         start_airflow=True,
+        image_tag=image_tag,
         extra_args=extra_args,
         answer=answer,
     )
@@ -398,7 +408,7 @@ def build_docs(
     """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)
+    rebuild_or_pull_ci_image_if_needed(command_params=params, dry_run=dry_run, verbose=verbose)
     ci_image_name = params.airflow_image_name
     doc_builder = DocBuildParams(
         package_filter=package_filter,
@@ -581,7 +591,7 @@ def enter_shell(**kwargs) -> RunCommandResult:
         get_console().print(CHEATSHEET, style=CHEATSHEET_STYLE)
     enter_shell_params = ShellParams(**filter_out_none(**kwargs))
     enter_shell_params.include_mypy_volume = True
-    rebuild_ci_image_if_needed(build_params=enter_shell_params, dry_run=dry_run, verbose=verbose)
+    rebuild_or_pull_ci_image_if_needed(command_params=enter_shell_params, dry_run=dry_run, verbose=verbose)
     return run_shell(verbose, dry_run, enter_shell_params)
 
 

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

@@ -24,7 +24,7 @@
 
 import click
 
-from airflow_breeze.commands.ci_image_commands import rebuild_ci_image_if_needed
+from airflow_breeze.commands.ci_image_commands import rebuild_or_pull_ci_image_if_needed
 from airflow_breeze.commands.main_command import main
 from airflow_breeze.global_constants import (
     ALLOWED_PLATFORMS,
@@ -260,7 +260,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)
+    rebuild_or_pull_ci_image_if_needed(command_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"],
@@ -299,7 +299,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)
+    rebuild_or_pull_ci_image_if_needed(command_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,
@@ -351,7 +351,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)
+    rebuild_or_pull_ci_image_if_needed(command_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,
@@ -540,7 +540,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)
+    rebuild_or_pull_ci_image_if_needed(command_params=shell_params, dry_run=dry_run, verbose=verbose)
     cmd_to_run = [
         "-c",
         "python /opt/airflow/scripts/in_container/verify_providers.py",
@@ -621,8 +621,8 @@ def release_prod_images(
     dry_run: bool,
 ):
     perform_environment_checks(verbose=verbose)
-    rebuild_ci_image_if_needed(
-        build_params=ShellParams(verbose=verbose, python=DEFAULT_PYTHON_MAJOR_MINOR_VERSION),
+    rebuild_or_pull_ci_image_if_needed(
+        command_params=ShellParams(verbose=verbose, python=DEFAULT_PYTHON_MAJOR_MINOR_VERSION),
         dry_run=dry_run,
         verbose=verbose,
     )