rebase conflict resolved (#20338)

GitOrigin-RevId: 95740a87083c703968ce3da45b15113851ef09f7

dev/breeze/src/airflow_breeze/branch_defaults.py

@@ -0,0 +1,19 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+AIRFLOW_BRANCH = "main"
+DEFAULT_AIRFLOW_CONSTRAINTS_BRANCH = "constraints-main"

dev/breeze/src/airflow_breeze/breeze.py

@@ -21,8 +21,8 @@
 
 import click
 from click import ClickException
-from rich.console import Console
 
+from airflow_breeze.console import console
 from airflow_breeze.visuals import ASCIIART, ASCIIART_STYLE
 
 NAME = "Breeze2"
@@ -32,8 +32,6 @@
 
 __AIRFLOW_CFG_FILE = "setup.cfg"
 
-console = Console(force_terminal=True, color_system="standard", width=180)
-
 
 def get_airflow_sources_root():
     return __AIRFLOW_SOURCES_ROOT
@@ -95,13 +93,107 @@ def shell(verbose: bool):
 
 
 @option_verbose
-@main.command()
-def build_ci_image(verbose: bool):
-    """Builds breeze.ci image for breeze.py."""
+@main.command(name='build-ci-image')
+@click.option(
+    '--additional-extras',
+    help='This installs additional extra package while installing airflow in the image.',
+)
+@click.option('-p', '--python', help='Choose your python version')
+@click.option(
+    '--additional-dev-apt-deps', help='Additional apt dev dependencies to use when building the images.'
+)
+@click.option(
+    '--additional-runtime-apt-deps',
+    help='Additional apt runtime dependencies to use when building the images.',
+)
+@click.option(
+    '--additional-python-deps', help='Additional python dependencies to use when building the images.'
+)
+@click.option(
+    '--additional_dev_apt_command', help='Additional command executed before dev apt deps are installed.'
+)
+@click.option(
+    '--additional_runtime_apt_command',
+    help='Additional command executed before runtime apt deps are installed.',
+)
+@click.option(
+    '--additional_dev_apt_env', help='Additional environment variables set when adding dev dependencies.'
+)
+@click.option(
+    '--additional_runtime_apt_env',
+    help='Additional environment variables set when adding runtime dependencies.',
+)
+@click.option('--dev-apt-command', help='The basic command executed before dev apt deps are installed.')
+@click.option(
+    '--dev-apt-deps',
+    help='The basic apt dev dependencies to use when building the images.',
+)
+@click.option(
+    '--runtime-apt-command', help='The basic command executed before runtime apt deps are installed.'
+)
+@click.option(
+    '--runtime-apt-deps',
+    help='The basic apt runtime dependencies to use when building the images.',
+)
+@click.option('--github-repository', help='Choose repository to push/pull image.')
+@click.option('--build-cache', help='Cache option')
+@click.option('--upgrade-to-newer-dependencies', is_flag=True)
+def build_ci_image(
+    verbose: bool,
+    additional_extras: Optional[str],
+    python: Optional[float],
+    additional_dev_apt_deps: Optional[str],
+    additional_runtime_apt_deps: Optional[str],
+    additional_python_deps: Optional[str],
+    additional_dev_apt_command: Optional[str],
+    additional_runtime_apt_command: Optional[str],
+    additional_dev_apt_env: Optional[str],
+    additional_runtime_apt_env: Optional[str],
+    dev_apt_command: Optional[str],
+    dev_apt_deps: Optional[str],
+    runtime_apt_command: Optional[str],
+    runtime_apt_deps: Optional[str],
+    github_repository: Optional[str],
+    build_cache: Optional[str],
+    upgrade_to_newer_dependencies: bool,
+):
+    """Builds docker CI image without entering the container."""
+    from airflow_breeze.ci.build_image import build_image
+
     if verbose:
         console.print(f"\n[blue]Building image of airflow from {__AIRFLOW_SOURCES_ROOT}[/]\n")
-    raise ClickException("\nPlease implement building the CI image\n")
+    build_image(
+        verbose,
+        additional_extras=additional_extras,
+        python_version=python,
+        additional_dev_apt_deps=additional_dev_apt_deps,
+        additional_runtime_apt_deps=additional_runtime_apt_deps,
+        additional_python_deps=additional_python_deps,
+        additional_runtime_apt_command=additional_runtime_apt_command,
+        additional_dev_apt_command=additional_dev_apt_command,
+        additional_dev_apt_env=additional_dev_apt_env,
+        additional_runtime_apt_env=additional_runtime_apt_env,
+        dev_apt_command=dev_apt_command,
+        dev_apt_deps=dev_apt_deps,
+        runtime_apt_command=runtime_apt_command,
+        runtime_apt_deps=runtime_apt_deps,
+        github_repository=github_repository,
+        docker_cache=build_cache,
+        upgrade_to_newer_dependencies=upgrade_to_newer_dependencies,
+    )
+
+
+@option_verbose
+@main.command(name='build-prod-image')
+def build_prod_image(verbose: bool):
+    """Builds docker Production image without entering the container."""
+    if verbose:
+        console.print("\n[blue]Building image[/]\n")
+    raise ClickException("\nPlease implement building the Production image\n")
 
 
 if __name__ == '__main__':
+    from airflow_breeze.cache import BUILD_CACHE_DIR
+
+    BUILD_CACHE_DIR.mkdir(parents=True, exist_ok=True)
     main()

dev/breeze/src/airflow_breeze/cache.py

@@ -0,0 +1,79 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import sys
+from pathlib import Path
+from typing import Any, List, Optional, Tuple
+
+from airflow_breeze import global_constants
+from airflow_breeze.breeze import get_airflow_sources_root
+from airflow_breeze.console import console
+
+BUILD_CACHE_DIR = Path(get_airflow_sources_root(), '.build')
+
+
+def check_if_cache_exists(param_name: str) -> bool:
+    return (Path(BUILD_CACHE_DIR) / f".{param_name}").exists()
+
+
+def read_from_cache_file(param_name: str) -> Optional[str]:
+    cache_exists = check_if_cache_exists(param_name)
+    if cache_exists:
+        return (Path(BUILD_CACHE_DIR) / f".{param_name}").read_text().strip()
+    else:
+        return None
+
+
+def write_to_cache_file(param_name: str, param_value: str, check_allowed_values: bool = True) -> None:
+    allowed = False
+    if check_allowed_values:
+        allowed, allowed_values = check_if_values_allowed(param_name, param_value)
+    if allowed or not check_allowed_values:
+        Path(BUILD_CACHE_DIR, f".{param_name}").write_text(param_value)
+    else:
+        console.print(f'[cyan]You have sent the {param_value} for {param_name}')
+        console.print(f'[cyan]Allowed value for the {param_name} are {allowed_values}')
+        console.print('[cyan]Provide one of the supported params. Write to cache dir failed')
+        sys.exit()
+
+
+def check_cache_and_write_if_not_cached(
+    param_name: str, default_param_value: str
+) -> Tuple[bool, Optional[str]]:
+    is_cached = False
+    allowed = False
+    cached_value = read_from_cache_file(param_name)
+    if cached_value is None:
+        write_to_cache_file(param_name, default_param_value)
+        cached_value = default_param_value
+    else:
+        allowed, allowed_values = check_if_values_allowed(param_name, cached_value)
+        if allowed:
+            is_cached = True
+        else:
+            write_to_cache_file(param_name, default_param_value)
+            cached_value = default_param_value
+    return is_cached, cached_value
+
+
+def check_if_values_allowed(param_name: str, param_value: str) -> Tuple[bool, List[Any]]:
+    allowed = False
+    allowed_values: List[Any] = []
+    allowed_values = getattr(global_constants, f'ALLOWED_{param_name.upper()}')
+    if param_value in allowed_values:
+        allowed = True
+    return allowed, allowed_values

dev/breeze/src/airflow_breeze/ci/build_image.py

@@ -0,0 +1,94 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# from airflow_breeze.utils import run_command
+from pathlib import Path
+from typing import List
+
+from airflow_breeze.breeze import get_airflow_sources_root
+from airflow_breeze.cache import check_cache_and_write_if_not_cached
+from airflow_breeze.ci.build_params import BuildParams
+from airflow_breeze.console import console
+from airflow_breeze.utils import filter_out_none, run_command
+
+PARAMS_CI_IMAGE = [
+    "python_base_image",
+    "airflow_version",
+    "airflow_branch",
+    "airflow_extras",
+    "airflow_pre_cached_pip_packages",
+    "additional_airflow_extras",
+    "additional_python_deps",
+    "additional_dev_apt_command",
+    "additional_dev_apt_deps",
+    "additional_dev_apt_env",
+    "additional_runtime_apt_command",
+    "additional_runtime_apt_deps",
+    "additional_runtime_apt_env",
+    "upgrade_to_newer_dependencies",
+    "constraints_github_repository",
+    "airflow_constraints_reference",
+    "airflow_constraints",
+    "airflow_image_repository",
+    "airflow_image_date_created",
+    "build_id",
+    "commit_sha",
+]
+
+PARAMS_TO_VERIFY_CI_IMAGE = [
+    "dev_apt_command",
+    "dev_apt_deps",
+    "runtime_apt_command",
+    "runtime_apt_deps",
+]
+
+
+def construct_arguments_docker_command(ci_image: BuildParams) -> List[str]:
+    args_command = []
+    for param in PARAMS_CI_IMAGE:
+        args_command.append("--build-arg")
+        args_command.append(param.upper() + "=" + str(getattr(ci_image, param)))
+    for verify_param in PARAMS_TO_VERIFY_CI_IMAGE:
+        param_value = str(getattr(ci_image, verify_param))
+        if len(param_value) > 0:
+            args_command.append("--build-arg")
+            args_command.append(verify_param.upper() + "=" + param_value)
+    docker_cache = ci_image.docker_cache_ci_directive
+    if len(docker_cache) > 0:
+        args_command.extend(ci_image.docker_cache_ci_directive)
+    return args_command
+
+
+def construct_docker_command(ci_image: BuildParams) -> List[str]:
+    arguments = construct_arguments_docker_command(ci_image)
+    final_command = []
+    final_command.extend(["docker", "build"])
+    final_command.extend(arguments)
+    final_command.extend(["-t", ci_image.airflow_ci_image_name, "--target", "main", "."])
+    final_command.extend(["-f", str(Path(get_airflow_sources_root(), 'Dockerfile.ci').resolve())])
+    return final_command
+
+
+def build_image(verbose, **kwargs):
+    ci_image_params = BuildParams(filter_out_none(**kwargs))
+    is_cached, value = check_cache_and_write_if_not_cached(
+        "PYTHON_MAJOR_MINOR_VERSION", ci_image_params.python_version
+    )
+    if is_cached:
+        ci_image_params.python_version = value
+    cmd = construct_docker_command(ci_image_params)
+    output = run_command(cmd, verbose=verbose, text=True)
+    console.print(f"[blue]{output}")