Skip to content

[v3-0-test] Add more diagnostics for Airflow installation inside CI i… #52232

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
potiuk merged 1 commit into from
Jun 25, 2025
Merged

[v3-0-test] Add more diagnostics for Airflow installation inside CI i… #52232

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 60 additions & 9 deletions scripts/in_container/install_airflow_and_providers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

from __future__ import annotations

import os
import re
import sys
from pathlib import Path
Expand Down Expand Up @@ -120,7 +121,9 @@ def calculate_constraints_location(
):
constraints_base = f"https://raw.githubusercontent.com/{github_repository}/{constraints_reference}"
location = f"{constraints_base}/{constraints_mode}-{python_version}.txt"
console.print(f"[info]Determined {'providers' if providers else 'airflow'} constraints as: {location}")
console.print(
f"[bright_blue]Determined {'providers' if providers else 'airflow'} constraints as: {location}"
)
return location


Expand All @@ -142,25 +145,27 @@ def get_airflow_constraints_location(
* DEFAULT_CONSTRAINT_BRANCH + constraints mode - as fallback
* constraints-main + constraints mode - as fallback
"""
console.print("[bright_blue]Determining airflow constraints location")
if airflow_skip_constraints:
console.print("[bright_blue]Skipping airflow constraints.")
return None
if airflow_constraints_location:
console.print(f"[info]Using constraints from location: {airflow_constraints_location}")
console.print(f"[bright_blue]Using constraints from location: {airflow_constraints_location}")
return airflow_constraints_location
if airflow_constraints_reference:
console.print(
f"[info]Building constraints location from constraints reference: {airflow_constraints_reference}"
f"[bright_blue]Building constraints location from constraints reference: {airflow_constraints_reference}"
)
elif airflow_package_version:
if re.match(r"[0-9]+\.[0-9]+\.[0-9]+[0-9a-z.]*|main|v[0-9]_.*", airflow_package_version):
airflow_constraints_reference = f"constraints-{airflow_package_version}"
console.print(
f"[info]Determined constraints reference from airflow package version "
f"[bright_blue]Determined constraints reference from airflow package version "
f"{airflow_package_version} as: {airflow_constraints_reference}"
)
else:
airflow_constraints_reference = default_constraints_branch
console.print(f"[info]Falling back tp: {default_constraints_branch}")
console.print(f"[bright_blue]Falling back to constraints branch: {default_constraints_branch}")
return calculate_constraints_location(
constraints_mode=airflow_constraints_mode,
constraints_reference=airflow_constraints_reference,
Expand All @@ -187,9 +192,10 @@ def get_providers_constraints_location(
* constraints-main + constraints mode - as fallback
"""
if providers_skip_constraints:
console.print("[bright_blue]Skipping providers constraints.")
return None
if providers_constraints_location:
console.print(f"[info]Using constraints from location: {providers_constraints_location}")
console.print(f"[bright_blue]Using constraints from location: {providers_constraints_location}")
return providers_constraints_location
if not providers_constraints_reference:
providers_constraints_reference = default_constraints_branch
Expand All @@ -212,6 +218,7 @@ class InstallationSpec(NamedTuple):
airflow_ctl_constraints_location: str | None
provider_distributions: list[str]
provider_constraints_location: str | None
pre_release: bool = os.environ.get("ALLOW_PRE_RELEASES", "false").lower() == "true"


ALLOWED_VCS_PROTOCOLS = ("git+file://", "git+https://", "git+ssh://", "git+http://", "git+git://", "git://")
Expand All @@ -235,6 +242,7 @@ def find_installation_spec(
use_airflow_version: str,
use_distributions_from_dist: bool,
) -> InstallationSpec:
console.print("[bright_blue]Finding installation specification")
if use_distributions_from_dist:
console.print("[bright_blue]Using distributions from dist folder")
else:
Expand All @@ -247,6 +255,8 @@ def find_installation_spec(
)
sys.exit(1)
extension = "whl" if distribution_format == "wheel" else "tar.gz"
pre_release = os.environ.get("ALLOW_PRE_RELEASES", "false").lower() == "true"
console.print("[bright_blue]Pre-release: ", pre_release)
if airflow_extras:
console.print(f"[bright_blue]Using airflow extras: {airflow_extras}")
airflow_extras = f"[{airflow_extras}]"
Expand All @@ -260,6 +270,7 @@ def find_installation_spec(
console.print("[red]This is not supported. Please use --mount-sources=remove flag in breeze.")
sys.exit(1)
if use_airflow_version in ["wheel", "sdist"] and use_distributions_from_dist:
console.print("[bright_blue]Finding specification from local distribution files in dist folder")
airflow_distribution_spec, airflow_core_distribution_spec = find_airflow_package(extension)
if use_airflow_version != distribution_format:
console.print(
Expand All @@ -270,6 +281,7 @@ def find_installation_spec(
if airflow_distribution_spec:
airflow_version = get_airflow_version_from_package(airflow_distribution_spec)
if airflow_version:
console.print(f"[bright_blue]Using airflow version retrieved from package: {airflow_version}")
airflow_constraints_location = get_airflow_constraints_location(
airflow_skip_constraints=airflow_skip_constraints,
airflow_constraints_mode=airflow_constraints_mode,
Expand Down Expand Up @@ -399,6 +411,18 @@ def find_installation_spec(
)
airflow_ctl_distribution = None
airflow_ctl_constraints_location = None
try:
from packaging.version import Version

if Version(use_airflow_version).is_prerelease:
console.print(
f"[yellow]Forcing pre-release installation for pre-release "
f"airflow version {use_airflow_version}"
)
pre_release = True
except Exception:
# In case parsing version fails here we just skip setting pre-release
pass
provider_distributions_list = []
if use_distributions_from_dist:
selected_providers_list = install_selected_providers.split(",") if install_selected_providers else []
Expand All @@ -407,7 +431,7 @@ def find_installation_spec(
else:
console.print("\n[bright_blue]No preselected providers\n")
provider_distributions_list = find_provider_distributions(extension, selected_providers_list)
return InstallationSpec(
installation_spec = InstallationSpec(
airflow_distribution=airflow_distribution_spec,
airflow_core_distribution=airflow_core_distribution_spec,
airflow_constraints_location=airflow_constraints_location,
Expand All @@ -425,12 +449,15 @@ def find_installation_spec(
github_repository=github_repository,
python_version=python_version,
),
pre_release=pre_release,
)
console.print("[bright_blue]Installation specification:[/]", installation_spec)
return installation_spec


ALLOWED_DISTRIBUTION_FORMAT = ["wheel", "sdist", "both"]
ALLOWED_CONSTRAINTS_MODE = ["constraints-source-providers", "constraints", "constraints-no-providers"]
ALLOWED_MOUNT_SOURCES = ["remove", "tests", "providers-and-tests"]
ALLOWED_MOUNT_SOURCES = ["remove", "tests", "providers-and-tests", "selected"]

INIT_CONTENT = '__path__ = __import__("pkgutil").extend_path(__path__, __name__) # type: ignore'
FUTURE_CONTENT = "from __future__ import annotations"
Expand Down Expand Up @@ -617,12 +644,33 @@ def install_airflow_and_providers(
use_distributions_from_dist=use_distributions_from_dist,
)
if installation_spec.airflow_distribution and install_airflow_with_constraints:
console.print("[bright_blue]Installing airflow with constraints")
console.print(
"[bright_blue]Airflow constraints location: ", installation_spec.airflow_constraints_location
)
console.print("[bright_blue]Airflow distribution", installation_spec.airflow_distribution)
console.print("[bright_blue]Airflow core distribution", installation_spec.airflow_core_distribution)
console.print(
"[bright_blue]Airflow task-sdk constraints location: ",
installation_spec.airflow_task_sdk_constraints_location,
)
console.print(
"[bright_blue]Airflow task-sdk distribution", installation_spec.airflow_task_sdk_distribution
)
console.print(
"[bright_blue]Airflow ctl constraints location: ",
installation_spec.airflow_ctl_constraints_location,
)
console.print("[bright_blue]Airflow ctl distribution", installation_spec.airflow_ctl_distribution)
base_install_airflow_cmd = [
"/usr/local/bin/uv",
"pip",
"install",
installation_spec.airflow_distribution,
]
if installation_spec.pre_release:
console.print("[bright_blue]Allowing pre-release versions of airflow")
base_install_airflow_cmd.append("--pre")
base_install_airflow_cmd.append(installation_spec.airflow_distribution)
console.print(
f"\n[bright_blue]Installing airflow distribution: {installation_spec.airflow_distribution} with constraints"
)
Expand Down Expand Up @@ -700,11 +748,14 @@ def install_airflow_and_providers(
)
run_command(base_install_airflow_cmd, github_actions=github_actions, check=True)
if installation_spec.provider_distributions or not install_airflow_with_constraints:
console.print("[bright_blue]Installing airflow without constraints")
base_install_providers_cmd = [
"/usr/local/bin/uv",
"pip",
"install",
]
if installation_spec.pre_release:
base_install_providers_cmd.append("--pre")
if not install_airflow_with_constraints and installation_spec.airflow_distribution:
console.print(
f"\n[bright_blue]Installing airflow distribution: {installation_spec.airflow_distribution} without constraints"
Expand Down