diff --git a/scripts/in_container/install_airflow_and_providers.py b/scripts/in_container/install_airflow_and_providers.py index 49fb52ffb62ea..1b3341ef954eb 100755 --- a/scripts/in_container/install_airflow_and_providers.py +++ b/scripts/in_container/install_airflow_and_providers.py @@ -393,6 +393,11 @@ def find_installation_spec( github_repository=github_repository, python_version=python_version, ) + if airflow_version.startswith("2."): + # We need to make sure that pydantic is added as extra for Airflow 2.x + # because it's not in extras by default, and pydantic incompatible version is + # installed by default in the container and breaks serialization import + airflow_extras = _add_pydantic_to_extras(airflow_extras) if airflow_extras: airflow_distribution_spec += airflow_extras # We always install latest task-sdk - it's independent from Airflow @@ -490,6 +495,8 @@ def find_installation_spec( sys.exit(1) else: compile_ui_assets = False + if use_airflow_version.startswith("2"): + airflow_extras = _add_pydantic_to_extras(airflow_extras) console.print(f"\nInstalling airflow via apache-airflow=={use_airflow_version}") airflow_distribution_spec = f"apache-airflow{airflow_extras}=={use_airflow_version}" airflow_core_distribution_spec = ( @@ -559,6 +566,16 @@ def find_installation_spec( return installation_spec +def _add_pydantic_to_extras(airflow_extras: str) -> str: + console.print("[yellow]Adding pydantic to airflow extras for Airflow 2.x") + if airflow_extras: + airflow_extras = "[" + airflow_extras.strip("[]") + ",pydantic]" + else: + airflow_extras = "[pydantic]" + console.print(f"[yellow]New extras: {airflow_extras.strip('[]')}") + return airflow_extras + + def download_airflow_source_tarball(installation_spec: InstallationSpec): """Download Airflow source tarball from GitHub.""" if not installation_spec.compile_ui_assets: