From 797ca08fc5573867385633f99834e040eccd29c6 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 1 Sep 2026 22:01:41 +0000 Subject: [PATCH 1/2] Bound and harden the Glue task-UUID scan fallback The no-match retry path walked the job's entire GetJobRuns history before giving up, and AWS lookup failures were swallowed as warnings behind a broad except Exception. Anyone hitting that fallback then submitted a fresh run into a still-active job with no error-level signal. --- providers/amazon/docs/operators/glue.rst | 7 + .../providers/amazon/aws/operators/glue.py | 44 ++++- .../unit/amazon/aws/operators/test_glue.py | 182 +++++++++++++++++- 3 files changed, 222 insertions(+), 11 deletions(-) diff --git a/providers/amazon/docs/operators/glue.rst b/providers/amazon/docs/operators/glue.rst index eac6e270a1d91..babf12bd24716 100644 --- a/providers/amazon/docs/operators/glue.rst +++ b/providers/amazon/docs/operators/glue.rst @@ -189,6 +189,13 @@ with this task instance's identity, and reconnects if it finds one that is still XCom check only ever succeeds on Airflow 2.x -- every Airflow 3 release clears task XComs before each non-deferral attempt, so on Airflow 3.0-3.2 every retry pays the full scan. +.. note:: + The task-UUID scan fallback requires ``glue:GetJobRuns`` in the task's IAM policy, in addition + to ``glue:StartJobRun`` and ``glue:GetJobRun``. Anyone relying on this fallback -- via + ``resume_glue_job_on_retry=True`` below Airflow 3.3, or when task state store has nothing for + this task on Airflow 3.3+ -- needs that extra action. Without it the scan fails, is logged as + an error, and the operator submits a fresh run instead of reconnecting. + That older mechanism only activates below 3.3 when ``resume_glue_job_on_retry=True`` is set explicitly -- ``durable=True`` does not turn it on there. Upgrading the provider alone, with no DAG change, does not turn it on either: below 3.3, behavior is unchanged from before this feature 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 c488965ddde01..4398bb94fbea7 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py @@ -21,6 +21,7 @@ import urllib.parse import warnings from collections.abc import Sequence +from datetime import datetime from typing import TYPE_CHECKING, Any, cast from botocore.exceptions import ClientError @@ -180,6 +181,9 @@ class GlueJobOperator(ResumableJobMixin, AwsBaseOperator[GlueJobHook]): operator_extra_links = (GlueJobRunDetailsLink(),) TASK_UUID_ARG = "--airflow_task_uuid" + # GetJobRuns MaxResults maxes out at 200; the page cap is a backstop when no age cutoff applies. + TASK_UUID_SCAN_PAGE_SIZE = 200 + TASK_UUID_SCAN_MAX_PAGES = 10 external_id_key = "glue_job_run_id" def __init__( @@ -328,12 +332,12 @@ def _prepare_script_args_with_task_uuid( script_args[self.TASK_UUID_ARG] = task_uuid 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. + def _find_job_run_id_by_task_uuid( + self, task_uuid: str, *, cutoff: datetime | None = None + ) -> tuple[str, str] | None: 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) @@ -344,9 +348,19 @@ def _find_job_run_id_by_task_uuid(self, task_uuid: str) -> tuple[str, str] | Non job_run_state = job_run.get("JobRunState") if job_run_id and job_run_state: return job_run_id, job_run_state + started_on = job_run.get("StartedOn") + if cutoff is not None and started_on is not None and started_on < cutoff: + return None next_token = response.get("NextToken") if not next_token: return None + self.log.error( + "Glue job run UUID scan for job_name=%s exceeded %s pages without a match; " + "submitting a fresh Glue job run", + self.job_name, + self.TASK_UUID_SCAN_MAX_PAGES, + ) + return None def execute(self, context: Context) -> str | None: """ @@ -447,11 +461,17 @@ def _find_previous_job_run(self, context: Context, task_uuid: str) -> str | None 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: + self.log.exception( + "Failed to get previous Glue job run state for run_id=%s; submitting a fresh Glue job run", + previous_job_run_id, + ) else: try: - existing = self._find_job_run_id_by_task_uuid(task_uuid) + dag_run = context.get("dag_run") + start_date = getattr(dag_run, "start_date", None) if dag_run is not None else None + cutoff = start_date if isinstance(start_date, datetime) else None + existing = self._find_job_run_id_by_task_uuid(task_uuid, cutoff=cutoff) if existing: existing_job_run_id, existing_job_run_state = existing self.log.info( @@ -461,8 +481,12 @@ def _find_previous_job_run(self, context: Context, task_uuid: str) -> str | None ) 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 for job_name=%s; " + "submitting a fresh Glue job run", + self.job_name, + ) 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 6ea43e6006490..85e02cd8a54d2 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py @@ -19,7 +19,7 @@ import re import warnings from collections.abc import Generator -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING from unittest import mock @@ -696,6 +696,186 @@ def test_find_job_run_by_task_uuid_reconnects(self, mock_initialize_job, mock_ge mock_initialize_job.assert_not_called() mock_ti.xcom_push.assert_any_call(key="glue_job_run_id", value="existing_run_123") + @mock.patch.object(GlueJobHook, "get_conn") + def test_task_uuid_scan_paginates_until_match(self, mock_get_conn): + glue = GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME) + 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: "test-task-uuid"}, + "JobRunState": "RUNNING", + } + ] + }, + ] + + assert glue._find_job_run_id_by_task_uuid("test-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", + ), + ] + + @mock.patch.object(GlueJobHook, "get_conn") + def test_task_uuid_scan_stops_after_page_limit(self, mock_get_conn, caplog): + glue = GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME) + glue.TASK_UUID_SCAN_MAX_PAGES = 2 + glue.hook.conn = mock.MagicMock() + glue.hook.conn.get_job_runs.return_value = {"JobRuns": [], "NextToken": "next-page"} + + with caplog.at_level("ERROR"): + assert glue._find_job_run_id_by_task_uuid("missing-task-uuid") is None + + assert glue.hook.conn.get_job_runs.call_count == 2 + assert "pages without a match" in caplog + + @mock.patch.object(GlueJobHook, "get_conn") + def test_task_uuid_scan_stops_when_started_on_predates_cutoff(self, mock_get_conn): + cutoff = datetime(2024, 6, 1, 12, 0, tzinfo=timezone.utc) + glue = GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME) + glue.hook.conn = mock.MagicMock() + glue.hook.conn.get_job_runs.return_value = { + "JobRuns": [ + { + "Id": "older-run", + "JobRunState": "SUCCEEDED", + "StartedOn": datetime(2024, 5, 31, 12, 0, tzinfo=timezone.utc), + "Arguments": {}, + } + ], + "NextToken": "next-page", + } + + assert glue._find_job_run_id_by_task_uuid("test-task-uuid", cutoff=cutoff) is None + glue.hook.conn.get_job_runs.assert_called_once_with( + JobName=JOB_NAME, MaxResults=GlueJobOperator.TASK_UUID_SCAN_PAGE_SIZE + ) + + @mock.patch.object(GlueJobHook, "get_conn") + def test_find_previous_job_run_uses_dag_run_start_date_as_cutoff(self, mock_get_conn): + cutoff = datetime(2024, 6, 1, 12, 0, tzinfo=timezone.utc) + glue = GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME) + glue.hook.conn = mock.MagicMock() + glue.hook.conn.get_job_runs.return_value = { + "JobRuns": [ + { + "Id": "older-run", + "JobRunState": "SUCCEEDED", + "StartedOn": datetime(2024, 5, 31, 12, 0, tzinfo=timezone.utc), + "Arguments": {}, + } + ], + "NextToken": "next-page", + } + mock_ti = mock.MagicMock() + mock_ti.task_id = TASK_ID + mock_ti.xcom_pull.return_value = None + dag_run = mock.MagicMock() + dag_run.start_date = cutoff + + assert glue._find_previous_job_run({"ti": mock_ti, "dag_run": dag_run}, "test-task-uuid") is None + glue.hook.conn.get_job_runs.assert_called_once_with( + JobName=JOB_NAME, MaxResults=GlueJobOperator.TASK_UUID_SCAN_PAGE_SIZE + ) + + @mock.patch.object(GlueJobHook, "get_conn") + def test_task_uuid_scan_continues_when_started_on_is_not_before_cutoff(self, mock_get_conn): + cutoff = datetime(2024, 6, 1, 12, 0, tzinfo=timezone.utc) + glue = GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME) + glue.hook.conn = mock.MagicMock() + glue.hook.conn.get_job_runs.side_effect = [ + { + "JobRuns": [ + { + "Id": "newer-unrelated", + "JobRunState": "SUCCEEDED", + "StartedOn": cutoff, + "Arguments": {}, + } + ], + "NextToken": "next-page", + }, + { + "JobRuns": [ + { + "Id": JOB_RUN_ID, + "JobRunState": "RUNNING", + "StartedOn": cutoff, + "Arguments": {GlueJobOperator.TASK_UUID_ARG: "test-task-uuid"}, + } + ] + }, + ] + + assert glue._find_job_run_id_by_task_uuid("test-task-uuid", cutoff=cutoff) == (JOB_RUN_ID, "RUNNING") + assert glue.hook.conn.get_job_runs.call_count == 2 + + @mock.patch.object(GlueJobHook, "get_conn") + def test_find_previous_job_run_scan_client_error_returns_none(self, mock_get_conn, caplog): + glue = GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME) + glue.hook.conn = mock.MagicMock() + glue.hook.conn.get_job_runs.side_effect = ClientError( + {"Error": {"Code": "AccessDeniedException", "Message": "not authorized"}}, + "GetJobRuns", + ) + mock_ti = mock.MagicMock() + mock_ti.task_id = TASK_ID + mock_ti.xcom_pull.return_value = None + + with caplog.at_level("ERROR"): + assert glue._find_previous_job_run({"ti": mock_ti}, "test-task-uuid") is None + + assert "Failed to find previous Glue job run by task UUID" in caplog + + @mock.patch.object(GlueJobHook, "get_conn") + def test_find_previous_job_run_scan_non_client_error_propagates(self, mock_get_conn): + glue = GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME) + glue.hook.conn = mock.MagicMock() + glue.hook.conn.get_job_runs.side_effect = RuntimeError("scan bug") + mock_ti = mock.MagicMock() + mock_ti.task_id = TASK_ID + mock_ti.xcom_pull.return_value = None + + with pytest.raises(RuntimeError, match="scan bug"): + glue._find_previous_job_run({"ti": mock_ti}, "test-task-uuid") + + @mock.patch.object(GlueJobHook, "get_conn") + def test_find_previous_job_run_xcom_client_error_returns_none(self, mock_get_conn, caplog): + glue = GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME) + glue.hook.conn = mock.MagicMock() + glue.hook.conn.get_job_run.side_effect = ClientError( + {"Error": {"Code": "EntityNotFoundException", "Message": "missing"}}, + "GetJobRun", + ) + mock_ti = mock.MagicMock() + mock_ti.task_id = TASK_ID + mock_ti.xcom_pull.return_value = "previous_run_12345" + + with caplog.at_level("ERROR"): + assert glue._find_previous_job_run({"ti": mock_ti}, "unused-uuid") is None + + glue.hook.conn.get_job_runs.assert_not_called() + assert "Failed to get previous Glue job run state" in caplog + + @mock.patch.object(GlueJobHook, "get_conn") + def test_find_previous_job_run_xcom_non_client_error_propagates(self, mock_get_conn): + glue = GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME) + glue.hook.conn = mock.MagicMock() + glue.hook.conn.get_job_run.side_effect = RuntimeError("lookup bug") + mock_ti = mock.MagicMock() + mock_ti.task_id = TASK_ID + mock_ti.xcom_pull.return_value = "previous_run_12345" + + with pytest.raises(RuntimeError, match="lookup bug"): + glue._find_previous_job_run({"ti": mock_ti}, "unused-uuid") + class TestGlueJobOperatorOpenLineageInjection: """Tests for OpenLineage parent job info and transport info injection in GlueJobOperator.""" From 0e8c11c338ec58fc25c412fbce9f0362becd9975 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 1 Sep 2026 22:19:29 +0000 Subject: [PATCH 2/2] Fix Glue scan tests to use structured log assertions Structlog caplog matches the full event string, so substring checks missed the page-cap error log even when the scan stopped correctly. --- .../unit/amazon/aws/operators/test_glue.py | 24 ++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) 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 85e02cd8a54d2..8a0a59fb18ac2 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_glue.py @@ -734,7 +734,13 @@ def test_task_uuid_scan_stops_after_page_limit(self, mock_get_conn, caplog): assert glue._find_job_run_id_by_task_uuid("missing-task-uuid") is None assert glue.hook.conn.get_job_runs.call_count == 2 - assert "pages without a match" in caplog + assert { + "event": ( + f"Glue job run UUID scan for job_name={JOB_NAME} exceeded 2 pages without a match; " + "submitting a fresh Glue job run" + ), + "log_level": "error", + } in caplog @mock.patch.object(GlueJobHook, "get_conn") def test_task_uuid_scan_stops_when_started_on_predates_cutoff(self, mock_get_conn): @@ -832,7 +838,13 @@ def test_find_previous_job_run_scan_client_error_returns_none(self, mock_get_con with caplog.at_level("ERROR"): assert glue._find_previous_job_run({"ti": mock_ti}, "test-task-uuid") is None - assert "Failed to find previous Glue job run by task UUID" in caplog + assert { + "event": ( + f"Failed to find previous Glue job run by task UUID for job_name={JOB_NAME}; " + "submitting a fresh Glue job run" + ), + "log_level": "error", + } in caplog @mock.patch.object(GlueJobHook, "get_conn") def test_find_previous_job_run_scan_non_client_error_propagates(self, mock_get_conn): @@ -862,7 +874,13 @@ def test_find_previous_job_run_xcom_client_error_returns_none(self, mock_get_con assert glue._find_previous_job_run({"ti": mock_ti}, "unused-uuid") is None glue.hook.conn.get_job_runs.assert_not_called() - assert "Failed to get previous Glue job run state" in caplog + assert { + "event": ( + "Failed to get previous Glue job run state for run_id=previous_run_12345; " + "submitting a fresh Glue job run" + ), + "log_level": "error", + } in caplog @mock.patch.object(GlueJobHook, "get_conn") def test_find_previous_job_run_xcom_non_client_error_propagates(self, mock_get_conn):