dabla commented on code in PR #55068:
URL: https://github.com/apache/airflow/pull/55068#discussion_r2774220089
##########
airflow-core/src/airflow/models/dagbag.py:
##########
@@ -46,24 +46,27 @@ 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
+ def _read_dag(self, serialized_dag_model: SerializedDagModel) ->
SerializedDAG | None:
+ serialized_dag_model.load_op_links = self.load_op_links
+ if dag := serialized_dag_model.dag:
+ self._dags[serialized_dag_model.dag_version_id] =
serialized_dag_model
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 (serialized_dag_model := self._dags.get(version_id)):
+ dag_version = session.get(DagVersion, version_id,
options=[joinedload(DagVersion.serialized_dag)])
+ if not dag_version or not (serialized_dag_model :=
dag_version.serialized_dag):
+ return None
+ self._read_dag(serialized_dag_model)
+ return serialized_dag_model
+
+ def get_dag(self, version_id: str, session: Session) -> SerializedDAG |
None:
+ if serialized_dag_model := self.get_dag_model(version_id=version_id,
session=session):
+ return self._read_dag(serialized_dag_model)
Review Comment:
Good catch forgot to adapt it after refactoring!
--
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]
