Add common code and utils to implement ER pre-commit in a clearer way (…

…apache#35875)

* Add common code and utils to implement ER pre-commit in a clearer way

Extracts comon initialization code for pre-commits that allows to use
a simpler way of running pre-commits with breeze ci image. Rather
than running a custom `docker` command, we are using the
`breeze shell` command to run the command inside the image. We are
also fixing a problem where the ER diagram was generated using
random backend - we force it now to use Postgres as backend.

There are few small problems solved as part of it as well:

* We are also switching the console used to be the standard colored
  one so that the output will use colors also in git hook and CI.

* When image needs update we are going to warn about it in the
  builds that are "serial" (i.e. do not run in parallell) that
  the image should be upgraded. We are not doing it for checks
  that use parallelism because it would pollute the output too
  much (and those are usually not run with images anyway).

* There was an error where pre-commit run with breeze would
  also start port-forwarding if database was started together
  with it, so pre-commit project no longer adds port forwarding
  for backends.

* If there was a database stated as part of pre-commit project, it
  would continue running after pre-commit completed, which is
  undesireable. The pre-commit will now stop pre-commit project
  after completion.

* Update dev/breeze/src/airflow_breeze/utils/md5_build_check.py

Co-authored-by: Pankaj Koti <[email protected]>

---------

Co-authored-by: Pankaj Koti <[email protected]>

.github/workflows/ci.yml

@@ -611,7 +611,7 @@ jobs:
           --skip-image-upgrade-check --commit-ref "${{ github.sha }}"
         env:
           VERBOSE: "false"
-          SKIP_IMAGE_PRE_COMMITS: "true"
+          SKIP_BREEZE_PRE_COMMITS: "true"
           SKIP: ${{ needs.build-info.outputs.skip-pre-commits }}
           COLUMNS: "250"
 

.pre-commit-config.yaml

@@ -429,6 +429,7 @@ repos:
         language: python
         files: ^docs
         pass_filenames: false
+        additional_dependencies: ['rich>=12.4.4']
       - id: check-pydevd-left-in-code
         language: pygrep
         name: Check for pydevd debug statements accidentally left
@@ -1070,5 +1071,5 @@ repos:
         entry: ./scripts/ci/pre_commit/pre_commit_update_er_diagram.py
         pass_filenames: false
         files: ^airflow/migrations/versions/.*\.py$|^docs/apache-airflow/migrations-ref\.rst$
-        additional_dependencies: ['rich>=12.4.4', 'inputimeout', 'pyyaml', 'jsonschema', 'filelock', 'markdown-it-py']
+        additional_dependencies: ['rich>=12.4.4']
         ## ONLY ADD PRE-COMMITS HERE THAT REQUIRE CI IMAGE

STATIC_CODE_CHECKS.rst

@@ -123,8 +123,8 @@ require Breeze Docker image to be built locally.
   ``export SKIP=ruff,mypy-core,``. You can also add this to your ``.bashrc`` or ``.zshrc`` if you
   do not want to set it manually every time you enter the terminal.
 
-  In case you do not have breeze image configured locally, you can also disable all checks that require
-  the image by setting ``SKIP_IMAGE_PRE_COMMITS`` to "true". This will mark the tests as "green" automatically
+  In case you do not have breeze image configured locally, you can also disable all checks that require breeze
+  the image by setting ``SKIP_BREEZE_PRE_COMMITS`` to "true". This will mark the tests as "green" automatically
   when run locally (note that those checks will anyway run in CI).
 
 .. note:: Mypy volume cache

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

@@ -132,8 +132,6 @@ def check_if_image_building_is_needed(ci_image_params: BuildCiParams, output: Ou
     )
     if result.returncode != 0:
         return True
-    if ci_image_params.skip_image_upgrade_check:
-        return False
     if not ci_image_params.force_build and not ci_image_params.upgrade_to_newer_dependencies:
         if not should_we_run_the_build(build_ci_params=ci_image_params):
             return False
@@ -638,6 +636,7 @@ def should_we_run_the_build(build_ci_params: BuildCiParams) -> bool:
     from inputimeout import TimeoutOccurred
 
     if not md5sum_check_if_build_is_needed(
+        build_ci_params=build_ci_params,
         md5sum_cache_dir=build_ci_params.md5sum_cache_dir,
         skip_provider_dependencies_check=build_ci_params.skip_provider_dependencies_check,
     ):
@@ -808,6 +807,8 @@ def rebuild_or_pull_ci_image_if_needed(command_params: ShellParams | BuildCiPara
         platform=command_params.platform,
         force_build=command_params.force_build,
         skip_provider_dependencies_check=command_params.skip_provider_dependencies_check,
+        skip_image_upgrade_check=command_params.skip_image_upgrade_check,
+        warn_image_upgrade_needed=command_params.warn_image_upgrade_needed,
     )
     if command_params.image_tag is not None and command_params.image_tag != "latest":
         return_code, message = run_pull_image(

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

@@ -85,6 +85,7 @@
     option_use_airflow_version,
     option_use_packages_from_dist,
     option_verbose,
+    option_warn_image_upgrade_needed,
 )
 from airflow_breeze.utils.console import get_console
 from airflow_breeze.utils.custom_param_types import BetterChoice
@@ -203,6 +204,7 @@ def run(self):
 @option_skip_db_tests
 @option_skip_environment_initialization
 @option_skip_image_upgrade_check
+@option_warn_image_upgrade_needed
 @option_standalone_dag_processor
 @option_upgrade_boto
 @option_use_airflow_version
@@ -248,6 +250,7 @@ def shell(
     use_airflow_version: str | None,
     use_packages_from_dist: bool,
     verbose_commands: bool,
+    warn_image_upgrade_needed: bool,
 ):
     """Enter breeze environment. this is the default command use when no other is selected."""
     if get_verbose() or get_dry_run() and not quiet:
@@ -298,6 +301,7 @@ def shell(
         use_packages_from_dist=use_packages_from_dist,
         verbose_commands=verbose_commands,
         restart=restart,
+        warn_image_upgrade_needed=warn_image_upgrade_needed,
     )
     fix_ownership_using_docker()
     sys.exit(result.returncode)
@@ -693,7 +697,7 @@ def static_checks(
         text=True,
         env=env,
     )
-    if not os.environ.get("SKIP_IMAGE_PRE_COMMITS"):
+    if not os.environ.get("SKIP_BREEZE_PRE_COMMITS"):
         fix_ownership_using_docker()
     if static_checks_result.returncode != 0:
         if os.environ.get("CI"):

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

@@ -94,6 +94,7 @@
             "options": [
                 "--quiet",
                 "--skip-image-upgrade-check",
+                "--warn-image-upgrade-needed",
                 "--skip-environment-initialization",
                 "--tty",
             ],

dev/breeze/src/airflow_breeze/params/build_ci_params.py

@@ -42,6 +42,7 @@ class BuildCiParams(CommonBuildParams):
     eager_upgrade_additional_requirements: str | None = None
     skip_provider_dependencies_check: bool = False
     skip_image_upgrade_check: bool = False
+    warn_image_upgrade_needed: bool = False
 
     @property
     def airflow_version(self):

dev/breeze/src/airflow_breeze/params/shell_params.py

@@ -183,6 +183,7 @@ class ShellParams:
     verbose: bool = False
     verbose_commands: bool = False
     version_suffix_for_pypi: str = ""
+    warn_image_upgrade_needed: bool = False
 
     def clone_with_test(self, test_type: str) -> ShellParams:
         new_params = deepcopy(self)
@@ -275,6 +276,10 @@ def get_backend_compose_files(self, backend: str) -> list[Path]:
         backend_docker_compose_file = DOCKER_COMPOSE_DIR / f"backend-{backend}.yml"
         if backend in ("sqlite", "none") or not self.forward_ports:
             return [backend_docker_compose_file]
+        if self.project_name == "pre-commit":
+            # do not forward ports for pre-commit runs - to not clash with running containers from
+            # breeze
+            return [backend_docker_compose_file]
         return [backend_docker_compose_file, DOCKER_COMPOSE_DIR / f"backend-{backend}-port.yml"]
 
     @cached_property

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

@@ -758,6 +758,12 @@ def _set_default_from_parent(ctx: click.core.Context, option: click.core.Option,
     is_flag=True,
     envvar="SKIP_IMAGE_UPGRADE_CHECK",
 )
+option_warn_image_upgrade_needed = click.option(
+    "--warn-image-upgrade-needed",
+    help="Warn when image upgrade is needed even if --skip-upgrade-check is used.",
+    is_flag=True,
+    envvar="WARN_IMAGE_UPGRADE_NEEDED",
+)
 option_skip_environment_initialization = click.option(
     "--skip-environment-initialization",
     help="Skip running breeze entrypoint initialization - no user output, no db checks.",

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

@@ -23,12 +23,16 @@
 import os
 import sys
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 from airflow_breeze.global_constants import ALL_PROVIDER_YAML_FILES, FILES_FOR_REBUILD_CHECK
 from airflow_breeze.utils.console import get_console
 from airflow_breeze.utils.path_utils import AIRFLOW_SOURCES_ROOT
 from airflow_breeze.utils.run_utils import run_command
 
+if TYPE_CHECKING:
+    from airflow_breeze.params.build_ci_params import BuildCiParams
+
 
 def check_md5checksum_in_cache_modified(file_hash: str, cache_path: Path, update: bool) -> bool:
     """
@@ -126,34 +130,49 @@ def calculate_md5_checksum_for_files(
     return modified_files, not_modified_files
 
 
-def md5sum_check_if_build_is_needed(md5sum_cache_dir: Path, skip_provider_dependencies_check: bool) -> bool:
+def md5sum_check_if_build_is_needed(
+    build_ci_params: BuildCiParams, md5sum_cache_dir: Path, skip_provider_dependencies_check: bool
+) -> bool:
     """
     Checks if build is needed based on whether important files were modified.
 
+    :param build_ci_params: parameters for the build
     :param md5sum_cache_dir: directory where cached md5 sums are stored
     :param skip_provider_dependencies_check: whether to skip regeneration of the provider dependencies
 
     :return: True if build is needed.
     """
-    build_needed = False
     modified_files, not_modified_files = calculate_md5_checksum_for_files(
         md5sum_cache_dir, update=False, skip_provider_dependencies_check=skip_provider_dependencies_check
     )
     if modified_files:
+        if build_ci_params.skip_image_upgrade_check:
+            if build_ci_params.warn_image_upgrade_needed:
+                get_console().print(
+                    "\n[warning]You are skipping the image upgrade check, but the image needs an upgrade. "
+                    "This might lead to out-dated results of the check![/]"
+                )
+                get_console().print(
+                    f"[info]Consider running `breeze ci-image build --python {build_ci_params.python} "
+                    f"at earliest convenience![/]\n"
+                )
+            return False
         get_console().print(
             f"[warning]The following important files are modified in {AIRFLOW_SOURCES_ROOT} "
             f"since last time image was built: [/]\n\n"
         )
         for file in modified_files:
             get_console().print(f" * [info]{file}[/]")
         get_console().print("\n[warning]Likely CI image needs rebuild[/]\n")
-        build_needed = True
+        return True
     else:
+        if build_ci_params.skip_image_upgrade_check:
+            return False
         get_console().print(
             "[info]Docker image build is not needed for CI build as no important files are changed! "
             "You can add --force-build to force it[/]"
         )
-    return build_needed
+    return False
 
 
 def save_md5_file(cache_path: Path, file_content: str) -> None: