jedcunningham commented on code in PR #46626:
URL: https://github.com/apache/airflow/pull/46626#discussion_r1952875855


##########
airflow/models/dagrun.py:
##########
@@ -300,6 +297,21 @@ def validate_run_id(self, key: str, run_id: str) -> str | 
None:
             f"The run_id provided '{run_id}' does not match regex pattern 
'{regex}' or '{RUN_ID_REGEX}'"
         )
 
+    @provide_session
+    def dag_versions(self, session: Session = NEW_SESSION) -> list[DagVersion]:
+        """
+        Return the DAG versions associated with the TIs of this DagRun.
+
+        :param session: database session
+        """
+        tis = session.scalars(
+            select(TI.dag_version_id).where(TI.run_id == self.run_id, 
TI.dag_id == self.dag_id).distinct()
+        ).all()
+        tih = session.scalars(
+            select(TIH.dag_version_id).where(TIH.run_id == self.run_id, 
TIH.dag_id == self.dag_id).distinct()
+        ).all()
+        return list(tis + tih)

Review Comment:
   Looking more closely, if we just add a tih relationship to dagrun, we should 
be able to joinedload them - and we don't need the whole version record, just 
the ids are enough, and those are already on ti/tih.
   
   But, instead of having a method on dagrun to do this, it might be better to 
have a helper/classmethod that knows how to do this in bulk, e.g. `def 
get_dag_version_ids(list[DagRun], dict[DagRun, list[uuids]])`. 



-- 
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