fix(serialized_dag): handle compressing serialized_dag for get_dag_dependencies (#48924)

Author: Lee-W

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

@@ -23,7 +23,7 @@
 import zlib
 from collections.abc import Iterable, Iterator, Sequence
 from datetime import timedelta
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Callable, Literal
 
 import sqlalchemy_jsonfield
 import uuid6
@@ -635,39 +635,43 @@ def get_dag_dependencies(cls, session: Session = NEW_SESSION) -> dict[str, list[
 
         :param session: ORM Session
         """
+        load_json: Callable | None
+        if COMPRESS_SERIALIZED_DAGS is False:
+            if session.bind.dialect.name 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 []
+            else:
+                data_col_to_select = func.json_extract_path(cls._data, "dag", "dag_dependencies")
+                load_json = None
+        else:
+            data_col_to_select = cls._data_compressed
+
+            def load_json(deps_data):
+                return json.loads(zlib.decompress(deps_data))["dag"]["dag_dependencies"] if deps_data else []
+
         latest_sdag_subquery = (
             select(cls.dag_id, func.max(cls.created_at).label("max_created")).group_by(cls.dag_id).subquery()
         )
-        if session.bind.dialect.name in ["sqlite", "mysql"]:
-            query = session.execute(
-                select(cls.dag_id, func.json_extract(cls._data, "$.dag.dag_dependencies"))
-                .join(
-                    latest_sdag_subquery,
-                    (cls.dag_id == latest_sdag_subquery.c.dag_id)
-                    & (cls.created_at == latest_sdag_subquery.c.max_created),
-                )
-                .join(cls.dag_model)
-                .where(~DagModel.is_stale)
+        query = session.execute(
+            select(cls.dag_id, data_col_to_select)
+            .join(
+                latest_sdag_subquery,
+                (cls.dag_id == latest_sdag_subquery.c.dag_id)
+                & (cls.created_at == latest_sdag_subquery.c.max_created),
             )
-            iterator = [(dag_id, json.loads(deps_data) if deps_data else []) for dag_id, deps_data in query]
-        else:
-            iterator = session.execute(
-                select(
-                    cls.dag_id,
-                    func.json_extract_path(cls._data, "dag", "dag_dependencies"),
-                )
-                .join(
-                    latest_sdag_subquery,
-                    (cls.dag_id == latest_sdag_subquery.c.dag_id)
-                    & (cls.created_at == latest_sdag_subquery.c.max_created),
-                )
-                .join(cls.dag_model)
-                .where(~DagModel.is_stale)
-            ).all()
+            .join(cls.dag_model)
+            .where(~DagModel.is_stale)
+        )
+        iterator = (
+            [(dag_id, load_json(deps_data)) for dag_id, deps_data in query]
+            if load_json is not None
+            else query.all()
+        )
 
         resolver = _DagDependenciesResolver(dag_id_dependencies=iterator, session=session)
         dag_depdendencies_by_dag = resolver.resolve()
-
         return dag_depdendencies_by_dag
 
     @staticmethod

airflow-core/tests/unit/models/test_serialized_dag.py

@@ -35,6 +35,7 @@
 from airflow.providers.standard.operators.empty import EmptyOperator
 from airflow.providers.standard.operators.python import PythonOperator
 from airflow.sdk import DAG, Asset, task as task_decorator
+from airflow.serialization.dag_dependency import DagDependency
 from airflow.serialization.serialized_objects import LazyDeserializedDAG, SerializedDAG
 from airflow.settings import json
 from airflow.utils.hashlib_wrapper import md5
@@ -386,6 +387,28 @@ def test_get_dependencies(self, session):
         dependencies = SDM.get_dag_dependencies(session=session)
         assert dag_id not in dependencies
 
+    def test_get_dependencies_with_asset_ref(self, dag_maker, session):
+        with dag_maker(
+            dag_id="test_get_dependencies_with_asset_ref_example",
+            start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
+            schedule=[Asset.ref(uri="test://asset1")],
+        ) as dag:
+            BashOperator(task_id="any", bash_command="sleep 5")
+        dag.sync_to_db()
+        SDM.write_dag(dag, bundle_name="testing")
+        dependencies = SDM.get_dag_dependencies(session=session)
+        assert dependencies == {
+            "test_get_dependencies_with_asset_ref_example": [
+                DagDependency(
+                    source="asset-uri-ref",
+                    target="test_get_dependencies_with_asset_ref_example",
+                    label="test://asset1",
+                    dependency_type="asset-uri-ref",
+                    dependency_id="test://asset1",
+                )
+            ]
+        }
+
     @pytest.mark.parametrize("min_update_interval", [0, 10])
     @mock.patch.object(DagVersion, "get_latest_version")
     def test_min_update_interval_is_respected(