jedcunningham commented on code in PR #46626:
URL: https://github.com/apache/airflow/pull/46626#discussion_r1949886001
##########
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:
We should probably union this instead so it's a single query vs 2.
##########
airflow/models/dag.py:
##########
@@ -287,7 +286,11 @@ def _create_orm_dagrun(
run.dag = dag
# create the associated task instances
# state is None at the moment of creation
- run.verify_integrity(session=session)
+ if not dag_version:
+ dag_version = DagVersion.get_latest_version(dag.dag_id,
session=session)
+ if not dag_version:
+ raise AirflowException(f"Could not find a version for DAG
{dag.dag_id}")
Review Comment:
Is this possible?
##########
tests/models/test_dagrun.py:
##########
@@ -135,11 +135,11 @@ def create_dag_run(
return dag_run
- def test_clear_task_instances_for_backfill_running_dagrun(self, session):
+ def test_clear_task_instances_for_backfill_running_dagrun(self, dag_maker,
session):
now = timezone.utcnow()
state = DagRunState.RUNNING
dag_id = "test_clear_task_instances_for_backfill_running_dagrun"
- with DAG(dag_id=dag_id) as dag:
+ with dag_maker(dag_id=dag_id) as dag:
EmptyOperator(task_id="backfill_task_0")
self.create_dag_run(dag, logical_date=now, is_backfill=True,
state=state, session=session)
Review Comment:
Can probably use the dagmaker create dagrun tooling too.
--
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]
