Add the upgrade_sqlalchemy breeze flag to run tests with SQLA2

Dockerfile.ci

@@ -995,13 +995,13 @@ function handle_mount_sources() {
         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/
+        rm -rf /usr/local/lib/python"${PYTHON_MAJOR_MINOR_VERSION}"/site-packages/apache_airflow*.dist-info/
     fi
 }
 
 function determine_airflow_to_use() {
     USE_AIRFLOW_VERSION="${USE_AIRFLOW_VERSION:=""}"
-    if [[ ${USE_AIRFLOW_VERSION} == "" && ${USE_DISTRIBUTIONS_FROM_DIST=} != "true" ]]; then
+    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}"
@@ -1032,7 +1032,7 @@ function determine_airflow_to_use() {
         # for the use in parallel runs in docker containers--no-cache is needed - otherwise there is
         # possibility of overriding temporary environments by multiple parallel processes
         uv run --no-cache /opt/airflow/scripts/in_container/install_development_dependencies.py \
-           --constraint https://raw.githubusercontent.com/apache/airflow/constraints-main/constraints-${PYTHON_MAJOR_MINOR_VERSION}.txt
+           --constraint https://raw.githubusercontent.com/apache/airflow/constraints-main/constraints-"${PYTHON_MAJOR_MINOR_VERSION}".txt
         # Some packages might leave legacy typing module which causes test issues
         # shellcheck disable=SC2086
         ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} typing || true
@@ -1067,6 +1067,19 @@ function check_boto_upgrade() {
     set +x
 }
 
+function check_upgrade_sqlalchemy() {
+    if [[ "${UPGRADE_SQLALCHEMY}" != "true" ]]; then
+        return
+    fi
+    echo
+    echo "${COLOR_BLUE}Upgrading sqlalchemy to the latest version to run tests with it${COLOR_RESET}"
+    echo
+    set -x
+    # shellcheck disable=SC2086
+    ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} --upgrade "sqlalchemy[asyncio]<2.1" "databricks-sqlalchemy>=2"
+    set +x
+}
+
 function check_downgrade_sqlalchemy() {
     if [[ ${DOWNGRADE_SQLALCHEMY=} != "true" ]]; then
         return
@@ -1200,6 +1213,7 @@ 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

dev/breeze/doc/images/output_shell.txt

@@ -1 +1 @@
-74abdb31412e7ee706db5f3ece5beabc
+963bec258918bc3fa63ac8fc0a1eb433

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

@@ -1 +1 @@
-13a8da75e6aa2615690b1c23c98fd8b8
+8179cec98112f74c7e612d9ce1534afc

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

@@ -1 +1 @@
-7e1ef96b4f18b709f93d71e988cd2887
+fed6b33e1f70b07c02f02483b66eec5a

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

@@ -396,6 +396,12 @@ def _set_default_from_parent(ctx: click.core.Context, option: click.core.Option,
     is_flag=True,
     envvar="UPGRADE_BOTO",
 )
+option_upgrade_sqlalchemy = click.option(
+    "--upgrade-sqlalchemy",
+    help="Upgrade SQLAlchemy to the latest version.",
+    is_flag=True,
+    envvar="UPGRADE_SQLALCHEMY",
+)
 option_use_uv = click.option(
     "--use-uv/--no-use-uv",
     is_flag=True,

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

@@ -67,6 +67,7 @@
     option_standalone_dag_processor,
     option_tty,
     option_upgrade_boto,
+    option_upgrade_sqlalchemy,
     option_use_airflow_version,
     option_use_uv,
     option_uv_http_timeout,
@@ -314,6 +315,7 @@ def run(self):
 @option_warn_image_upgrade_needed
 @option_standalone_dag_processor
 @option_upgrade_boto
+@option_upgrade_sqlalchemy
 @option_use_airflow_version
 @option_allow_pre_releases
 @option_use_distributions_from_dist
@@ -373,6 +375,7 @@ def shell(
     test_type: str | None,
     tty: str,
     upgrade_boto: bool,
+    upgrade_sqlalchemy: bool,
     use_airflow_version: str | None,
     allow_pre_releases: bool,
     use_distributions_from_dist: bool,
@@ -445,6 +448,7 @@ def shell(
         test_type=test_type,
         tty=tty,
         upgrade_boto=upgrade_boto,
+        upgrade_sqlalchemy=upgrade_sqlalchemy,
         use_airflow_version=use_airflow_version,
         use_distributions_from_dist=use_distributions_from_dist,
         use_uv=use_uv,

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

@@ -169,6 +169,7 @@
             "name": "Upgrading/downgrading/removing selected packages",
             "options": [
                 "--upgrade-boto",
+                "--upgrade-sqlalchemy",
                 "--downgrade-sqlalchemy",
                 "--downgrade-pendulum",
             ],

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

@@ -57,6 +57,7 @@
     option_skip_cleanup,
     option_skip_db_tests,
     option_upgrade_boto,
+    option_upgrade_sqlalchemy,
     option_use_airflow_version,
     option_verbose,
 )
@@ -620,6 +621,7 @@ def _verify_parallelism_parameters(
 @option_test_type_core_group
 @option_total_test_timeout
 @option_upgrade_boto
+@option_upgrade_sqlalchemy
 @option_use_airflow_version
 @option_allow_pre_releases
 @option_use_distributions_from_dist
@@ -685,6 +687,7 @@ def core_tests(**kwargs):
 @option_test_type_providers_group
 @option_total_test_timeout
 @option_upgrade_boto
+@option_upgrade_sqlalchemy
 @option_use_airflow_version
 @option_allow_pre_releases
 @option_use_distributions_from_dist
@@ -748,6 +751,7 @@ def task_sdk_tests(**kwargs):
         test_type=ALL_TEST_TYPE,
         total_test_timeout=DEFAULT_TOTAL_TEST_TIMEOUT,
         upgrade_boto=False,
+        upgrade_sqlalchemy=False,
         use_airflow_version=None,
         use_distributions_from_dist=False,
         **kwargs,
@@ -1291,6 +1295,7 @@ def _run_test_command(
     test_type: str,
     total_test_timeout: int,
     upgrade_boto: bool,
+    upgrade_sqlalchemy: bool,
     use_airflow_version: str | None,
     use_distributions_from_dist: bool,
     use_xdist: bool,
@@ -1337,6 +1342,7 @@ def _run_test_command(
         test_type=test_type,
         test_group=test_group,
         upgrade_boto=upgrade_boto,
+        upgrade_sqlalchemy=upgrade_sqlalchemy,
         use_airflow_version=use_airflow_version,
         use_distributions_from_dist=use_distributions_from_dist,
         use_xdist=use_xdist,

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

@@ -65,6 +65,7 @@
     "name": "Upgrading/downgrading/removing selected packages",
     "options": [
         "--upgrade-boto",
+        "--upgrade-sqlalchemy",
         "--downgrade-sqlalchemy",
         "--downgrade-pendulum",
     ],

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

@@ -211,6 +211,7 @@ class ShellParams:
     test_group: GroupOfTests | None = None
     tty: str = "auto"
     upgrade_boto: bool = False
+    upgrade_sqlalchemy: bool = False
     use_airflow_version: str | None = None
     use_distributions_from_dist: bool = False
     use_uv: bool = False
@@ -642,6 +643,7 @@ def env_variables_for_docker_commands(self) -> dict[str, str]:
         _set_var(_env, "TEST_TYPE", self.test_type, "")
         _set_var(_env, "TEST_GROUP", str(self.test_group.value) if self.test_group else "")
         _set_var(_env, "UPGRADE_BOTO", self.upgrade_boto)
+        _set_var(_env, "UPGRADE_SQLALCHEMY", self.upgrade_sqlalchemy)
         _set_var(_env, "USE_AIRFLOW_VERSION", self.use_airflow_version, "")
         _set_var(_env, "USE_DISTRIBUTIONS_FROM_DIST", self.use_distributions_from_dist)
         _set_var(_env, "USE_UV", self.use_uv)

scripts/ci/docker-compose/devcontainer.env

@@ -67,6 +67,7 @@ START_AIRFLOW="false"
 SUSPENDED_PROVIDERS_FOLDERS=""
 TEST_TYPE=
 UPGRADE_BOTO="false"
+UPGRADE_SQLALCHEMY="false"
 UPGRADE_RANDOM_INDICATOR_STRING=""
 VERBOSE="false"
 VERBOSE_COMMANDS="false"

scripts/docker/entrypoint_ci.sh

@@ -201,14 +201,14 @@ function handle_mount_sources() {
         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/
+        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
+    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}"
@@ -239,7 +239,7 @@ function determine_airflow_to_use() {
         # for the use in parallel runs in docker containers--no-cache is needed - otherwise there is
         # possibility of overriding temporary environments by multiple parallel processes
         uv run --no-cache /opt/airflow/scripts/in_container/install_development_dependencies.py \
-           --constraint https://raw.githubusercontent.com/apache/airflow/constraints-main/constraints-${PYTHON_MAJOR_MINOR_VERSION}.txt
+           --constraint https://raw.githubusercontent.com/apache/airflow/constraints-main/constraints-"${PYTHON_MAJOR_MINOR_VERSION}".txt
         # Some packages might leave legacy typing module which causes test issues
         # shellcheck disable=SC2086
         ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} typing || true
@@ -275,6 +275,20 @@ function check_boto_upgrade() {
     set +x
 }
 
+# Upgrade sqlalchemy to the latest version to run tests with it
+function check_upgrade_sqlalchemy() {
+    if [[ "${UPGRADE_SQLALCHEMY}" != "true" ]]; then
+        return
+    fi
+    echo
+    echo "${COLOR_BLUE}Upgrading sqlalchemy to the latest version to run tests with it${COLOR_RESET}"
+    echo
+    set -x
+    # shellcheck disable=SC2086
+    ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} --upgrade "sqlalchemy[asyncio]<2.1" "databricks-sqlalchemy>=2"
+    set +x
+}
+
 # Download minimum supported version of sqlalchemy to run tests with it
 function check_downgrade_sqlalchemy() {
     if [[ ${DOWNGRADE_SQLALCHEMY=} != "true" ]]; then
@@ -411,6 +425,7 @@ 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