Bound and harden the Glue task-UUID scan fallback

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

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:

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,204 @@ 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 {
+            "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):
+        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 {
+            "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):
+        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 {
+            "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):
+        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."""