Skip to content

Commit 57dc55c

Browse files
authored
Add the upgrade_sqlalchemy breeze flag (#52559)
+ Fix some existing shellcheck violations
1 parent 4c7d43c commit 57dc55c

15 files changed

+200
-138
lines changed

Dockerfile.ci

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -995,13 +995,13 @@ function handle_mount_sources() {
995995
echo
996996
echo "${COLOR_BLUE}Mounted sources are removed, cleaning up mounted dist-info files${COLOR_RESET}"
997997
echo
998-
rm -rf /usr/local/lib/python${PYTHON_MAJOR_MINOR_VERSION}/site-packages/apache_airflow*.dist-info/
998+
rm -rf /usr/local/lib/python"${PYTHON_MAJOR_MINOR_VERSION}"/site-packages/apache_airflow*.dist-info/
999999
fi
10001000
}
10011001

10021002
function determine_airflow_to_use() {
10031003
USE_AIRFLOW_VERSION="${USE_AIRFLOW_VERSION:=""}"
1004-
if [[ ${USE_AIRFLOW_VERSION} == "" && ${USE_DISTRIBUTIONS_FROM_DIST=} != "true" ]]; then
1004+
if [[ "${USE_AIRFLOW_VERSION}" == "" && "${USE_DISTRIBUTIONS_FROM_DIST}" != "true" ]]; then
10051005
export PYTHONPATH=${AIRFLOW_SOURCES}
10061006
echo
10071007
echo "${COLOR_BLUE}Using airflow version from current sources${COLOR_RESET}"
@@ -1032,7 +1032,7 @@ function determine_airflow_to_use() {
10321032
# for the use in parallel runs in docker containers--no-cache is needed - otherwise there is
10331033
# possibility of overriding temporary environments by multiple parallel processes
10341034
uv run --no-cache /opt/airflow/scripts/in_container/install_development_dependencies.py \
1035-
--constraint https://raw.githubusercontent.com/apache/airflow/constraints-main/constraints-${PYTHON_MAJOR_MINOR_VERSION}.txt
1035+
--constraint https://raw.githubusercontent.com/apache/airflow/constraints-main/constraints-"${PYTHON_MAJOR_MINOR_VERSION}".txt
10361036
# Some packages might leave legacy typing module which causes test issues
10371037
# shellcheck disable=SC2086
10381038
${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} typing || true
@@ -1067,6 +1067,19 @@ function check_boto_upgrade() {
10671067
set +x
10681068
}
10691069

1070+
function check_upgrade_sqlalchemy() {
1071+
if [[ "${UPGRADE_SQLALCHEMY}" != "true" ]]; then
1072+
return
1073+
fi
1074+
echo
1075+
echo "${COLOR_BLUE}Upgrading sqlalchemy to the latest version to run tests with it${COLOR_RESET}"
1076+
echo
1077+
set -x
1078+
# shellcheck disable=SC2086
1079+
${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} --upgrade "sqlalchemy[asyncio]<2.1" "databricks-sqlalchemy>=2"
1080+
set +x
1081+
}
1082+
10701083
function check_downgrade_sqlalchemy() {
10711084
if [[ ${DOWNGRADE_SQLALCHEMY=} != "true" ]]; then
10721085
return
@@ -1200,6 +1213,7 @@ handle_mount_sources
12001213
determine_airflow_to_use
12011214
environment_initialization
12021215
check_boto_upgrade
1216+
check_upgrade_sqlalchemy
12031217
check_downgrade_sqlalchemy
12041218
check_downgrade_pendulum
12051219
check_force_lowest_dependencies

dev/breeze/doc/images/output_shell.svg

Lines changed: 28 additions & 24 deletions
Loading
Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
74abdb31412e7ee706db5f3ece5beabc
1+
963bec258918bc3fa63ac8fc0a1eb433

dev/breeze/doc/images/output_testing_core-tests.svg

Lines changed: 52 additions & 48 deletions
Loading
Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
13a8da75e6aa2615690b1c23c98fd8b8
1+
8179cec98112f74c7e612d9ce1534afc

dev/breeze/doc/images/output_testing_providers-tests.svg

Lines changed: 61 additions & 57 deletions
Loading
Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
7e1ef96b4f18b709f93d71e988cd2887
1+
fed6b33e1f70b07c02f02483b66eec5a

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -396,6 +396,12 @@ def _set_default_from_parent(ctx: click.core.Context, option: click.core.Option,
396396
is_flag=True,
397397
envvar="UPGRADE_BOTO",
398398
)
399+
option_upgrade_sqlalchemy = click.option(
400+
"--upgrade-sqlalchemy",
401+
help="Upgrade SQLAlchemy to the latest version.",
402+
is_flag=True,
403+
envvar="UPGRADE_SQLALCHEMY",
404+
)
399405
option_use_uv = click.option(
400406
"--use-uv/--no-use-uv",
401407
is_flag=True,

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@
6767
option_standalone_dag_processor,
6868
option_tty,
6969
option_upgrade_boto,
70+
option_upgrade_sqlalchemy,
7071
option_use_airflow_version,
7172
option_use_uv,
7273
option_uv_http_timeout,
@@ -314,6 +315,7 @@ def run(self):
314315
@option_warn_image_upgrade_needed
315316
@option_standalone_dag_processor
316317
@option_upgrade_boto
318+
@option_upgrade_sqlalchemy
317319
@option_use_airflow_version
318320
@option_allow_pre_releases
319321
@option_use_distributions_from_dist
@@ -373,6 +375,7 @@ def shell(
373375
test_type: str | None,
374376
tty: str,
375377
upgrade_boto: bool,
378+
upgrade_sqlalchemy: bool,
376379
use_airflow_version: str | None,
377380
allow_pre_releases: bool,
378381
use_distributions_from_dist: bool,
@@ -445,6 +448,7 @@ def shell(
445448
test_type=test_type,
446449
tty=tty,
447450
upgrade_boto=upgrade_boto,
451+
upgrade_sqlalchemy=upgrade_sqlalchemy,
448452
use_airflow_version=use_airflow_version,
449453
use_distributions_from_dist=use_distributions_from_dist,
450454
use_uv=use_uv,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,7 @@
169169
"name": "Upgrading/downgrading/removing selected packages",
170170
"options": [
171171
"--upgrade-boto",
172+
"--upgrade-sqlalchemy",
172173
"--downgrade-sqlalchemy",
173174
"--downgrade-pendulum",
174175
],

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@
5757
option_skip_cleanup,
5858
option_skip_db_tests,
5959
option_upgrade_boto,
60+
option_upgrade_sqlalchemy,
6061
option_use_airflow_version,
6162
option_verbose,
6263
)
@@ -620,6 +621,7 @@ def _verify_parallelism_parameters(
620621
@option_test_type_core_group
621622
@option_total_test_timeout
622623
@option_upgrade_boto
624+
@option_upgrade_sqlalchemy
623625
@option_use_airflow_version
624626
@option_allow_pre_releases
625627
@option_use_distributions_from_dist
@@ -685,6 +687,7 @@ def core_tests(**kwargs):
685687
@option_test_type_providers_group
686688
@option_total_test_timeout
687689
@option_upgrade_boto
690+
@option_upgrade_sqlalchemy
688691
@option_use_airflow_version
689692
@option_allow_pre_releases
690693
@option_use_distributions_from_dist
@@ -748,6 +751,7 @@ def task_sdk_tests(**kwargs):
748751
test_type=ALL_TEST_TYPE,
749752
total_test_timeout=DEFAULT_TOTAL_TEST_TIMEOUT,
750753
upgrade_boto=False,
754+
upgrade_sqlalchemy=False,
751755
use_airflow_version=None,
752756
use_distributions_from_dist=False,
753757
**kwargs,
@@ -1291,6 +1295,7 @@ def _run_test_command(
12911295
test_type: str,
12921296
total_test_timeout: int,
12931297
upgrade_boto: bool,
1298+
upgrade_sqlalchemy: bool,
12941299
use_airflow_version: str | None,
12951300
use_distributions_from_dist: bool,
12961301
use_xdist: bool,
@@ -1337,6 +1342,7 @@ def _run_test_command(
13371342
test_type=test_type,
13381343
test_group=test_group,
13391344
upgrade_boto=upgrade_boto,
1345+
upgrade_sqlalchemy=upgrade_sqlalchemy,
13401346
use_airflow_version=use_airflow_version,
13411347
use_distributions_from_dist=use_distributions_from_dist,
13421348
use_xdist=use_xdist,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@
6565
"name": "Upgrading/downgrading/removing selected packages",
6666
"options": [
6767
"--upgrade-boto",
68+
"--upgrade-sqlalchemy",
6869
"--downgrade-sqlalchemy",
6970
"--downgrade-pendulum",
7071
],

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -211,6 +211,7 @@ class ShellParams:
211211
test_group: GroupOfTests | None = None
212212
tty: str = "auto"
213213
upgrade_boto: bool = False
214+
upgrade_sqlalchemy: bool = False
214215
use_airflow_version: str | None = None
215216
use_distributions_from_dist: bool = False
216217
use_uv: bool = False
@@ -642,6 +643,7 @@ def env_variables_for_docker_commands(self) -> dict[str, str]:
642643
_set_var(_env, "TEST_TYPE", self.test_type, "")
643644
_set_var(_env, "TEST_GROUP", str(self.test_group.value) if self.test_group else "")
644645
_set_var(_env, "UPGRADE_BOTO", self.upgrade_boto)
646+
_set_var(_env, "UPGRADE_SQLALCHEMY", self.upgrade_sqlalchemy)
645647
_set_var(_env, "USE_AIRFLOW_VERSION", self.use_airflow_version, "")
646648
_set_var(_env, "USE_DISTRIBUTIONS_FROM_DIST", self.use_distributions_from_dist)
647649
_set_var(_env, "USE_UV", self.use_uv)

scripts/ci/docker-compose/devcontainer.env

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ START_AIRFLOW="false"
6767
SUSPENDED_PROVIDERS_FOLDERS=""
6868
TEST_TYPE=
6969
UPGRADE_BOTO="false"
70+
UPGRADE_SQLALCHEMY="false"
7071
UPGRADE_RANDOM_INDICATOR_STRING=""
7172
VERBOSE="false"
7273
VERBOSE_COMMANDS="false"

scripts/docker/entrypoint_ci.sh

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -201,14 +201,14 @@ function handle_mount_sources() {
201201
echo
202202
echo "${COLOR_BLUE}Mounted sources are removed, cleaning up mounted dist-info files${COLOR_RESET}"
203203
echo
204-
rm -rf /usr/local/lib/python${PYTHON_MAJOR_MINOR_VERSION}/site-packages/apache_airflow*.dist-info/
204+
rm -rf /usr/local/lib/python"${PYTHON_MAJOR_MINOR_VERSION}"/site-packages/apache_airflow*.dist-info/
205205
fi
206206
}
207207

208208
# Determine which airflow version to use
209209
function determine_airflow_to_use() {
210210
USE_AIRFLOW_VERSION="${USE_AIRFLOW_VERSION:=""}"
211-
if [[ ${USE_AIRFLOW_VERSION} == "" && ${USE_DISTRIBUTIONS_FROM_DIST=} != "true" ]]; then
211+
if [[ "${USE_AIRFLOW_VERSION}" == "" && "${USE_DISTRIBUTIONS_FROM_DIST}" != "true" ]]; then
212212
export PYTHONPATH=${AIRFLOW_SOURCES}
213213
echo
214214
echo "${COLOR_BLUE}Using airflow version from current sources${COLOR_RESET}"
@@ -239,7 +239,7 @@ function determine_airflow_to_use() {
239239
# for the use in parallel runs in docker containers--no-cache is needed - otherwise there is
240240
# possibility of overriding temporary environments by multiple parallel processes
241241
uv run --no-cache /opt/airflow/scripts/in_container/install_development_dependencies.py \
242-
--constraint https://raw.githubusercontent.com/apache/airflow/constraints-main/constraints-${PYTHON_MAJOR_MINOR_VERSION}.txt
242+
--constraint https://raw.githubusercontent.com/apache/airflow/constraints-main/constraints-"${PYTHON_MAJOR_MINOR_VERSION}".txt
243243
# Some packages might leave legacy typing module which causes test issues
244244
# shellcheck disable=SC2086
245245
${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} typing || true
@@ -275,6 +275,20 @@ function check_boto_upgrade() {
275275
set +x
276276
}
277277

278+
# Upgrade sqlalchemy to the latest version to run tests with it
279+
function check_upgrade_sqlalchemy() {
280+
if [[ "${UPGRADE_SQLALCHEMY}" != "true" ]]; then
281+
return
282+
fi
283+
echo
284+
echo "${COLOR_BLUE}Upgrading sqlalchemy to the latest version to run tests with it${COLOR_RESET}"
285+
echo
286+
set -x
287+
# shellcheck disable=SC2086
288+
${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} --upgrade "sqlalchemy[asyncio]<2.1" "databricks-sqlalchemy>=2"
289+
set +x
290+
}
291+
278292
# Download minimum supported version of sqlalchemy to run tests with it
279293
function check_downgrade_sqlalchemy() {
280294
if [[ ${DOWNGRADE_SQLALCHEMY=} != "true" ]]; then
@@ -411,6 +425,7 @@ handle_mount_sources
411425
determine_airflow_to_use
412426
environment_initialization
413427
check_boto_upgrade
428+
check_upgrade_sqlalchemy
414429
check_downgrade_sqlalchemy
415430
check_downgrade_pendulum
416431
check_force_lowest_dependencies

0 commit comments

Comments
 (0)