Upgrade dependencies to allow us to use SQLAlchemy v2 in Airflow 3.0/main

Dockerfile

@@ -429,6 +429,7 @@ COPY <<"EOF" /install_airflow_dependencies_from_branch_tip.sh
 : "${INSTALL_POSTGRES_CLIENT:?Should be true or false}"
 
 function install_airflow_dependencies_from_branch_tip() {
+    set -e
     echo
     echo "${COLOR_BLUE}Installing airflow from ${AIRFLOW_BRANCH}. It is used to cache dependencies${COLOR_RESET}"
     echo
@@ -456,7 +457,7 @@ function install_airflow_dependencies_from_branch_tip() {
     # Uninstall airflow and providers to keep only the dependencies. In the future when
     # planned https://github.com/pypa/pip/issues/11440 is implemented in pip we might be able to use this
     # flag and skip the remove step.
-    pip freeze | grep apache-airflow-providers | xargs ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} || true
+    pip freeze | (grep apache-airflow-providers || true) | xargs --no-run-if-empty ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS}
     set +x
     echo
     echo "${COLOR_BLUE}Uninstalling just airflow. Dependencies remain. Now target airflow can be reinstalled using mostly cached dependencies${COLOR_RESET}"
@@ -473,14 +474,14 @@ function install_airflow_dependencies_from_branch_tip() {
     # and increase the AIRFLOW_CI_BUILD_EPOCH in Dockerfile.ci to make sure your cache is rebuilt.
     local DEPENDENCIES_TO_REMOVE
     # IMPORTANT!! Make sure to increase AIRFLOW_CI_BUILD_EPOCH in Dockerfile.ci when you remove a dependency from that list
-    DEPENDENCIES_TO_REMOVE=()
+    DEPENDENCIES_TO_REMOVE=("sqlalchemy-redshift")
     if [[ "${DEPENDENCIES_TO_REMOVE[*]}" != "" ]]; then
         echo
         echo "${COLOR_BLUE}Uninstalling just removed dependencies (temporary until cache refreshes)${COLOR_RESET}"
         echo "${COLOR_BLUE}Dependencies to uninstall: ${DEPENDENCIES_TO_REMOVE[*]}${COLOR_RESET}"
         echo
         set +x
-        ${PACKAGING_TOOL_CMD} uninstall "${DEPENDENCIES_TO_REMOVE[@]}" || true
+        ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} "${DEPENDENCIES_TO_REMOVE[@]}"
         set -x
         # make sure that the dependency is not needed by something else
         pip check
@@ -901,7 +902,7 @@ function install_airflow() {
         echo "${COLOR_BLUE}Remove airflow and all provider packages installed before potentially${COLOR_RESET}"
         echo
         set -x
-        ${PACKAGING_TOOL_CMD} freeze | grep apache-airflow | xargs ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} 2>/dev/null || true
+        ${PACKAGING_TOOL_CMD} freeze | grep apache-airflow | xargs --no-run-if-empty ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} 2>/dev/null || true
         set +x
         echo
         echo "${COLOR_BLUE}Installing all packages in eager upgrade mode. Installation method: ${AIRFLOW_INSTALLATION_METHOD}${COLOR_RESET}"

Dockerfile.ci

@@ -375,6 +375,7 @@ COPY <<"EOF" /install_airflow_dependencies_from_branch_tip.sh
 : "${INSTALL_POSTGRES_CLIENT:?Should be true or false}"
 
 function install_airflow_dependencies_from_branch_tip() {
+    set -e
     echo
     echo "${COLOR_BLUE}Installing airflow from ${AIRFLOW_BRANCH}. It is used to cache dependencies${COLOR_RESET}"
     echo
@@ -402,7 +403,7 @@ function install_airflow_dependencies_from_branch_tip() {
     # Uninstall airflow and providers to keep only the dependencies. In the future when
     # planned https://github.com/pypa/pip/issues/11440 is implemented in pip we might be able to use this
     # flag and skip the remove step.
-    pip freeze | grep apache-airflow-providers | xargs ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} || true
+    pip freeze | (grep apache-airflow-providers || true) | xargs --no-run-if-empty ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS}
     set +x
     echo
     echo "${COLOR_BLUE}Uninstalling just airflow. Dependencies remain. Now target airflow can be reinstalled using mostly cached dependencies${COLOR_RESET}"
@@ -419,14 +420,14 @@ function install_airflow_dependencies_from_branch_tip() {
     # and increase the AIRFLOW_CI_BUILD_EPOCH in Dockerfile.ci to make sure your cache is rebuilt.
     local DEPENDENCIES_TO_REMOVE
     # IMPORTANT!! Make sure to increase AIRFLOW_CI_BUILD_EPOCH in Dockerfile.ci when you remove a dependency from that list
-    DEPENDENCIES_TO_REMOVE=()
+    DEPENDENCIES_TO_REMOVE=("sqlalchemy-redshift")
     if [[ "${DEPENDENCIES_TO_REMOVE[*]}" != "" ]]; then
         echo
         echo "${COLOR_BLUE}Uninstalling just removed dependencies (temporary until cache refreshes)${COLOR_RESET}"
         echo "${COLOR_BLUE}Dependencies to uninstall: ${DEPENDENCIES_TO_REMOVE[*]}${COLOR_RESET}"
         echo
         set +x
-        ${PACKAGING_TOOL_CMD} uninstall "${DEPENDENCIES_TO_REMOVE[@]}" || true
+        ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} "${DEPENDENCIES_TO_REMOVE[@]}"
         set -x
         # make sure that the dependency is not needed by something else
         pip check
@@ -700,7 +701,7 @@ function install_airflow() {
         echo "${COLOR_BLUE}Remove airflow and all provider packages installed before potentially${COLOR_RESET}"
         echo
         set -x
-        ${PACKAGING_TOOL_CMD} freeze | grep apache-airflow | xargs ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} 2>/dev/null || true
+        ${PACKAGING_TOOL_CMD} freeze | grep apache-airflow | xargs --no-run-if-empty ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} 2>/dev/null || true
         set +x
         echo
         echo "${COLOR_BLUE}Installing all packages in eager upgrade mode. Installation method: ${AIRFLOW_INSTALLATION_METHOD}${COLOR_RESET}"
@@ -1264,7 +1265,7 @@ ARG AIRFLOW_CONSTRAINTS_LOCATION=""
 ARG DEFAULT_CONSTRAINTS_BRANCH="constraints-main"
 # By changing the epoch we can force reinstalling Airflow and pip all dependencies
 # It can also be overwritten manually by setting the AIRFLOW_CI_BUILD_EPOCH environment variable.
-ARG AIRFLOW_CI_BUILD_EPOCH="10"
+ARG AIRFLOW_CI_BUILD_EPOCH="11"
 ARG AIRFLOW_PRE_CACHED_PIP_PACKAGES="true"
 ARG AIRFLOW_PIP_VERSION=24.2
 ARG AIRFLOW_UV_VERSION=0.4.17
@@ -1291,7 +1292,6 @@ ARG AIRFLOW_VERSION=""
 ARG ADDITIONAL_PIP_INSTALL_FLAGS=""
 
 ARG AIRFLOW_PIP_VERSION=24.2
-ARG AIRFLOW_UV_VERSION=0.4.7
 ARG AIRFLOW_USE_UV="true"
 
 ENV AIRFLOW_REPO=${AIRFLOW_REPO}\

airflow/migrations/versions/0035_3_0_0_add_name_field_to_dataset_model.py

@@ -49,7 +49,7 @@
 
 _STRING_COLUMN_TYPE = sa.String(length=1500).with_variant(
     sa.String(length=1500, collation="latin1_general_cs"),
-    dialect_name="mysql",
+    "mysql",
 )
 
 
@@ -85,7 +85,7 @@ def downgrade():
             "uri",
             type_=sa.String(length=3000).with_variant(
                 sa.String(length=3000, collation="latin1_general_cs"),
-                dialect_name="mysql",
+                "mysql",
             ),
             nullable=False,
         )

airflow/models/base.py

@@ -20,7 +20,7 @@
 from typing import TYPE_CHECKING, Any
 
 from sqlalchemy import Column, Integer, MetaData, String, text
-from sqlalchemy.orm import registry
+from sqlalchemy.orm import DeclarativeBase
 
 from airflow.configuration import conf
 
@@ -45,13 +45,23 @@ def _get_schema():
 
 
 metadata = MetaData(schema=_get_schema(), naming_convention=naming_convention)
-mapper_registry = registry(metadata=metadata)
 _sentinel = object()
 
 if TYPE_CHECKING:
     Base = Any
 else:
-    Base = mapper_registry.generate_base()
+
+    class Base(DeclarativeBase):
+        """
+        Base class to ease transition to SQLAv2.
+
+        :meta private:
+        """
+
+        metadata = metadata
+        # https://docs.sqlalchemy.org/en/20/changelog/migration_20.html#migration-20-step-six
+        __allow_unmapped__ = True
+
 
 ID_LEN = 250
 

airflow/utils/db.py

@@ -53,6 +53,7 @@
     select,
     text,
 )
+from sqlalchemy.engine import Engine
 
 import airflow
 from airflow import settings
@@ -733,6 +734,7 @@ def create_default_connections(session: Session = NEW_SESSION):
     )
 
 
+@contextlib.contextmanager
 def _get_flask_db(sql_database_uri):
     from flask import Flask
     from flask_sqlalchemy import SQLAlchemy
@@ -744,7 +746,8 @@ def _get_flask_db(sql_database_uri):
     flask_app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
     db = SQLAlchemy(flask_app)
     AirflowDatabaseSessionInterface(app=flask_app, db=db, table="session", key_prefix="")
-    return db
+    with flask_app.app_context():
+        yield db
 
 
 def _create_db_from_orm(session):
@@ -753,8 +756,8 @@ def _create_db_from_orm(session):
     from airflow.models.base import Base
 
     def _create_flask_session_tbl(sql_database_uri):
-        db = _get_flask_db(sql_database_uri)
-        db.create_all()
+        with _get_flask_db(sql_database_uri) as db:
+            db.create_all()
 
     with create_global_lock(session=session, lock=DBLocks.MIGRATIONS):
         engine = session.get_bind().engine
@@ -1284,8 +1287,8 @@ def drop_airflow_models(connection):
     from airflow.models.base import Base
 
     Base.metadata.drop_all(connection)
-    db = _get_flask_db(connection.engine.url)
-    db.drop_all()
+    with _get_flask_db(connection.engine.url) as db:
+        db.drop_all()
     # alembic adds significant import time, so we import it lazily
     from alembic.migration import MigrationContext
 
@@ -1340,13 +1343,18 @@ def create_global_lock(
     lock_timeout: int = 1800,
 ) -> Generator[None, None, None]:
     """Contextmanager that will create and teardown a global db lock."""
-    conn = session.get_bind().connect()
+    bind = session.get_bind()
+    if isinstance(bind, Engine):
+        conn = bind.connect()
+    else:
+        conn = bind
     dialect = conn.dialect
+    mysql_supports_locks = dialect.name == "mysql" and dialect.server_version_info and dialect.server_version_info >= (5, 6)
     try:
         if dialect.name == "postgresql":
             conn.execute(text("SET LOCK_TIMEOUT to :timeout"), {"timeout": lock_timeout})
             conn.execute(text("SELECT pg_advisory_lock(:id)"), {"id": lock.value})
-        elif dialect.name == "mysql" and dialect.server_version_info >= (5, 6):
+        elif mysql_supports_locks:
             conn.execute(text("SELECT GET_LOCK(:id, :timeout)"), {"id": str(lock), "timeout": lock_timeout})
 
         yield
@@ -1356,7 +1364,7 @@ def create_global_lock(
             (unlocked,) = conn.execute(text("SELECT pg_advisory_unlock(:id)"), {"id": lock.value}).fetchone()
             if not unlocked:
                 raise RuntimeError("Error releasing DB lock!")
-        elif dialect.name == "mysql" and dialect.server_version_info >= (5, 6):
+        elif mysql_supports_locks:
             conn.execute(text("select RELEASE_LOCK(:id)"), {"id": str(lock)})
 
 

dev/breeze/src/airflow_breeze/global_constants.py

@@ -508,7 +508,7 @@ def get_airflow_extras():
     # END OF EXTRAS LIST UPDATED BY PRE COMMIT
 ]
 
-CHICKEN_EGG_PROVIDERS = " ".join(["standard"])
+CHICKEN_EGG_PROVIDERS = " ".join(["standard", "fab", "amazon"])
 
 
 BASE_PROVIDERS_COMPATIBILITY_CHECKS: list[dict[str, str | list[str]]] = [

dev/breeze/tests/test_packages.py

@@ -169,7 +169,7 @@ def test_get_documentation_package_path():
             """
     "apache-airflow-providers-common-compat>=1.2.0",
     "apache-airflow>=2.9.0",
-    "flask-appbuilder==4.5.0",
+    "flask-appbuilder==5.0.0rc1",
     "flask-login>=0.6.2",
     "flask>=2.2,<2.3",
     "google-re2>=1.0",
@@ -183,7 +183,7 @@ def test_get_documentation_package_path():
             """
     "apache-airflow-providers-common-compat>=1.2.0.dev0",
     "apache-airflow>=2.9.0.dev0",
-    "flask-appbuilder==4.5.0",
+    "flask-appbuilder==5.0.0rc1",
     "flask-login>=0.6.2",
     "flask>=2.2,<2.3",
     "google-re2>=1.0",
@@ -197,7 +197,7 @@ def test_get_documentation_package_path():
             """
     "apache-airflow-providers-common-compat>=1.2.0b0",
     "apache-airflow>=2.9.0b0",
-    "flask-appbuilder==4.5.0",
+    "flask-appbuilder==5.0.0rc1",
     "flask-login>=0.6.2",
     "flask>=2.2,<2.3",
     "google-re2>=1.0",

docs/apache-airflow-providers-amazon/index.rst

@@ -119,7 +119,6 @@ PIP package                                 Version required
 ``watchtower``                              ``>=3.0.0,!=3.3.0,<4``
 ``jsonpath_ng``                             ``>=1.5.3``
 ``redshift_connector``                      ``>=2.0.918``
-``sqlalchemy_redshift``                     ``>=0.8.6``
 ``asgiref``                                 ``>=2.3.0``
 ``PyAthena``                                ``>=3.0.10``
 ``jmespath``                                ``>=0.7.0``

docs/apache-airflow/img/airflow_erd.sha256

@@ -1 +1 @@
-8bd129828ba299ef05d70305eee66d15b6c0c79dc6ae82f654b9657464e3682a
+c4498d5a4d0f05418a13d74b267539927c3f89860610c5973fd040d2b34038e2