Add SQLA-version-dependent dialect kwarg generator

Author: Dev-iL

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -75,7 +75,12 @@
 from airflow.utils.retries import MAX_DB_RETRIES, retry_db_transaction, run_with_db_retries
 from airflow.utils.session import NEW_SESSION, create_session, provide_session
 from airflow.utils.span_status import SpanStatus
-from airflow.utils.sqlalchemy import is_lock_not_available_error, prohibit_commit, with_row_locks
+from airflow.utils.sqlalchemy import (
+    is_lock_not_available_error,
+    prohibit_commit,
+    make_dialect_kwarg,
+    with_row_locks,
+)
 from airflow.utils.state import DagRunState, JobState, State, TaskInstanceState
 from airflow.utils.thread_safe_dict import ThreadSafeDict
 from airflow.utils.types import DagRunTriggeredByType, DagRunType
@@ -102,6 +107,8 @@
 TASK_STUCK_IN_QUEUED_RESCHEDULE_EVENT = "stuck in queued reschedule"
 """:meta private:"""
 
+dialect_kwarg = make_dialect_kwarg("mysql")
+
 
 class SchedulerDagBag:
     """
@@ -400,7 +407,7 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
 
             query = (
                 select(TI)
-                .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+                .with_hint(TI, "USE INDEX (ti_state)", **dialect_kwarg)
                 .join(TI.dag_run)
                 .where(DR.state == DagRunState.RUNNING)
                 .join(TI.dag_model)
@@ -2237,7 +2244,7 @@ def _find_task_instances_without_heartbeats(self, *, session: Session) -> list[T
             .options(selectinload(TI.dag_model))
             .options(selectinload(TI.dag_run))
             .options(selectinload(TI.dag_version))
-            .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+            .with_hint(TI, "USE INDEX (ti_state)", **dialect_kwarg)
             .join(DM, TI.dag_id == DM.dag_id)
             .where(
                 TI.state.in_((TaskInstanceState.RUNNING, TaskInstanceState.RESTARTING)),

airflow-core/src/airflow/migrations/versions/0036_3_0_0_add_name_field_to_dataset_model.py

@@ -48,7 +48,7 @@
 
 _STRING_COLUMN_TYPE = sa.String(length=1500).with_variant(
     sa.String(length=1500, collation="latin1_general_cs"),
-    dialect_name="mysql",
+    "mysql",
 )
 
 
@@ -128,7 +128,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-core/src/airflow/migrations/versions/0038_3_0_0_add_asset_active.py

@@ -39,7 +39,7 @@
 
 _STRING_COLUMN_TYPE = sa.String(length=1500).with_variant(
     sa.String(length=1500, collation="latin1_general_cs"),
-    dialect_name="mysql",
+    "mysql",
 )
 
 

airflow-core/src/airflow/migrations/versions/0039_3_0_0_tweak_assetaliasmodel_to_match_asset.py

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

airflow-core/src/airflow/models/dagrun.py

@@ -552,10 +552,11 @@ def get_running_dag_runs_to_examine(cls, session: Session) -> Query:
         """
         from airflow.models.backfill import BackfillDagRun
         from airflow.models.dag import DagModel
+        from airflow.utils.sqlalchemy import make_dialect_kwarg
 
         query = (
             select(cls)
-            .with_hint(cls, "USE INDEX (idx_dag_run_running_dags)", dialect_name="mysql")
+            .with_hint(cls, "USE INDEX (idx_dag_run_running_dags)", **make_dialect_kwarg("mysql"))
             .where(cls.state == DagRunState.RUNNING)
             .join(DagModel, DagModel.dag_id == cls.dag_id)
             .join(BackfillDagRun, BackfillDagRun.dag_run_id == DagRun.id, isouter=True)

airflow-core/src/airflow/utils/sqlalchemy.py

@@ -36,6 +36,8 @@
 from airflow.utils.timezone import make_naive, utc
 
 if TYPE_CHECKING:
+    from collections.abc import Iterable
+
     from kubernetes.client.models.v1_pod import V1Pod
     from sqlalchemy.exc import OperationalError
     from sqlalchemy.orm import Query, Session
@@ -448,3 +450,8 @@ def get_orm_mapper():
 
 def is_sqlalchemy_v1() -> bool:
     return version.parse(metadata.version("sqlalchemy")).major == 1
+
+
+def make_dialect_kwarg(dialect: str) -> dict[str, str | Iterable[str]]:
+    """Create an SQLAlchemy-version-aware dialect keyword argument."""
+    return {"dialect_name": dialect} if is_sqlalchemy_v1() else {"dialect_names": (dialect,)}