SQLA2: fix mypy issue with getting the dialect name (apache#56941)

* SQLA: add util to get the dialect name type-safely

* Replace unsafe dialect retrievals with `get_dialect_name`

Author: Dev-iL

airflow-core/src/airflow/api_fastapi/core_api/services/ui/calendar.py

@@ -39,6 +39,7 @@
 from airflow.timetables._cron import CronMixin
 from airflow.timetables.base import DataInterval, TimeRestriction
 from airflow.timetables.simple import ContinuousTimetable
+from airflow.utils.sqlalchemy import get_dialect_name
 
 log = structlog.get_logger(logger_name=__name__)
 
@@ -92,7 +93,7 @@ def _get_historical_dag_runs(
         granularity: Literal["hourly", "daily"],
     ) -> tuple[list[CalendarTimeRangeResponse], Sequence[Row]]:
         """Get historical DAG runs from the database."""
-        dialect = session.bind.dialect.name
+        dialect = get_dialect_name(session)
 
         time_expression = self._get_time_truncation_expression(DagRun.logical_date, granularity, dialect)
 

airflow-core/src/airflow/assets/manager.py

@@ -38,6 +38,7 @@
 )
 from airflow.stats import Stats
 from airflow.utils.log.logging_mixin import LoggingMixin
+from airflow.utils.sqlalchemy import get_dialect_name
 
 if TYPE_CHECKING:
     from sqlalchemy.orm.session import Session
@@ -245,7 +246,7 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
         if not dags_to_queue:
             return
 
-        if session.bind.dialect.name == "postgresql":
+        if get_dialect_name(session) == "postgresql":
             return cls._postgres_queue_dagruns(asset_id, dags_to_queue, session)
         return cls._slow_path_queue_dagruns(asset_id, dags_to_queue, session)
 

airflow-core/src/airflow/dag_processing/collection.py

@@ -60,7 +60,7 @@
 from airflow.serialization.serialized_objects import BaseSerialization, LazyDeserializedDAG, SerializedDAG
 from airflow.triggers.base import BaseEventTrigger
 from airflow.utils.retries import MAX_DB_RETRIES, run_with_db_retries
-from airflow.utils.sqlalchemy import with_row_locks
+from airflow.utils.sqlalchemy import get_dialect_name, with_row_locks
 from airflow.utils.types import DagRunType
 
 if TYPE_CHECKING:
@@ -756,7 +756,7 @@ def activate_assets_if_possible(self, models: Iterable[AssetModel], *, session:
         there's a conflict. The scheduler makes a more comprehensive pass
         through all assets in ``_update_asset_orphanage``.
         """
-        if session.bind is not None and (dialect_name := session.bind.dialect.name) == "postgresql":
+        if (dialect_name := get_dialect_name(session)) == "postgresql":
             from sqlalchemy.dialects.postgresql import insert as postgresql_insert
 
             stmt: Any = postgresql_insert(AssetActive).on_conflict_do_nothing()

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

@@ -79,7 +79,14 @@
 from airflow.utils.retries import retry_db_transaction
 from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.utils.span_status import SpanStatus
-from airflow.utils.sqlalchemy import ExtendedJSON, UtcDateTime, mapped_column, nulls_first, with_row_locks
+from airflow.utils.sqlalchemy import (
+    ExtendedJSON,
+    UtcDateTime,
+    get_dialect_name,
+    mapped_column,
+    nulls_first,
+    with_row_locks,
+)
 from airflow.utils.state import DagRunState, State, TaskInstanceState
 from airflow.utils.strings import get_random_string
 from airflow.utils.thread_safe_dict import ThreadSafeDict
@@ -399,7 +406,7 @@ def duration(self) -> float | None:
     @duration.expression  # type: ignore[no-redef]
     @provide_session
     def duration(cls, session: Session = NEW_SESSION) -> Case:
-        dialect_name = session.bind.dialect.name
+        dialect_name = get_dialect_name(session)
         if dialect_name == "mysql":
             return func.timestampdiff(text("SECOND"), cls.start_date, cls.end_date)
 

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

@@ -40,7 +40,7 @@
 from airflow.triggers.deadline import PAYLOAD_BODY_KEY, PAYLOAD_STATUS_KEY, DeadlineCallbackTrigger
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.session import provide_session
-from airflow.utils.sqlalchemy import UtcDateTime, mapped_column
+from airflow.utils.sqlalchemy import UtcDateTime, get_dialect_name, mapped_column
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
@@ -411,7 +411,7 @@ def _evaluate_with(self, *, session: Session, **kwargs: Any) -> datetime | None:
             dag_id = kwargs["dag_id"]
 
             # Get database dialect to use appropriate time difference calculation
-            dialect = getattr(session.bind.dialect, "name", None)
+            dialect = get_dialect_name(session)
 
             # Create database-specific expression for calculating duration in seconds
             if dialect == "postgresql":

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

@@ -49,7 +49,7 @@
 from airflow.settings import COMPRESS_SERIALIZED_DAGS, json
 from airflow.utils.hashlib_wrapper import md5
 from airflow.utils.session import NEW_SESSION, provide_session
-from airflow.utils.sqlalchemy import UtcDateTime, mapped_column
+from airflow.utils.sqlalchemy import UtcDateTime, get_dialect_name, mapped_column
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
@@ -591,12 +591,13 @@ def get_dag_dependencies(cls, session: Session = NEW_SESSION) -> dict[str, list[
         """
         load_json: Callable | None
         if COMPRESS_SERIALIZED_DAGS is False:
-            if session.bind.dialect.name in ["sqlite", "mysql"]:
+            dialect = get_dialect_name(session)
+            if dialect in ["sqlite", "mysql"]:
                 data_col_to_select = func.json_extract(cls._data, "$.dag.dag_dependencies")
 
                 def load_json(deps_data):
                     return json.loads(deps_data) if deps_data else []
-            elif session.bind.dialect.name == "postgresql":
+            elif dialect == "postgresql":
                 # Use #> operator which works for both JSON and JSONB types
                 # Returns the JSON sub-object at the specified path
                 data_col_to_select = cls._data.op("#>")(literal('{"dag","dag_dependencies"}'))

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

@@ -37,7 +37,7 @@
 from airflow.triggers.base import BaseTaskEndEvent
 from airflow.utils.retries import run_with_db_retries
 from airflow.utils.session import NEW_SESSION, provide_session
-from airflow.utils.sqlalchemy import UtcDateTime, mapped_column, with_row_locks
+from airflow.utils.sqlalchemy import UtcDateTime, get_dialect_name, mapped_column, with_row_locks
 from airflow.utils.state import TaskInstanceState
 
 if TYPE_CHECKING:
@@ -231,7 +231,7 @@ def clean_unused(cls, session: Session = NEW_SESSION) -> None:
             .group_by(cls.id)
             .having(func.count(TaskInstance.trigger_id) == 0)
         )
-        if session.bind.dialect.name == "mysql":
+        if get_dialect_name(session) == "mysql":
             # MySQL doesn't support DELETE with JOIN, so we need to do it in two steps
             ids = session.scalars(ids).all()
         session.execute(

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

@@ -58,6 +58,13 @@ def mapped_column(*args, **kwargs):
         return Column(*args, **kwargs)
 
 
+def get_dialect_name(session: Session) -> str | None:
+    """Safely get the name of the dialect associated with the given session."""
+    if (bind := session.get_bind()) is None:
+        raise ValueError("No bind/engine is associated with the provided Session")
+    return getattr(bind.dialect, "name", None)
+
+
 class UtcDateTime(TypeDecorator):
     """
     Similar to :class:`~sqlalchemy.types.TIMESTAMP` with ``timezone=True`` option, with some differences.
@@ -312,7 +319,7 @@ def nulls_first(col, session: Session) -> dict[str, Any]:
     Other databases do not need it since NULL values are considered lower than
     any other values, and appear first when the order is ASC (ascending).
     """
-    if session.bind.dialect.name == "postgresql":
+    if get_dialect_name(session) == "postgresql":
         return nullsfirst(col)
     return col
 

airflow-core/tests/unit/utils/test_sqlalchemy.py

@@ -21,7 +21,6 @@
 import pickle
 from copy import deepcopy
 from unittest import mock
-from unittest.mock import MagicMock
 
 import pytest
 from kubernetes.client import models as k8s
@@ -37,6 +36,7 @@
 from airflow.utils.sqlalchemy import (
     ExecutorConfigType,
     ensure_pod_is_valid_after_unpickling,
+    get_dialect_name,
     is_sqlalchemy_v1,
     prohibit_commit,
     with_row_locks,
@@ -52,13 +52,40 @@
 TEST_POD = k8s.V1Pod(spec=k8s.V1PodSpec(containers=[k8s.V1Container(name="base")]))
 
 
+class TestGetDialectName:
+    def test_returns_dialect_name_when_present(self, mocker):
+        mock_session = mocker.Mock()
+        mock_bind = mocker.Mock()
+        mock_bind.dialect.name = "postgresql"
+        mock_session.get_bind.return_value = mock_bind
+
+        assert get_dialect_name(mock_session) == "postgresql"
+
+    def test_raises_when_no_bind(self, mocker):
+        mock_session = mocker.Mock()
+        mock_session.get_bind.return_value = None
+
+        with pytest.raises(ValueError, match="No bind/engine is associated"):
+            get_dialect_name(mock_session)
+
+    def test_returns_none_when_dialect_has_no_name(self, mocker):
+        mock_session = mocker.Mock()
+        mock_bind = mocker.Mock()
+        # simulate dialect object without `name` attribute
+        mock_bind.dialect = mock.Mock()
+        delattr(mock_bind.dialect, "name") if hasattr(mock_bind.dialect, "name") else None
+        mock_session.get_bind.return_value = mock_bind
+
+        assert get_dialect_name(mock_session) is None
+
+
 class TestSqlAlchemyUtils:
     def setup_method(self):
         session = Session()
 
         # make sure NOT to run in UTC. Only postgres supports storing
         # timezone information in the datetime field
-        if session.bind.dialect.name == "postgresql":
+        if get_dialect_name(session) == "postgresql":
             session.execute(text("SET timezone='Europe/Amsterdam'"))
 
         self.session = session
@@ -124,7 +151,7 @@ def test_process_bind_param_naive(self):
         dag.clear()
 
     @pytest.mark.parametrize(
-        "dialect, supports_for_update_of, use_row_level_lock_conf, expected_use_row_level_lock",
+        ("dialect", "supports_for_update_of", "use_row_level_lock_conf", "expected_use_row_level_lock"),
         [
             ("postgresql", True, True, True),
             ("postgresql", True, False, False),
@@ -192,7 +219,7 @@ def teardown_method(self):
 
 class TestExecutorConfigType:
     @pytest.mark.parametrize(
-        "input, expected",
+        ("input", "expected"),
         [
             ("anything", "anything"),
             (
@@ -206,13 +233,13 @@ class TestExecutorConfigType:
             ),
         ],
     )
-    def test_bind_processor(self, input, expected):
+    def test_bind_processor(self, input, expected, mocker):
         """
         The returned bind processor should pickle the object as is, unless it is a dictionary with
         a pod_override node, in which case it should run it through BaseSerialization.
         """
         config_type = ExecutorConfigType()
-        mock_dialect = MagicMock()
+        mock_dialect = mocker.MagicMock()
         mock_dialect.dbapi = None
         process = config_type.bind_processor(mock_dialect)
         assert pickle.loads(process(input)) == expected
@@ -239,13 +266,13 @@ def test_bind_processor(self, input, expected):
             ),
         ],
     )
-    def test_result_processor(self, input):
+    def test_result_processor(self, input, mocker):
         """
         The returned bind processor should pickle the object as is, unless it is a dictionary with
         a pod_override node whose value was serialized with BaseSerialization.
         """
         config_type = ExecutorConfigType()
-        mock_dialect = MagicMock()
+        mock_dialect = mocker.MagicMock()
         mock_dialect.dbapi = None
         process = config_type.result_processor(mock_dialect, None)
         result = process(input)
@@ -277,7 +304,7 @@ def __eq__(self, other):
         assert instance.compare_values(a, a) is False
         assert instance.compare_values("a", "a") is True
 
-    def test_result_processor_bad_pickled_obj(self):
+    def test_result_processor_bad_pickled_obj(self, mocker):
         """
         If unpickled obj is missing attrs that curr lib expects
         """
@@ -309,7 +336,7 @@ def test_result_processor_bad_pickled_obj(self):
 
         # get the result processor method
         config_type = ExecutorConfigType()
-        mock_dialect = MagicMock()
+        mock_dialect = mocker.MagicMock()
         mock_dialect.dbapi = None
         process = config_type.result_processor(mock_dialect, None)
 
@@ -322,13 +349,13 @@ def test_result_processor_bad_pickled_obj(self):
 
 
 @pytest.mark.parametrize(
-    "mock_version, expected_result",
+    ("mock_version", "expected_result"),
     [
         ("1.0.0", True),  # Test 1: v1 identified as v1
         ("2.3.4", False),  # Test 2: v2 not identified as v1
     ],
 )
-def test_is_sqlalchemy_v1(mock_version, expected_result):
-    with mock.patch("airflow.utils.sqlalchemy.metadata") as mock_metadata:
-        mock_metadata.version.return_value = mock_version
-        assert is_sqlalchemy_v1() == expected_result
+def test_is_sqlalchemy_v1(mock_version, expected_result, mocker):
+    mock_metadata = mocker.patch("airflow.utils.sqlalchemy.metadata")
+    mock_metadata.version.return_value = mock_version
+    assert is_sqlalchemy_v1() == expected_result