This is an automated email from the ASF dual-hosted git repository.
vincbeck 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 1d66775efcb Limit Glue retry task UUID scan (#71546)
1d66775efcb is described below
commit 1d66775efcbc2b31ecb476bef3198ca51ef7b4e3
Author: Vincent Hsiao <[email protected]>
AuthorDate: Mon Oct 19 22:12:00 2026 +0800
Limit Glue retry task UUID scan (#71546)
---
providers/amazon/docs/operators/glue.rst | 4 +-
.../airflow/providers/amazon/aws/operators/glue.py | 27 ++++--
.../tests/unit/amazon/aws/operators/test_glue.py | 107 +++++++++++++++++++++
3 files changed, 128 insertions(+), 10 deletions(-)
diff --git a/providers/amazon/docs/operators/glue.rst
b/providers/amazon/docs/operators/glue.rst
index eac6e270a1d..6a4459b4898 100644
--- a/providers/amazon/docs/operators/glue.rst
+++ b/providers/amazon/docs/operators/glue.rst
@@ -227,7 +227,9 @@ resubmit, and with ``concurrent_run_limit=1`` that fails
with ``ConcurrentRunsEx
against the run it can't see. To avoid that, ``durable=True`` (or, below
Airflow 3.3,
``resume_glue_job_on_retry=True``) tags the job's arguments with this task
instance's identity on
every attempt and, on retry, scans the job's run history for that tag before
submitting -- the
-same mechanism used as a fallback on the synchronous path.
+same mechanism used as a fallback on the synchronous path. This scan requires
``glue:GetJobRuns``
+in the task's IAM policy, in addition to the usual ``glue:StartJobRun`` and
``glue:GetJobRun``
+permissions.
``durable`` supersedes the deprecated ``resume_glue_job_on_retry`` parameter
on Airflow 3.3+,
where passing ``resume_glue_job_on_retry`` still works and maps its value onto
``durable``. Below
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py
index c8b1d617e0f..b4dc5903415 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py
@@ -23,7 +23,7 @@ import warnings
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, cast
-from botocore.exceptions import ClientError
+from botocore.exceptions import BotoCoreError, ClientError
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.exceptions import GlueJobRunStoppedError
@@ -145,6 +145,8 @@ class GlueJobOperator(ResumableJobMixin,
AwsBaseOperator[GlueJobHook]):
operator_extra_links = (GlueJobRunDetailsLink(),)
TASK_UUID_ARG = "--airflow_task_uuid"
+ TASK_UUID_SCAN_PAGE_SIZE = 200
+ TASK_UUID_SCAN_MAX_PAGES = 10
external_id_key = "glue_job_run_id"
def __init__(
@@ -294,11 +296,9 @@ class GlueJobOperator(ResumableJobMixin,
AwsBaseOperator[GlueJobHook]):
return script_args, task_uuid
def _find_job_run_id_by_task_uuid(self, task_uuid: str) -> tuple[str, str]
| None:
- # Unbounded walk with no page cap; a no-match run is the common retry
shape and pays the
- # full scan. Tracked at https://github.com/apache/airflow/issues/71489.
next_token: str | None = None
- while True:
- request = {"JobName": self.job_name, "MaxResults": 50}
+ for _ in range(self.TASK_UUID_SCAN_MAX_PAGES):
+ request = {"JobName": self.job_name, "MaxResults":
self.TASK_UUID_SCAN_PAGE_SIZE}
if next_token:
request["NextToken"] = next_token
response = self.hook.conn.get_job_runs(**request)
@@ -312,6 +312,11 @@ class GlueJobOperator(ResumableJobMixin,
AwsBaseOperator[GlueJobHook]):
next_token = response.get("NextToken")
if not next_token:
return None
+ self.log.warning(
+ "Stopped scanning Glue job runs for task UUID after %s pages
without a match",
+ self.TASK_UUID_SCAN_MAX_PAGES,
+ )
+ return None
def execute(self, context: Context) -> str | None:
"""
@@ -414,8 +419,10 @@ class GlueJobOperator(ResumableJobMixin,
AwsBaseOperator[GlueJobHook]):
self.log.info("Previous Glue job_run_id: %s, state: %s",
previous_job_run_id, state)
if self.is_job_active(state):
return previous_job_run_id
- except Exception:
- self.log.warning("Failed to get previous Glue job run state",
exc_info=True)
+ except (ClientError, BotoCoreError):
+ self.log.exception(
+ "Failed to get previous Glue job run state; submitting a
fresh Glue job run"
+ )
else:
try:
existing = self._find_job_run_id_by_task_uuid(task_uuid)
@@ -428,8 +435,10 @@ class GlueJobOperator(ResumableJobMixin,
AwsBaseOperator[GlueJobHook]):
)
if self.is_job_active(existing_job_run_state):
return existing_job_run_id
- except Exception:
- self.log.warning("Failed to find previous Glue job run by task
UUID", exc_info=True)
+ except ClientError:
+ self.log.exception(
+ "Failed to find previous Glue job run by task UUID;
submitting a fresh Glue job run"
+ )
return None
def _has_stored_external_id(self, context: Context) -> bool:
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py
index 4a1852b919c..a38f901edac 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py
@@ -275,6 +275,113 @@ class TestGlueJobOperator:
assert defer.value.trigger.run_id == JOB_RUN_ID
+ def test_task_uuid_scan_paginates_until_match(self):
+ glue = GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME)
+ task_uuid = "test-task-uuid"
+ glue.hook.conn = mock.MagicMock()
+ glue.hook.conn.get_job_runs.side_effect = [
+ {"JobRuns": [], "NextToken": "next-page"},
+ {
+ "JobRuns": [
+ {
+ "Id": JOB_RUN_ID,
+ "Arguments": {GlueJobOperator.TASK_UUID_ARG:
task_uuid},
+ "JobRunState": "RUNNING",
+ }
+ ]
+ },
+ ]
+
+ assert glue._find_job_run_id_by_task_uuid(task_uuid) == (JOB_RUN_ID,
"RUNNING")
+ assert glue.hook.conn.get_job_runs.call_args_list == [
+ mock.call(JobName=JOB_NAME,
MaxResults=GlueJobOperator.TASK_UUID_SCAN_PAGE_SIZE),
+ mock.call(
+ JobName=JOB_NAME,
+ MaxResults=GlueJobOperator.TASK_UUID_SCAN_PAGE_SIZE,
+ NextToken="next-page",
+ ),
+ ]
+
+ def test_task_uuid_scan_stops_after_page_limit(self):
+ glue = GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME)
+ glue.hook.conn = mock.MagicMock()
+ glue.hook.conn.get_job_runs.return_value = {"JobRuns": [],
"NextToken": "next-page"}
+
+ with mock.patch.object(glue.log, "warning") as mock_log_warning:
+ assert glue._find_job_run_id_by_task_uuid("missing-task-uuid") is
None
+
+ assert glue.hook.conn.get_job_runs.call_count ==
GlueJobOperator.TASK_UUID_SCAN_MAX_PAGES
+ mock_log_warning.assert_called_once_with(
+ "Stopped scanning Glue job runs for task UUID after %s pages
without a match",
+ GlueJobOperator.TASK_UUID_SCAN_MAX_PAGES,
+ )
+
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_deferrable_retry_submits_fresh_when_task_uuid_scan_client_error(
+ self, mock_get_conn, mock_initialize_job
+ ):
+ with pytest.warns(
+ AirflowProviderDeprecationWarning,
match=f"^{re.escape(EXPECTED_DEPRECATION_MESSAGE)}$"
+ ):
+ glue = GlueJobOperator(
+ task_id=TASK_ID,
+ job_name=JOB_NAME,
+ script_location="s3://folder/file",
+ deferrable=True,
+ resume_glue_job_on_retry=True,
+ )
+ mock_ti = mock.MagicMock()
+ mock_ti.dag_id = "test_dag_id"
+ mock_ti.task_id = TASK_ID
+ mock_ti.run_id = "manual__2024-01-01T00:00:00+00:00"
+ mock_ti.map_index = -1
+ mock_ti.try_number = 2
+ mock_ti.xcom_pull.return_value = None
+ mock_initialize_job.return_value = {"JobRunId": JOB_RUN_ID}
+ glue.hook.conn = mock.MagicMock()
+ glue.hook.conn.get_job_runs.side_effect = ClientError(
+ {"Error": {"Code": "ThrottlingException", "Message": "slow
down"}}, "GetJobRuns"
+ )
+
+ with mock.patch.object(glue.log, "exception") as mock_log_exception:
+ with pytest.raises(TaskDeferred) as defer:
+ glue.execute({"ti": mock_ti})
+
+ assert defer.value.trigger.run_id == JOB_RUN_ID
+ mock_initialize_job.assert_called_once()
+ mock_log_exception.assert_called_once_with(
+ "Failed to find previous Glue job run by task UUID; submitting a
fresh Glue job run"
+ )
+
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_deferrable_retry_does_not_swallow_task_uuid_scan_bug(self,
mock_get_conn, mock_initialize_job):
+ with pytest.warns(
+ AirflowProviderDeprecationWarning,
match=f"^{re.escape(EXPECTED_DEPRECATION_MESSAGE)}$"
+ ):
+ glue = GlueJobOperator(
+ task_id=TASK_ID,
+ job_name=JOB_NAME,
+ script_location="s3://folder/file",
+ deferrable=True,
+ resume_glue_job_on_retry=True,
+ )
+ mock_ti = mock.MagicMock()
+ mock_ti.dag_id = "test_dag_id"
+ mock_ti.task_id = TASK_ID
+ mock_ti.run_id = "manual__2024-01-01T00:00:00+00:00"
+ mock_ti.map_index = -1
+ mock_ti.try_number = 2
+ mock_ti.xcom_pull.return_value = None
+ glue.hook.conn = mock.MagicMock()
+ glue.hook.conn.get_job_runs.side_effect = RuntimeError("scan bug")
+
+ with pytest.raises(RuntimeError, match="scan bug"):
+ glue.execute({"ti": mock_ti})
+
+ mock_initialize_job.assert_not_called()
+
@pytest.mark.skipif(
not AIRFLOW_V_3_3_PLUS,
reason="task_state_store only exists as an execute() context key on
Airflow 3.3+",
