Fix failures on main related to DagRun validation

Follow-up of #45834
apache · Jan 22, 2025 · d2243a7 · d2243a7
1 parent a9bad24
commit d2243a7
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 2 deletions.
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
@@ -2991,6 +2991,9 @@ def run(
         raise_on_defer: bool = False,
     ) -> None:
         """Run TaskInstance."""
+        if self.dag_run is not None:
+            self.dag_run.set_state(DagRunState.RUNNING)
+
         res = self.check_and_change_state_before_execution(
             verbose=verbose,
             ignore_all_deps=ignore_all_deps,

diff --git a/providers/tests/apache/kylin/operators/test_kylin_cube.py b/providers/tests/apache/kylin/operators/test_kylin_cube.py
@@ -177,13 +177,15 @@ def test_render_template(self, session):
                 run_id="kylin_test",
                 logical_date=DEFAULT_DATE,
                 run_type=DagRunType.MANUAL,
+                state="running",
             )
         else:
             ti.dag_run = DagRun(
                 dag_id=self.dag.dag_id,
                 run_id="kylin_test",
                 execution_date=DEFAULT_DATE,
                 run_type=DagRunType.MANUAL,
+                state="running",
             )
         session.add(ti)
         session.commit()

diff --git a/providers/tests/apache/spark/operators/test_spark_submit.py b/providers/tests/apache/spark/operators/test_spark_submit.py
@@ -201,14 +201,17 @@ def test_render_template(self, session):
                 run_id="spark_test",
                 logical_date=DEFAULT_DATE,
                 run_type=DagRunType.MANUAL,
+                state="running",
             )
         else:
             ti.dag_run = DagRun(
                 dag_id=self.dag.dag_id,
                 run_id="spark_test",
                 execution_date=DEFAULT_DATE,
                 run_type=DagRunType.MANUAL,
+                state="running",
             )
+
         session.add(ti)
         session.commit()
         # When

diff --git a/providers/tests/openlineage/plugins/test_listener.py b/providers/tests/openlineage/plugins/test_listener.py
@@ -101,7 +101,7 @@ def test_listener_does_not_change_task_instance(render_mock, xcom_push_mock):
         run_id=run_id,
         data_interval=(date, date),
         run_type=types.DagRunType.MANUAL,
-        state=DagRunState.QUEUED,
+        state=DagRunState.RUNNING,
         **dagrun_kwargs,
     )
     ti = TaskInstance(t, run_id=run_id)
@@ -187,7 +187,7 @@ def sample_callable(**kwargs):
         run_id=run_id,
         data_interval=(date, date),
         run_type=types.DagRunType.MANUAL,
-        state=DagRunState.QUEUED,
+        state=DagRunState.RUNNING,
         **dagrun_kwargs,
     )
     task_instance = TaskInstance(t, run_id=run_id)

diff --git a/providers/tests/redis/log/test_redis_task_handler.py b/providers/tests/redis/log/test_redis_task_handler.py
@@ -46,6 +46,7 @@ def ti(self):
         else:
             dag_run = DagRun(dag_id=dag.dag_id, execution_date=date, run_id="test", run_type="scheduled")
 
+        dag_run.set_state(State.RUNNING)
         with create_session() as session:
             session.add(dag_run)
             session.commit()