This is an automated email from the ASF dual-hosted git repository.
shahar1 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 162f629a01e Fix task state lookup for mapped tasks in Google triggers
(#76895)
162f629a01e is described below
commit 162f629a01ec4d61f3d4fd2d8b314048bb6d698b
Author: Peter Nguyen <[email protected]>
AuthorDate: Thu Nov 12 21:57:39 2026 -0800
Fix task state lookup for mapped tasks in Google triggers (#76895)
The execution API keys a mapped task instance's state as
"<task_id>_<map_index>", but the Google triggers indexed the response by
plain task_id. For a mapped task, safe_to_cancel() therefore always raised
"TaskInstance ... is not found" when the trigger was cancelled, failing the
task instead of leaving it deferred.
Build the key from map_index in get_task_state() for DataprocSubmitTrigger,
DataprocSubmitJobDirectTrigger, DataprocClusterTrigger and
BigQueryInsertJobTrigger.
---
.../providers/google/cloud/triggers/bigquery.py | 4 +-
.../providers/google/cloud/triggers/dataproc.py | 12 +++--
.../unit/google/cloud/triggers/test_bigquery.py | 24 ++++++++-
.../unit/google/cloud/triggers/test_dataproc.py | 60 ++++++++++++++++++++++
4 files changed, 95 insertions(+), 5 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
b/providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
index 74fc923eafa..8f7c961147f 100644
--- a/providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
+++ b/providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
@@ -151,8 +151,10 @@ class BigQueryInsertJobTrigger(BaseTrigger):
run_ids=[ti.run_id],
map_index=ti.map_index,
)
+ # The execution API keys a mapped task instance as
``<task_id>_<map_index>``.
+ task_key = ti.task_id if ti.map_index < 0 else
f"{ti.task_id}_{ti.map_index}"
try:
- task_state = task_states_response[ti.run_id][ti.task_id]
+ task_state = task_states_response[ti.run_id][task_key]
except Exception:
msg = (
f"TaskInstance with dag_id: {ti.dag_id}, task_id:
{ti.task_id}, "
diff --git
a/providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py
b/providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py
index cd877516028..6b9b1c35058 100644
--- a/providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py
+++ b/providers/google/src/airflow/providers/google/cloud/triggers/dataproc.py
@@ -178,8 +178,10 @@ class DataprocSubmitTrigger(DataprocBaseTrigger):
run_ids=[ti.run_id],
map_index=ti.map_index,
)
+ # The execution API keys a mapped task instance as
``<task_id>_<map_index>``.
+ task_key = ti.task_id if ti.map_index < 0 else
f"{ti.task_id}_{ti.map_index}"
try:
- task_state = task_states_response[ti.run_id][ti.task_id]
+ task_state = task_states_response[ti.run_id][task_key]
except Exception:
msg = (
f"TaskInstance with dag_id: {ti.dag_id}, task_id:
{ti.task_id}, "
@@ -386,8 +388,10 @@ class DataprocSubmitJobDirectTrigger(DataprocBaseTrigger):
run_ids=[ti.run_id],
map_index=ti.map_index,
)
+ # The execution API keys a mapped task instance as
``<task_id>_<map_index>``.
+ task_key = ti.task_id if ti.map_index < 0 else
f"{ti.task_id}_{ti.map_index}"
try:
- task_state = task_states_response[ti.run_id][ti.task_id]
+ task_state = task_states_response[ti.run_id][task_key]
except Exception:
msg = (
f"TaskInstance with dag_id: {ti.dag_id}, task_id:
{ti.task_id}, "
@@ -546,8 +550,10 @@ class DataprocClusterTrigger(DataprocBaseTrigger):
run_ids=[ti.run_id],
map_index=ti.map_index,
)
+ # The execution API keys a mapped task instance as
``<task_id>_<map_index>``.
+ task_key = ti.task_id if ti.map_index < 0 else
f"{ti.task_id}_{ti.map_index}"
try:
- task_state = task_states_response[ti.run_id][ti.task_id]
+ task_state = task_states_response[ti.run_id][task_key]
except Exception:
msg = (
f"TaskInstance with dag_id: {ti.dag_id}, task_id:
{ti.task_id}, "
diff --git a/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
b/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
index a60a7b5daf8..f8de032ec03 100644
--- a/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
+++ b/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
@@ -19,6 +19,7 @@ from __future__ import annotations
import asyncio
import contextlib
import logging
+from types import SimpleNamespace
from typing import Any
from unittest import mock
from unittest.mock import AsyncMock
@@ -42,8 +43,9 @@ from airflow.providers.google.cloud.triggers.bigquery import (
BigQueryValueCheckTrigger,
)
from airflow.triggers.base import TriggerEvent
+from airflow.utils.state import TaskInstanceState
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_3_PLUS
TEST_CONN_ID = "bq_default"
TEST_JOB_ID = "1234"
@@ -163,6 +165,26 @@ def value_check_trigger():
class TestBigQueryInsertJobTrigger:
+ @pytest.mark.asyncio
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_0_PLUS or AIRFLOW_V_3_3_PLUS,
+ reason="The task state is read through the task SDK on Airflow 3.0 to
3.2 only",
+ )
+ @pytest.mark.parametrize(
+ ("map_index", "state_key"), [(-1, "test_task"), (3, "test_task_3")],
ids=["unmapped", "mapped"]
+ )
+
@mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_task_states",
autospec=True)
+ async def test_get_task_state_uses_map_index_key(
+ self, mock_get_task_states, insert_job_trigger, map_index, state_key
+ ):
+ """The execution API keys a mapped task instance as
``<task_id>_<map_index>``."""
+ insert_job_trigger.task_instance = SimpleNamespace(
+ dag_id="test_dag", task_id="test_task", run_id="test_run",
map_index=map_index
+ )
+ mock_get_task_states.return_value = {"test_run": {state_key:
TaskInstanceState.DEFERRED}}
+
+ assert await insert_job_trigger.get_task_state() ==
TaskInstanceState.DEFERRED
+
def test_serialization(self, insert_job_trigger):
"""
Asserts that the BigQueryInsertJobTrigger correctly serializes its
arguments and classpath.
diff --git a/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py
b/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py
index 946e455d2f6..7af9ae393c4 100644
--- a/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py
+++ b/providers/google/tests/unit/google/cloud/triggers/test_dataproc.py
@@ -39,6 +39,7 @@ from airflow.providers.google.cloud.triggers.dataproc import (
)
from airflow.providers.google.cloud.utils.dataproc import DataprocOperationType
from airflow.triggers.base import TriggerEvent
+from airflow.utils.state import TaskInstanceState
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_3_PLUS
@@ -449,6 +450,25 @@ class TestDataprocClusterTrigger:
f"run_id: test_run and map_index: {map_index} is not found"
)
+ @pytest.mark.asyncio
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_0_PLUS, reason="The task state is read through the
task SDK on Airflow 3 only"
+ )
+ @pytest.mark.parametrize(
+ ("map_index", "state_key"), [(-1, "test_task"), (3, "test_task_3")],
ids=["unmapped", "mapped"]
+ )
+
@mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_task_states",
autospec=True)
+ async def test_get_task_state_uses_map_index_key(
+ self, mock_get_task_states, cluster_trigger, map_index, state_key
+ ):
+ """The execution API keys a mapped task instance as
``<task_id>_<map_index>``."""
+ cluster_trigger.task_instance = SimpleNamespace(
+ dag_id="test_dag", task_id="test_task", run_id="test_run",
map_index=map_index
+ )
+ mock_get_task_states.return_value = {"test_run": {state_key:
TaskInstanceState.DEFERRED}}
+
+ assert await cluster_trigger.get_task_state() ==
TaskInstanceState.DEFERRED
+
@pytest.mark.asyncio
@mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocClusterTrigger.get_async_hook")
@mock.patch("airflow.providers.google.cloud.triggers.dataproc.DataprocClusterTrigger.get_sync_hook")
@@ -917,6 +937,26 @@ class TestDataprocSubmitTrigger:
await async_gen.aclose()
+ @pytest.mark.asyncio
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_0_PLUS or AIRFLOW_V_3_3_PLUS,
+ reason="The task state is read through the task SDK on Airflow 3.0 to
3.2 only",
+ )
+ @pytest.mark.parametrize(
+ ("map_index", "state_key"), [(-1, "test_task"), (3, "test_task_3")],
ids=["unmapped", "mapped"]
+ )
+
@mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_task_states",
autospec=True)
+ async def test_get_task_state_uses_map_index_key(
+ self, mock_get_task_states, submit_trigger, map_index, state_key
+ ):
+ """The execution API keys a mapped task instance as
``<task_id>_<map_index>``."""
+ submit_trigger.task_instance = SimpleNamespace(
+ dag_id="test_dag", task_id="test_task", run_id="test_run",
map_index=map_index
+ )
+ mock_get_task_states.return_value = {"test_run": {state_key:
TaskInstanceState.DEFERRED}}
+
+ assert await submit_trigger.get_task_state() ==
TaskInstanceState.DEFERRED
+
class TestDataprocSubmitJobDirectTrigger:
def test_serialization(self, submit_job_direct_trigger):
@@ -1299,3 +1339,23 @@ class TestDataprocSubmitJobDirectTrigger:
await async_gen.asend(None)
mock_get_sync_hook.return_value.cancel_job.assert_called_once()
+
+ @pytest.mark.asyncio
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_0_PLUS or AIRFLOW_V_3_3_PLUS,
+ reason="The task state is read through the task SDK on Airflow 3.0 to
3.2 only",
+ )
+ @pytest.mark.parametrize(
+ ("map_index", "state_key"), [(-1, "test_task"), (3, "test_task_3")],
ids=["unmapped", "mapped"]
+ )
+
@mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_task_states",
autospec=True)
+ async def test_get_task_state_uses_map_index_key(
+ self, mock_get_task_states, submit_job_direct_trigger, map_index,
state_key
+ ):
+ """The execution API keys a mapped task instance as
``<task_id>_<map_index>``."""
+ submit_job_direct_trigger.task_instance = SimpleNamespace(
+ dag_id="test_dag", task_id="test_task", run_id="test_run",
map_index=map_index
+ )
+ mock_get_task_states.return_value = {"test_run": {state_key:
TaskInstanceState.DEFERRED}}
+
+ assert await submit_job_direct_trigger.get_task_state() ==
TaskInstanceState.DEFERRED
