You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2025/04/11 13:14:15 UTC
(airflow) branch main updated: Fix GetDagRunState and GetTICount response conversion (#49095)
This is an automated email from the ASF dual-hosted git repository.
kaxilnaik pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new c2504f1afbd Fix GetDagRunState and GetTICount response conversion (#49095)
c2504f1afbd is described below
commit c2504f1afbd8b22e6269d90980e631244d26534c
Author: GPK <go...@gmail.com>
AuthorDate: Fri Apr 11 14:14:09 2025 +0100
Fix GetDagRunState and GetTICount response conversion (#49095)
---
airflow-core/src/airflow/jobs/triggerer_job_runner.py | 5 ++---
airflow-core/tests/unit/jobs/test_triggerer_job.py | 12 ++++++++----
task-sdk/src/airflow/sdk/execution_time/supervisor.py | 4 +++-
3 files changed, 13 insertions(+), 8 deletions(-)
diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
index 32e691eecb4..bef75cf5152 100644
--- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
+++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
@@ -424,10 +424,10 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
resp = dr_count
elif isinstance(msg, GetDagRunState):
dr_resp = self.client.dag_runs.get_state(msg.dag_id, msg.run_id)
- resp = dr_resp
+ resp = DagRunStateResult.from_api_response(dr_resp)
elif isinstance(msg, GetTICount):
- ti_count = self.client.task_instances.get_count(
+ resp = self.client.task_instances.get_count(
dag_id=msg.dag_id,
task_ids=msg.task_ids,
task_group_id=msg.task_group_id,
@@ -435,7 +435,6 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
run_ids=msg.run_ids,
states=msg.states,
)
- resp = ti_count.model_dump_json().encode()
else:
raise ValueError(f"Unknown message type {type(msg)}")
diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py
index 814aeaf7ba5..87302a4e6f7 100644
--- a/airflow-core/tests/unit/jobs/test_triggerer_job.py
+++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py
@@ -719,7 +719,11 @@ class CustomTriggerDagRun(BaseTrigger):
states=self.states,
logical_dates=self.logical_dates,
)
- yield TriggerEvent({"count": dag_run_states_count})
+ dag_run_state = await sync_to_async(RuntimeTaskInstance.get_dagrun_state)(
+ dag_id=self.trigger_dag_id,
+ run_id=self.run_ids[0],
+ )
+ yield TriggerEvent({"count": dag_run_states_count, "dag_run_state": dag_run_state})
@pytest.mark.xfail(
@@ -729,10 +733,10 @@ class CustomTriggerDagRun(BaseTrigger):
@pytest.mark.asyncio
@pytest.mark.flaky(reruns=2, reruns_delay=10)
@pytest.mark.execution_timeout(30)
-async def test_trigger_can_fetch_trigger_dag_run_count_in_deferrable(session, dag_maker):
+async def test_trigger_can_fetch_trigger_dag_run_count_and_state_in_deferrable(session, dag_maker):
"""Checks that the trigger will successfully fetch the count of trigger DAG runs."""
# Create the test DAG and task
- with dag_maker(dag_id="trigger_can_fetch_trigger_dag_run_count_in_deferrable", session=session):
+ with dag_maker(dag_id="trigger_can_fetch_trigger_dag_run_count_and_state_in_deferrable", session=session):
EmptyOperator(task_id="dummy1")
dr = dag_maker.create_dagrun()
task_instance = dr.task_instances[0]
@@ -766,7 +770,7 @@ async def test_trigger_can_fetch_trigger_dag_run_count_in_deferrable(session, da
task_instance.refresh_from_db()
assert task_instance.state == TaskInstanceState.SCHEDULED
assert task_instance.next_method != "__fail__"
- assert task_instance.next_kwargs == {"event": {"count": 1}}
+ assert task_instance.next_kwargs == {"event": {"count": 1, "dag_run_state": "running"}}
class CustomTriggerWorkflowStateTrigger(BaseTrigger):
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index 951475d3c76..7ade78b4ec1 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -64,6 +64,7 @@ from airflow.sdk.execution_time.comms import (
AssetEventsResult,
AssetResult,
ConnectionResult,
+ DagRunStateResult,
DeferTask,
DeleteXCom,
ErrorResponse,
@@ -1021,7 +1022,8 @@ class ActivitySubprocess(WatchedSubprocess):
msg.reset_dag_run,
)
elif isinstance(msg, GetDagRunState):
- resp = self.client.dag_runs.get_state(msg.dag_id, msg.run_id)
+ dr_resp = self.client.dag_runs.get_state(msg.dag_id, msg.run_id)
+ resp = DagRunStateResult.from_api_response(dr_resp)
elif isinstance(msg, GetTaskRescheduleStartDate):
resp = self.client.task_instances.get_reschedule_start_date(msg.ti_id, msg.try_number)
elif isinstance(msg, GetTICount):