dabla commented on code in PR #55068:
URL: https://github.com/apache/airflow/pull/55068#discussion_r2498982693


##########
airflow-core/src/airflow/models/dagbag.py:
##########
@@ -46,24 +46,29 @@ class DBDagBag:
     """
 
     def __init__(self, load_op_links: bool = True) -> None:
-        self._dags: dict[str, SerializedDAG] = {}  # dag_version_id to dag
+        self._dags: dict[str, SerializedDagModel] = {}  # dag_version_id to dag
         self.load_op_links = load_op_links
 
     def _read_dag(self, serdag: SerializedDagModel) -> SerializedDAG | None:
         serdag.load_op_links = self.load_op_links
         if dag := serdag.dag:
-            self._dags[serdag.dag_version_id] = dag
+            self._dags[serdag.dag_version_id] = serdag
         return dag
 
-    def _get_dag(self, version_id: str, session: Session) -> SerializedDAG | 
None:
-        if dag := self._dags.get(version_id):
-            return dag
-        dag_version = session.get(DagVersion, version_id, 
options=[joinedload(DagVersion.serialized_dag)])
-        if not dag_version:
-            return None
-        if not (serdag := dag_version.serialized_dag):
-            return None
-        return self._read_dag(serdag)
+    def get_dag_model(self, version_id: str, session: Session) -> 
SerializedDagModel | None:
+        if not (serdag := self._dags.get(version_id)):
+            dag_version = session.get(DagVersion, version_id, 
options=[joinedload(DagVersion.serialized_dag)])
+            if not dag_version:
+                return None
+            if not (serdag := dag_version.serialized_dag):
+                return None
+        self._read_dag(serdag)
+        return serdag
+
+    def get_dag(self, version_id: str, session: Session) -> SerializedDAG | 
None:

Review Comment:
   We could but it was already called like that originally so I didn't want to 
change it except for the fact that a changed it from protected to public, as 
the protected didn't make sense here imho as it's a responsibility of the 
DBDagBag.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to