#!/usr/bin/env bash

# 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.

function set_verbose() {

if [[ ${VERBOSE_COMMANDS:="false"} == "true" ]]; then

set -x

else

set +x

fi

}

set_verbose

# shellcheck source=scripts/in_container/_in_container_script_init.sh

. "${AIRFLOW_SOURCES:-/opt/airflow}"/scripts/in_container/_in_container_script_init.sh

# This one is to workaround https://github.com/apache/airflow/issues/17546

# issue with /usr/lib/<MACHINE>-linux-gnu/libstdc++.so.6: cannot allocate memory in static TLS block

# We do not yet a more "correct" solution to the problem but in order to avoid raising new issues

# by users of the prod image, we implement the workaround now.

# The side effect of this is slightly (in the range of 100s of milliseconds) slower load for any

# binary started and a little memory used for Heap allocated by initialization of libstdc++

# This overhead is not happening for binaries that already link dynamically libstdc++

LD_PRELOAD="/usr/lib/$(uname -m)-linux-gnu/libstdc++.so.6"

export LD_PRELOAD

# Add "other" and "group" write permission to the tmp folder

# Note that it will also change permissions in the /tmp folder on the host

# but this is necessary to enable some of our CLI tools to work without errors

chmod 1777 /tmp

AIRFLOW_SOURCES=$(cd "${IN_CONTAINER_DIR}/../.." || exit 1; pwd)

PYTHON_MAJOR_MINOR_VERSION=${PYTHON_MAJOR_MINOR_VERSION:=3.10}

export AIRFLOW_HOME=${AIRFLOW_HOME:=${HOME}}

# Create folder where sqlite database file will be stored if it does not exist (which happens when

# scripts are running rather than breeze shell)

mkdir "${AIRFLOW_HOME}/sqlite" -p || true

ASSET_COMPILATION_WAIT_MULTIPLIER=${ASSET_COMPILATION_WAIT_MULTIPLIER:=1}

if [[ "${CI=}" == "true" ]]; then

export COLUMNS="202"

fi

# shellcheck disable=SC1091

. "${IN_CONTAINER_DIR}/check_connectivity.sh"

# Make sure that asset compilation is completed before we proceed

function wait_for_asset_compilation() {

if [[ -f "${AIRFLOW_SOURCES}/.build/ui/.asset_compile.lock" ]]; then

echo

echo "${COLOR_YELLOW}Waiting for asset compilation to complete in the background.${COLOR_RESET}"

echo

local counter=0

while [[ -f "${AIRFLOW_SOURCES}/.build/ui/.asset_compile.lock" ]]; do

if (( counter % 5 == 2 )); then

echo "${COLOR_BLUE}Still waiting .....${COLOR_RESET}"

fi

sleep 1

((counter=counter+1))

if [[ ${counter} == 30*$ASSET_COMPILATION_WAIT_MULTIPLIER ]]; then

echo

echo "${COLOR_YELLOW}The asset compilation is taking too long.${COLOR_YELLOW}"

echo """

If it does not complete soon, you might want to stop it and remove file lock:

* press Ctrl-C

* run 'rm ${AIRFLOW_SOURCES}/.build/ui/.asset_compile.lock'

"""

fi

if [[ ${counter} == 60*$ASSET_COMPILATION_WAIT_MULTIPLIER ]]; then

echo

echo "${COLOR_RED}The asset compilation is taking too long. Exiting.${COLOR_RED}"

echo "${COLOR_RED}refer to dev/breeze/doc/04_troubleshooting.rst for resolution steps.${COLOR_RED}"

echo

exit 1

fi

done

fi

if [ -f "${AIRFLOW_SOURCES}/.build/ui/asset_compile.out" ]; then

echo

echo "${COLOR_RED}The asset compilation failed. Exiting.${COLOR_RESET}"

echo

cat "${AIRFLOW_SOURCES}/.build/ui/asset_compile.out"

rm "${AIRFLOW_SOURCES}/.build/ui/asset_compile.out"

echo

exit 1

fi

}

# Initialize environment variables to their default values that are used for running tests and

# interactive shell.

function environment_initialization() {

if [[ ${SKIP_ENVIRONMENT_INITIALIZATION=} == "true" ]]; then

return

fi

echo

echo "${COLOR_BLUE}Running Initialization. Your basic configuration is:${COLOR_RESET}"

echo

echo " * ${COLOR_BLUE}Airflow home:${COLOR_RESET} ${AIRFLOW_HOME}"

echo " * ${COLOR_BLUE}Airflow sources:${COLOR_RESET} ${AIRFLOW_SOURCES}"

echo " * ${COLOR_BLUE}Airflow core SQL connection:${COLOR_RESET} ${AIRFLOW__DATABASE__SQL_ALCHEMY_CONN:=}"

if [[ ${BACKEND=} == "postgres" ]]; then

echo " * ${COLOR_BLUE}Airflow backend:${COLOR_RESET} Postgres: ${POSTGRES_VERSION}"

elif [[ ${BACKEND=} == "mysql" ]]; then

echo " * ${COLOR_BLUE}Airflow backend:${COLOR_RESET} MySQL: ${MYSQL_VERSION}"

elif [[ ${BACKEND=} == "sqlite" ]]; then

echo " * ${COLOR_BLUE}Airflow backend:${COLOR_RESET} Sqlite"

elif [[ ${BACKEND=} == "custom" ]]; then

local _conn_url="${AIRFLOW__DATABASE__SQL_ALCHEMY_CONN:-}"

local _masked_url

_masked_url=$(echo "${_conn_url}" | sed -E 's|://([^:]+):([^@]+)@|://\1:***@|')

echo " * ${COLOR_BLUE}Airflow backend:${COLOR_RESET} Custom (${_masked_url})"

fi

echo

if [[ ${STANDALONE_DAG_PROCESSOR=} == "true" ]]; then

echo

echo "${COLOR_BLUE}Forcing scheduler/standalone_dag_processor to True${COLOR_RESET}"

echo

export AIRFLOW__SCHEDULER__STANDALONE_DAG_PROCESSOR=True

fi

if [[ ${GO_WORKER=} == "true" ]]; then

echo

echo "${COLOR_BLUE}Starting go worker${COLOR_RESET}"

echo

export AIRFLOW__SCHEDULER__GO_WORKER=True

fi

RUN_TESTS=${RUN_TESTS:="false"}

CI=${CI:="false"}

# Added to have run-tests on path

export PATH=${PATH}:${AIRFLOW_SOURCES}:/usr/local/go/bin/

mkdir -pv "${AIRFLOW_HOME}/logs/"

# Change the default worker_concurrency for tests

export AIRFLOW__CELERY__WORKER_CONCURRENCY=8

set +e

# shellcheck source=scripts/in_container/configure_environment.sh

. "${IN_CONTAINER_DIR}/configure_environment.sh"

# shellcheck source=scripts/in_container/run_init_script.sh

. "${IN_CONTAINER_DIR}/run_init_script.sh"

"${IN_CONTAINER_DIR}/check_environment.sh"

ENVIRONMENT_EXIT_CODE=$?

set -e

if [[ ${ENVIRONMENT_EXIT_CODE} != 0 ]]; then

echo

echo "Error: check_environment returned ${ENVIRONMENT_EXIT_CODE}. Exiting."

echo

exit ${ENVIRONMENT_EXIT_CODE}

fi

mkdir -p /usr/lib/google-cloud-sdk/bin

touch /usr/lib/google-cloud-sdk/bin/gcloud

ln -s -f /usr/bin/gcloud /usr/lib/google-cloud-sdk/bin/gcloud

if [[ ${SKIP_SSH_SETUP="false"} == "false" ]]; then

# Set up ssh keys

echo 'yes' | ssh-keygen -t rsa -C your_email@youremail.com -m PEM -P '' -f ~/.ssh/id_rsa \

>"${AIRFLOW_HOME}/logs/ssh-keygen.log" 2>&1

cat ~/.ssh/id_rsa.pub >> ~/.ssh/authorized_keys

ln -s -f ~/.ssh/authorized_keys ~/.ssh/authorized_keys2

chmod 600 ~/.ssh/*

# SSH Service

sudo service ssh restart >/dev/null 2>&1

# Sometimes the server is not quick enough to load the keys!

while [[ $(ssh-keyscan -H localhost 2>/dev/null | wc -l) != "3" ]] ; do

echo "Not all keys yet loaded by the server"

sleep 0.05

done

ssh-keyscan -H localhost >> ~/.ssh/known_hosts 2>/dev/null

fi

if [[ ${INTEGRATION_LOCALSTACK:-"false"} == "true" ]]; then

echo

echo "${COLOR_BLUE}Configuring LocalStack integration${COLOR_RESET}"

echo

# Define LocalStack AWS configuration

declare -A localstack_config=(

["AWS_ENDPOINT_URL"]="http://localstack:4566"

["AWS_ACCESS_KEY_ID"]="test"

["AWS_SECRET_ACCESS_KEY"]="test"

["AWS_DEFAULT_REGION"]="us-east-1"

)

# Export each configuration variable and log it

for key in "${!localstack_config[@]}"; do

export "$key"="${localstack_config[$key]}"

echo " * ${COLOR_BLUE}${key}:${COLOR_RESET} ${localstack_config[$key]}"

done

echo

fi

cd "${AIRFLOW_SOURCES}"

# Temporarily add /opt/airflow/providers/standard/tests to PYTHONPATH in order to see example dags

# in the UI when testing in Breeze. This might be solved differently in the future

if [[ -d /opt/airflow/providers/standard/tests ]]; then

export PYTHONPATH=${PYTHONPATH=}:/opt/airflow/providers/standard/tests

fi

if [[ ${START_AIRFLOW:="false"} == "true" || ${START_AIRFLOW} == "True" ]]; then

if [[ ${BREEZE_DEBUG_CELERY_WORKER=} == "true" ]]; then

export AIRFLOW__CELERY__POOL=${AIRFLOW__CELERY__POOL:-solo}

fi

export AIRFLOW__CORE__LOAD_EXAMPLES=${LOAD_EXAMPLES}

if [[ ${SKIP_ASSETS_COMPILATION:="false"} == "false" ]]; then

wait_for_asset_compilation

fi

if [[ ${TERMINAL_MULTIPLEXER:="mprocs"} == "mprocs" ]]; then

# shellcheck source=scripts/in_container/bin/run_mprocs

exec run_mprocs

else

# shellcheck source=scripts/in_container/bin/run_tmux

exec run_tmux

fi

fi

}

# Handle mount sources

function handle_mount_sources() {

if [[ ${MOUNT_SOURCES=} == "remove" ]]; then

echo

echo "${COLOR_BLUE}Mounted sources are removed, cleaning up mounted dist-info files${COLOR_RESET}"

echo

rm -rf /usr/local/lib/python"${PYTHON_MAJOR_MINOR_VERSION}"/site-packages/apache_airflow*.dist-info/

fi

}

# Determine which airflow version to use

function determine_airflow_to_use() {

USE_AIRFLOW_VERSION="${USE_AIRFLOW_VERSION:=""}"

if [[ "${USE_AIRFLOW_VERSION}" == "" && "${USE_DISTRIBUTIONS_FROM_DIST=}" != "true" ]]; then

export PYTHONPATH=${AIRFLOW_SOURCES}

echo

echo "${COLOR_BLUE}Using airflow version from current sources${COLOR_RESET}"

echo

# Cleanup the logs, tmp when entering the environment

sudo rm -rf "${AIRFLOW_SOURCES}"/logs/*

sudo rm -rf "${AIRFLOW_SOURCES}"/tmp/*

mkdir -p "${AIRFLOW_SOURCES}"/logs/

mkdir -p "${AIRFLOW_SOURCES}"/tmp/

else

if [[ ${CLEAN_AIRFLOW_INSTALLATION=} == "true" ]]; then

echo

echo "${COLOR_BLUE}Uninstalling all packages first${COLOR_RESET}"

echo

# shellcheck disable=SC2086

${PACKAGING_TOOL_CMD} freeze | grep -ve "^-e" | grep -ve "^#" | grep -ve "^uv" | grep -v "@" | \

xargs ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS}

# Now install rich ad click first to use the installation script

# shellcheck disable=SC2086

${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} rich rich-click click \

--constraint https://raw.githubusercontent.com/apache/airflow/constraints-main/constraints-${PYTHON_MAJOR_MINOR_VERSION}.txt

fi

echo

echo "${COLOR_BLUE}Reinstalling all development dependencies${COLOR_RESET}"

echo

# Generate constraints from uv.lock and use them to install development dependencies

# via the Python script. --no-cache is needed - otherwise there is possibility of

# overriding temporary environments by multiple parallel processes

local constraint_file="/tmp/constraints-from-lock.txt"

uv export --frozen --no-hashes --no-emit-project --no-emit-workspace --no-editable --no-header \

--no-annotate > "${constraint_file}" 2>/dev/null || true

uv run --no-cache /opt/airflow/scripts/in_container/install_development_dependencies.py \

--constraint "${constraint_file}"

# Some packages might leave legacy typing module which causes test issues

# shellcheck disable=SC2086

${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} typing || true

echo

echo "${COLOR_BLUE}Installing airflow and providers ${COLOR_RESET}"

echo

python "${IN_CONTAINER_DIR}/install_airflow_and_providers.py"

fi

if [[ "${USE_AIRFLOW_VERSION}" =~ ^2.* ]]; then

# Remove auth manager setting

unset AIRFLOW__CORE__AUTH_MANAGER

fi

if [[ "${USE_AIRFLOW_VERSION}" =~ ^2\.2\..*|^2\.1\..*|^2\.0\..* && "${AIRFLOW__DATABASE__SQL_ALCHEMY_CONN=}" != "" ]]; then

# make sure old variable is used for older airflow versions

export AIRFLOW__CORE__SQL_ALCHEMY_CONN="${AIRFLOW__DATABASE__SQL_ALCHEMY_CONN}"

fi

}

# Upgrade boto3 and botocore to latest version to run Amazon tests with them

function check_boto_upgrade() {

if [[ ${UPGRADE_BOTO=} != "true" ]]; then

return

fi

echo

echo "${COLOR_BLUE}Upgrading boto3, botocore to latest version to run Amazon tests with them${COLOR_RESET}"

echo

# shellcheck disable=SC2086

${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} aiobotocore s3fs || true

# Urllib 2.6.0 breaks kubernetes client because kubernetes client uses deprecated in 2.0.0 and

# removed in 2.6.0 `getheaders()` call (instead of `headers` property.

# Tracked in https://github.com/kubernetes-client/python/issues/2477

# shellcheck disable=SC2086

${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} --upgrade "boto3<1.38.3" "botocore<1.38.3" "urllib3<2.6.0"

}

# Upgrade sqlalchemy to the latest version to run tests with it

function check_upgrade_sqlalchemy() {

# The python version constraint is a TEMPORARY WORKAROUND to exclude all FAB tests. Is should be removed once we

# upgrade FAB to v5 (PR #50960).

if [[ "${UPGRADE_SQLALCHEMY=}" != "true" || ${PYTHON_MAJOR_MINOR_VERSION} != "3.13" ]]; then

return

fi

echo

echo "${COLOR_BLUE}Upgrading sqlalchemy to the latest version to run tests with it${COLOR_RESET}"

echo

uv sync --all-packages --no-install-package apache-airflow-providers-fab --resolution highest \

--no-python-downloads --no-managed-python

}

# Download minimum supported version of sqlalchemy to run tests with it

function check_downgrade_sqlalchemy() {

if [[ ${DOWNGRADE_SQLALCHEMY=} != "true" ]]; then

return

fi

local min_sqlalchemy_version

min_sqlalchemy_version=$(grep "sqlalchemy\[asyncio\]>=" airflow-core/pyproject.toml | sed "s/.*>=\([0-9\.]*\).*/\1/" | xargs)

echo

echo "${COLOR_BLUE}Downgrading sqlalchemy to minimum supported version: ${min_sqlalchemy_version}${COLOR_RESET}"

echo

# shellcheck disable=SC2086

${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} "sqlalchemy[asyncio]==${min_sqlalchemy_version}"

echo

echo "${COLOR_BLUE}Running 'pip check'${COLOR_RESET}"

echo

# We use pip check here to make sure that whatever `uv` installs, is also "correct" according to `pip`

pip check

}

# Download minimum supported version of pendulum to run tests with it

function check_downgrade_pendulum() {

if [[ ${DOWNGRADE_PENDULUM=} != "true" || ${PYTHON_MAJOR_MINOR_VERSION} == "3.12" ]]; then

return

fi

local min_pendulum_version

min_pendulum_version=$(grep "pendulum>=" airflow-core/pyproject.toml | head -1 | sed "s/.*>=\([0-9\.]*\).*/\1/" | xargs)

echo

echo "${COLOR_BLUE}Downgrading pendulum to minimum supported version: ${min_pendulum_version}${COLOR_RESET}"

echo

# shellcheck disable=SC2086

${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} "pendulum==${min_pendulum_version}"

echo

echo "${COLOR_BLUE}Running 'pip check'${COLOR_RESET}"

echo

# We use pip check here to make sure that whatever `uv` installs, is also "correct" according to `pip`

pip check

}

# Check if we should run tests and run them if needed

function check_run_tests() {

if [[ ${RUN_TESTS=} != "true" ]]; then

return

fi

if [[ ${TEST_GROUP:=""} == "system" ]]; then

exec "${IN_CONTAINER_DIR}/run_system_tests.sh" "${@}"

else

exec "${IN_CONTAINER_DIR}/run_ci_tests.sh" "${@}"

fi

}

function check_force_lowest_dependencies() {

if [[ ${FORCE_LOWEST_DEPENDENCIES=} != "true" ]]; then

return

fi

if [[ ${TEST_TYPE=} =~ Providers\[.*\] ]]; then

local provider_id

# shellcheck disable=SC2001

provider_id=$(echo "${TEST_TYPE}" | sed 's/Providers\[\(.*\)\]/\1/')

echo

echo "${COLOR_BLUE}Forcing dependencies to lowest versions for provider: ${provider_id}${COLOR_RESET}"

echo

if ! /opt/airflow/scripts/in_container/is_provider_excluded.py "${provider_id}"; then

echo

echo "S${COLOR_YELLOW}Skipping ${provider_id} provider check on Python ${PYTHON_MAJOR_MINOR_VERSION}!${COLOR_RESET}"

echo

exit 0

fi

cd "${AIRFLOW_SOURCES}/providers/${provider_id/.//}" || exit 1

# --no-binary is needed in order to avoid libxml and xmlsec using different version of libxml2

# (binary lxml embeds its own libxml2, while xmlsec uses system one).

# See https://bugs.launchpad.net/lxml/+bug/2110068

local sync_successful="false"

for attempt in 1 2 3; do

echo "Attempt ${attempt} of syncing to lowest dependencies"

set -x

if UV_LOCK_TIMEOUT=200 uv sync --resolution lowest-direct --no-binary-package lxml --no-binary-package xmlsec --all-extras \

--no-python-downloads --no-managed-python; then

set +x

sync_successful="true"

break

fi

set +x

echo "Sleeping 30s"

sleep 30

echo "Attempt ${attempt} failed. Retrying..."

done

if [[ "${sync_successful}" != "true" ]]; then

echo "${COLOR_RED}Failed to sync lowest dependencies after 3 attempts.${COLOR_RESET}"

exit 1

fi

else

echo

echo "${COLOR_BLUE}Forcing dependencies to lowest versions for Airflow.${COLOR_RESET}"

echo

cd "${AIRFLOW_SOURCES}/airflow-core"

# --no-binary is needed in order to avoid libxml and xmlsec using different version of libxml2

# (binary lxml embeds its own libxml2, while xmlsec uses system one).

# See https://bugs.launchpad.net/lxml/+bug/2110068

uv sync --resolution lowest-direct --no-binary-package lxml --no-binary-package xmlsec --all-extras \

--no-python-downloads --no-managed-python

fi

}

function check_airflow_python_client_installation() {

if [[ ${INSTALL_AIRFLOW_PYTHON_CLIENT=} != "true" ]]; then

return

fi

python "${IN_CONTAINER_DIR}/install_airflow_python_client.py"

}

function initialize_db() {

# If we are going to start the api server OR we are a system test (which may or may not start the api server,

# depending on the Airflow version being used to run the tests), then migrate the DB.

if [[ ${START_API_SERVER_WITH_EXAMPLES=} == "true" || ${TEST_GROUP:=""} == "system" ]]; then

echo

echo "${COLOR_BLUE}Initializing database${COLOR_RESET}"

echo

airflow db migrate

echo

echo "${COLOR_BLUE}Database initialized${COLOR_RESET}"

fi

}

function start_api_server_with_examples(){

USE_AIRFLOW_VERSION="${USE_AIRFLOW_VERSION:=""}"

# Do not start the api server if either START_API_SERVER_WITH_EXAMPLES is false or the TEST_GROUP env var is not

# equal to "system".

if [[ ${START_API_SERVER_WITH_EXAMPLES=} != "true" && ${TEST_GROUP:=""} != "system" ]]; then

return

fi

# If the use Airflow version is set and it is <= 3.0.0 (which does not have the API server anyway) also return

if [[ ${USE_AIRFLOW_VERSION} != "" && ${USE_AIRFLOW_VERSION} < "3.0.0" ]]; then

return

fi

export AIRFLOW__CORE__LOAD_EXAMPLES=True

export AIRFLOW__API__EXPOSE_CONFIG=True

airflow dags reserialize

echo "Example dags parsing finished"

if airflow config get-value core auth_manager | grep -q "FabAuthManager"; then

echo "Create admin user"

airflow users create -u admin -p admin -f Thor -l Administrator -r Admin -e admin@email.domain || true

echo "Admin user created"

else

echo "Skipping user creation as auth manager different from Fab is used"

fi

echo

echo "${COLOR_BLUE}Starting airflow api server${COLOR_RESET}"

echo

if [[ ${START_API_SERVER_DAEMON:-"true"} == "false" ]]; then

airflow api-server --port 8080 &

else

airflow api-server --port 8080 --daemon

fi

echo

echo "${COLOR_BLUE}Waiting for api-server to start${COLOR_RESET}"

echo

check_service_connection "Airflow api-server" "run_nc localhost 8080" 100

EXIT_CODE=$?

if [[ ${EXIT_CODE} != 0 ]]; then

echo

echo "${COLOR_RED}Api server did not start properly${COLOR_RESET}"

echo

exit ${EXIT_CODE}

fi

echo

echo "${COLOR_BLUE}Airflow api-server started${COLOR_RESET}"

}

handle_mount_sources

determine_airflow_to_use

environment_initialization

check_boto_upgrade

check_upgrade_sqlalchemy

check_downgrade_sqlalchemy

check_downgrade_pendulum

check_force_lowest_dependencies

check_airflow_python_client_installation

initialize_db

start_api_server_with_examples

check_run_tests "${@}"

# If we are not running tests - just exec to bash shell

exec /bin/bash "${@}"