Migrate Xcom model queries to sqlalchemy v2

Author: hussein-awala

airflow/models/taskinstance.py

@@ -3120,19 +3120,13 @@ def xcom_pull(
             session=session,
         )
 
-        # NOTE: Since we're only fetching the value field and not the whole
-        # class, the @recreate annotation does not kick in. Therefore we need to
-        # call XCom.deserialize_value() manually.
-
         # We are only pulling one single task.
         if (task_ids is None or isinstance(task_ids, str)) and not isinstance(map_indexes, Iterable):
-            first = query.with_entities(
-                XCom.run_id, XCom.task_id, XCom.dag_id, XCom.map_index, XCom.value
-            ).first()
+            first = query.one_or_none()
             if first is None:  # No matching XCom at all.
                 return default
             if map_indexes is not None or first.map_index < 0:
-                return XCom.deserialize_value(first)
+                return first.value
             query = query.order_by(None).order_by(XCom.map_index.asc())
             return LazyXComAccess.build_from_xcom_query(query)
 

airflow/models/xcom.py

@@ -39,6 +39,7 @@
     PrimaryKeyConstraint,
     String,
     delete,
+    select,
     text,
 )
 from sqlalchemy.ext.associationproxy import association_proxy
@@ -71,6 +72,7 @@
 
     import pendulum
     from sqlalchemy.orm import Session
+    from sqlalchemy.sql import Select
 
     from airflow.models.taskinstancekey import TaskInstanceKey
 
@@ -210,15 +212,15 @@ def set(
             message = "Passing 'execution_date' to 'XCom.set()' is deprecated. Use 'run_id' instead."
             warnings.warn(message, RemovedInAirflow3Warning, stacklevel=3)
             try:
-                dag_run_id, run_id = (
-                    session.query(DagRun.id, DagRun.run_id)
-                    .filter(DagRun.dag_id == dag_id, DagRun.execution_date == execution_date)
-                    .one()
-                )
+                dag_run_id, run_id = session.execute(
+                    select(DagRun.id, DagRun.run_id).where(
+                        DagRun.dag_id == dag_id, DagRun.execution_date == execution_date
+                    )
+                ).one()
             except NoResultFound:
                 raise ValueError(f"DAG run not found on DAG {dag_id!r} at {execution_date}") from None
         else:
-            dag_run_id = session.query(DagRun.id).filter_by(dag_id=dag_id, run_id=run_id).scalar()
+            dag_run_id = session.scalar(select(DagRun.id).filter_by(dag_id=dag_id, run_id=run_id))
             if dag_run_id is None:
                 raise ValueError(f"DAG run not found on DAG {dag_id!r} with ID {run_id!r}")
 
@@ -389,40 +391,99 @@ def get_one(
             raise ValueError("Exactly one of run_id or execution_date must be passed")
 
         if run_id:
-            query = BaseXCom.get_many(
+            stmt = BaseXCom._get_many_statement(
                 run_id=run_id,
                 key=key,
                 task_ids=task_id,
                 dag_ids=dag_id,
                 map_indexes=map_index,
                 include_prior_dates=include_prior_dates,
                 limit=1,
-                session=session,
             )
         elif execution_date is not None:
             message = "Passing 'execution_date' to 'XCom.get_one()' is deprecated. Use 'run_id' instead."
             warnings.warn(message, RemovedInAirflow3Warning, stacklevel=3)
 
             with warnings.catch_warnings():
                 warnings.simplefilter("ignore", RemovedInAirflow3Warning)
-                query = BaseXCom.get_many(
+                stmt = BaseXCom._get_many_statement(
                     execution_date=execution_date,
                     key=key,
                     task_ids=task_id,
                     dag_ids=dag_id,
                     map_indexes=map_index,
                     include_prior_dates=include_prior_dates,
                     limit=1,
-                    session=session,
                 )
         else:
             raise RuntimeError("Should not happen?")
 
-        result = query.with_entities(BaseXCom.value).first()
+        result = session.execute(stmt.with_only_columns(BaseXCom.value)).first()
         if result:
             return XCom.deserialize_value(result)
         return None
 
+    @staticmethod
+    def _get_many_statement(
+        execution_date: datetime.datetime | None = None,
+        key: str | None = None,
+        task_ids: str | Iterable[str] | None = None,
+        dag_ids: str | Iterable[str] | None = None,
+        map_indexes: int | Iterable[int] | None = None,
+        include_prior_dates: bool = False,
+        limit: int | None = None,
+        *,
+        run_id: str | None = None,
+    ) -> Select:
+        from airflow.models.dagrun import DagRun
+
+        if not exactly_one(execution_date is not None, run_id is not None):
+            raise ValueError(
+                f"Exactly one of run_id or execution_date must be passed. "
+                f"Passed execution_date={execution_date}, run_id={run_id}"
+            )
+        if execution_date is not None:
+            message = "Passing 'execution_date' to 'XCom.get_many()' is deprecated. Use 'run_id' instead."
+            warnings.warn(message, RemovedInAirflow3Warning, stacklevel=3)
+
+        stmt = select(BaseXCom).join(BaseXCom.dag_run)
+
+        if key:
+            stmt = stmt.where(BaseXCom.key == key)
+
+        if is_container(task_ids):
+            stmt = stmt.where(BaseXCom.task_id.in_(task_ids))
+        elif task_ids is not None:
+            stmt = stmt.where(BaseXCom.task_id == task_ids)
+
+        if is_container(dag_ids):
+            stmt = stmt.where(BaseXCom.dag_id.in_(dag_ids))
+        elif dag_ids is not None:
+            stmt = stmt.where(BaseXCom.dag_id == dag_ids)
+
+        if isinstance(map_indexes, range) and map_indexes.step == 1:
+            stmt = stmt.where(BaseXCom.map_index >= map_indexes.start, BaseXCom.map_index < map_indexes.stop)
+        elif is_container(map_indexes):
+            stmt = stmt.where(BaseXCom.map_index.in_(map_indexes))
+        elif map_indexes is not None:
+            stmt = stmt.where(BaseXCom.map_index == map_indexes)
+
+        if include_prior_dates:
+            if execution_date is not None:
+                stmt = stmt.where(DagRun.execution_date <= execution_date)
+            else:
+                dr = select(DagRun.execution_date).filter(DagRun.run_id == run_id).subquery()
+                stmt = stmt.where(BaseXCom.execution_date <= dr.c.execution_date)
+        elif execution_date is not None:
+            stmt = stmt.where(DagRun.execution_date == execution_date)
+        else:
+            stmt = stmt.where(BaseXCom.run_id == run_id)
+
+        stmt = stmt.order_by(DagRun.execution_date.desc(), BaseXCom.timestamp.desc())
+        if limit:
+            return stmt.limit(limit)
+        return stmt
+
     @overload
     @staticmethod
     def get_many(
@@ -498,56 +559,17 @@ def get_many(
 
         :sphinx-autoapi-skip:
         """
-        from airflow.models.dagrun import DagRun
-
-        if not exactly_one(execution_date is not None, run_id is not None):
-            raise ValueError(
-                f"Exactly one of run_id or execution_date must be passed. "
-                f"Passed execution_date={execution_date}, run_id={run_id}"
-            )
-        if execution_date is not None:
-            message = "Passing 'execution_date' to 'XCom.get_many()' is deprecated. Use 'run_id' instead."
-            warnings.warn(message, RemovedInAirflow3Warning, stacklevel=3)
-
-        query = session.query(BaseXCom).join(BaseXCom.dag_run)
-
-        if key:
-            query = query.filter(BaseXCom.key == key)
-
-        if is_container(task_ids):
-            query = query.filter(BaseXCom.task_id.in_(task_ids))
-        elif task_ids is not None:
-            query = query.filter(BaseXCom.task_id == task_ids)
-
-        if is_container(dag_ids):
-            query = query.filter(BaseXCom.dag_id.in_(dag_ids))
-        elif dag_ids is not None:
-            query = query.filter(BaseXCom.dag_id == dag_ids)
-
-        if isinstance(map_indexes, range) and map_indexes.step == 1:
-            query = query.filter(
-                BaseXCom.map_index >= map_indexes.start, BaseXCom.map_index < map_indexes.stop
-            )
-        elif is_container(map_indexes):
-            query = query.filter(BaseXCom.map_index.in_(map_indexes))
-        elif map_indexes is not None:
-            query = query.filter(BaseXCom.map_index == map_indexes)
-
-        if include_prior_dates:
-            if execution_date is not None:
-                query = query.filter(DagRun.execution_date <= execution_date)
-            else:
-                dr = session.query(DagRun.execution_date).filter(DagRun.run_id == run_id).subquery()
-                query = query.filter(BaseXCom.execution_date <= dr.c.execution_date)
-        elif execution_date is not None:
-            query = query.filter(DagRun.execution_date == execution_date)
-        else:
-            query = query.filter(BaseXCom.run_id == run_id)
-
-        query = query.order_by(DagRun.execution_date.desc(), BaseXCom.timestamp.desc())
-        if limit:
-            return query.limit(limit)
-        return query
+        stmt = BaseXCom._get_many_statement(
+            execution_date=execution_date,
+            key=key,
+            task_ids=task_ids,
+            dag_ids=dag_ids,
+            map_indexes=map_indexes,
+            include_prior_dates=include_prior_dates,
+            limit=limit,
+            run_id=run_id,
+        )
+        return session.scalars(stmt)
 
     @classmethod
     @provide_session
@@ -640,17 +662,15 @@ def clear(
         if execution_date is not None:
             message = "Passing 'execution_date' to 'XCom.clear()' is deprecated. Use 'run_id' instead."
             warnings.warn(message, RemovedInAirflow3Warning, stacklevel=3)
-            run_id = (
-                session.query(DagRun.run_id)
-                .filter(DagRun.dag_id == dag_id, DagRun.execution_date == execution_date)
-                .scalar()
+            run_id = session.scalar(
+                select(DagRun.run_id).where(DagRun.dag_id == dag_id, DagRun.execution_date == execution_date)
             )
 
-        query = session.query(BaseXCom).filter_by(dag_id=dag_id, task_id=task_id, run_id=run_id)
+        stmt = select(BaseXCom).filter_by(dag_id=dag_id, task_id=task_id, run_id=run_id)
         if map_index is not None:
-            query = query.filter_by(map_index=map_index)
+            stmt = stmt.filter_by(map_index=map_index)
 
-        for xcom in query:
+        for xcom in session.scalars(stmt):
             # print(f"Clearing XCOM {xcom} with value {xcom.value}")
             XCom.purge(xcom, session)
             session.delete(xcom)

tests/models/test_xcom.py

@@ -468,7 +468,7 @@ def test_xcom_get_many_from_prior_dates(self, session, tis_for_xcom_get_many_fro
             task_ids="task_1",
             include_prior_dates=True,
             session=session,
-        )
+        ).all()
 
         # The retrieved XComs should be ordered by logical date, latest first.
         assert [x.value for x in stored_xcoms] == [{"key2": "value2"}, {"key1": "value1"}]
@@ -488,7 +488,7 @@ def test_xcom_get_many_from_prior_dates_with_execution_date(
                 task_ids="task_1",
                 include_prior_dates=True,
                 session=session,
-            )
+            ).all()
 
         # The retrieved XComs should be ordered by logical date, latest first.
         assert [x.value for x in stored_xcoms] == [{"key2": "value2"}, {"key1": "value1"}]

tests/providers/common/io/xcom/test_backend.py

@@ -28,7 +28,7 @@
 from airflow.io.path import ObjectStoragePath
 from airflow.models.dagrun import DagRun
 from airflow.models.taskinstance import TaskInstance
-from airflow.models.xcom import BaseXCom, resolve_xcom_backend
+from airflow.models.xcom import resolve_xcom_backend
 from airflow.operators.empty import EmptyOperator
 from airflow.providers.common.io.xcom.backend import XComObjectStoreBackend
 from airflow.utils import timezone
@@ -151,20 +151,15 @@ def test_value_storage(self, task_instance, session):
             session=session,
         )
 
-        res = (
-            XCom.get_many(
-                key=XCOM_RETURN_KEY,
-                dag_ids=task_instance.dag_id,
-                task_ids=task_instance.task_id,
-                run_id=task_instance.run_id,
-                session=session,
-            )
-            .with_entities(BaseXCom.value)
-            .first()
-        )
+        res = XCom.get_many(
+            key=XCOM_RETURN_KEY,
+            dag_ids=task_instance.dag_id,
+            task_ids=task_instance.task_id,
+            run_id=task_instance.run_id,
+            session=session,
+        ).first()
 
-        data = BaseXCom.deserialize_value(res)
-        p = ObjectStoragePath(self.path) / XComObjectStoreBackend._get_key(data)
+        p = ObjectStoragePath(self.path) / XComObjectStoreBackend._get_key(res.value)
         assert p.exists() is True
 
         value = XCom.get_value(
@@ -197,20 +192,15 @@ def test_clear(self, task_instance, session):
             session=session,
         )
 
-        res = (
-            XCom.get_many(
-                key=XCOM_RETURN_KEY,
-                dag_ids=task_instance.dag_id,
-                task_ids=task_instance.task_id,
-                run_id=task_instance.run_id,
-                session=session,
-            )
-            .with_entities(BaseXCom.value)
-            .first()
-        )
+        res = XCom.get_many(
+            key=XCOM_RETURN_KEY,
+            dag_ids=task_instance.dag_id,
+            task_ids=task_instance.task_id,
+            run_id=task_instance.run_id,
+            session=session,
+        ).first()
 
-        data = BaseXCom.deserialize_value(res)
-        p = ObjectStoragePath(self.path) / XComObjectStoreBackend._get_key(data)
+        p = ObjectStoragePath(self.path) / XComObjectStoreBackend._get_key(res.value)
         assert p.exists() is True
 
         XCom.clear(
@@ -237,20 +227,15 @@ def test_compression(self, task_instance, session):
             session=session,
         )
 
-        res = (
-            XCom.get_many(
-                key=XCOM_RETURN_KEY,
-                dag_ids=task_instance.dag_id,
-                task_ids=task_instance.task_id,
-                run_id=task_instance.run_id,
-                session=session,
-            )
-            .with_entities(BaseXCom.value)
-            .first()
-        )
+        res = XCom.get_many(
+            key=XCOM_RETURN_KEY,
+            dag_ids=task_instance.dag_id,
+            task_ids=task_instance.task_id,
+            run_id=task_instance.run_id,
+            session=session,
+        ).first()
 
-        data = BaseXCom.deserialize_value(res)
-        p = ObjectStoragePath(self.path) / XComObjectStoreBackend._get_key(data)
+        p = ObjectStoragePath(self.path) / XComObjectStoreBackend._get_key(res.value)
         assert p.exists() is True
         assert p.suffix == ".gz"