kaxil commented on code in PR #71211:
URL: https://github.com/apache/airflow/pull/71211#discussion_r3736165132


##########
providers/amazon/docs/operators/glue.rst:
##########
@@ -139,6 +139,70 @@ To submit a new AWS Glue job you can use 
:class:`~airflow.providers.amazon.aws.o
   The same AWS IAM role used for the crawler can be used here as well, but it 
will need
   policies to provide access to the output location for result data.
 
+Durable execution
+==================
+
+``GlueJobOperator`` submits a job run and then polls it to completion on the 
worker. By default
+the operator runs in a *durable* mode that makes this crash-safe: the Glue job 
run id is
+persisted to :doc:`task state store 
<apache-airflow:core-concepts/task-state-store>` before
+polling begins, so if the worker crashes or is preempted and the task is 
retried, the operator
+reconnects to the run that is already executing in Glue instead of starting a 
new one.
+
+This matters more for Glue because a Glue job's ``concurrent_run_limit`` 
defaults to ``1``, so
+submitting a second run while the first is still active does not create a 
harmless duplicate, it
+fails outright with ``ConcurrentRunsExceededException`` and the task keeps 
retrying against a run
+it can never see. Durable execution turns that retry into a normal reconnect.
+
+On retry the operator checks the prior run's state:
+
+* if it is still starting, running, waiting for capacity, or being stopped, 
the operator
+  reconnects and continues polling
+* if it already succeeded, or was stopped outside Airflow, the operator 
returns immediately
+  without resubmitting

Review Comment:
   This path also fires when someone clears the task to force a re-run, not 
just on a retry. The store rows are scoped to `(dag_run_id, task_id, 
map_index)` and clearing a task doesn't delete them (the only cascade is on 
`dag_run` deletion), and `[state_store] clear_on_success` defaults to `False`, 
so the run id is still sitting there after a successful attempt. Clear the task 
in the UI and the next attempt reads `SUCCEEDED` and goes green in a second or 
two without submitting anything to Glue.
   
   That's a bigger change than the retry behaviour this section describes, 
given `durable` now defaults to `True` where `resume_glue_job_on_retry` 
defaulted to `False`. Was it weighed? At minimum it seems worth documenting 
here, with `clear_on_success = True` as the way to get clear-to-rerun back.



##########
providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py:
##########
@@ -384,6 +363,146 @@ def on_kill(self):
             if not response["SuccessfulSubmissions"]:
                 self.log.error("Failed to stop AWS Glue Job: %s. Run Id: %s", 
self.job_name, self._job_run_id)
 
+    def _set_job_run_id(self, context: Context, job_run_id: str) -> None:
+        """
+        Record the run id and surface its console link.
+
+        Called from every path that learns a run id (fresh submit, bootstrap 
search, reconnect) so
+        ``on_kill`` can stop the run and the link is available before polling 
starts. Guarded on the
+        id changing, so a single attempt logs the link once.
+        """
+        if self._job_run_id == job_run_id:
+            return
+        self._job_run_id = job_run_id
+        GlueJobRunDetailsLink.persist(
+            context=context,
+            operator=self,
+            region_name=self.hook.conn_region_name,
+            aws_partition=self.hook.conn_partition,
+            job_name=urllib.parse.quote(self.job_name, safe=""),
+            job_run_id=job_run_id,
+        )
+        self.log.info(
+            "You can monitor this Glue Job run at: %s",
+            GlueJobRunDetailsLink.format_str.format(
+                
aws_domain=GlueJobRunDetailsLink.get_aws_domain(self.hook.conn_partition),
+                region_name=self.hook.conn_region_name,
+                job_name=urllib.parse.quote(self.job_name, safe=""),
+                job_run_id=job_run_id,
+            ),
+        )
+
+    def _build_script_args(self, context: Context) -> dict:
+        script_args = dict(self.script_args)
+        if self.openlineage_inject_parent_job_info:
+            self.log.info("Injecting OpenLineage parent job information into 
Glue job arguments.")
+            script_args = 
inject_parent_job_information_into_glue_arguments(script_args, context)
+        if self.openlineage_inject_transport_info:
+            self.log.info("Injecting OpenLineage transport information into 
Glue job arguments.")
+            script_args = 
inject_transport_information_into_glue_arguments(script_args, context)
+        if self.durable:
+            script_args, _ = self._prepare_script_args_with_task_uuid(context, 
base_args=script_args)
+        return script_args
+
+    def _find_previous_job_run(self, context: Context, task_uuid: str) -> str 
| None:
+        """
+        Look for a Glue job run this task instance already started.
+
+        Checks XCom for a cached run id first; falls back to a task-UUID scan 
of the job's run
+        history when XCom has nothing (e.g. the cached id expired, or this is 
the first retry to
+        run under a provider version that persists it).
+        """
+        ti = context["ti"]
+        previous_job_run_id = ti.xcom_pull(key="glue_job_run_id", 
task_ids=ti.task_id)
+        if previous_job_run_id:
+            try:
+                job_run = self.hook.conn.get_job_run(JobName=self.job_name, 
RunId=previous_job_run_id)
+                state = job_run.get("JobRun", {}).get("JobRunState")
+                self.log.info("Previous Glue job_run_id: %s, state: %s", 
previous_job_run_id, state)
+                if state in ("RUNNING", "STARTING"):
+                    return previous_job_run_id
+            except Exception:
+                self.log.warning("Failed to get previous Glue job run state", 
exc_info=True)
+        else:
+            try:
+                existing = self._find_job_run_id_by_task_uuid(task_uuid)
+                if existing:
+                    existing_job_run_id, existing_job_run_state = existing
+                    self.log.info(
+                        "Found Glue job_run_id by task UUID: %s, state: %s",
+                        existing_job_run_id,
+                        existing_job_run_state,
+                    )
+                    if existing_job_run_state in ("RUNNING", "STARTING"):
+                        ti.xcom_push(key="glue_job_run_id", 
value=existing_job_run_id)
+                        return existing_job_run_id
+            except Exception:
+                self.log.warning("Failed to find previous Glue job run by task 
UUID", exc_info=True)
+        return None
+
+    def submit_job(self, context: Context) -> str:
+        """Start a Glue job run and return its run id, or reconnect to one 
this task already started."""
+        script_args = self._build_script_args(context)
+        if self.durable:
+            existing_job_run_id = self._find_previous_job_run(context, 
script_args[self.TASK_UUID_ARG])
+            if existing_job_run_id:
+                self._set_job_run_id(context, existing_job_run_id)
+                return existing_job_run_id
+        self.log.info(
+            "Initializing AWS Glue Job: %s. Wait for completion: %s",
+            self.job_name,
+            self.wait_for_completion,
+        )
+        glue_job_run = self.hook.initialize_job(script_args, 
self.run_job_kwargs)
+        # Set immediately (before any polling) so on_kill can stop the run 
even if the worker dies
+        # before poll_until_complete runs.
+        self._set_job_run_id(context, glue_job_run["JobRunId"])
+        return glue_job_run["JobRunId"]
+
+    def get_job_status(self, external_id: JsonValue, context: Context) -> str:
+        """Query the raw job run state; a run id Glue no longer knows about 
degrades to NOT_FOUND."""
+        job_run_id = cast("str", external_id)
+        # This is the first place a reconnecting attempt learns the run id.
+        self._set_job_run_id(context, job_run_id)
+        try:
+            return self.hook.get_job_state(self.job_name, job_run_id)
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "EntityNotFoundException":
+                return "NOT_FOUND"
+            raise
+
+    def is_job_active(self, status: str) -> bool:
+        return status not in (*JOB_RUN_TERMINAL_STATES, NOT_FOUND_STATE)
+
+    def is_job_succeeded(self, status: str) -> bool:
+        if status == "STOPPED":

Review Comment:
   `STOPPED` isn't only reachable from a manual console stop. `on_kill` stops 
the run itself when `stop_job_run_on_kill=True`, and `on_kill` fires both on 
SIGTERM and on `execution_timeout` (`task_runner.py` calls `task.on_kill()` in 
both), which leave the task `UP_FOR_RETRY` when retries remain. So a task that 
trips `execution_timeout` stops its own Glue run, and the retry then reads 
`STOPPED` and reports success with the ETL half done. Before this PR that retry 
resubmitted.
   
   The warning also says "stopped outside Airflow", which is the one case it 
isn't. Could `on_kill` drop the stored run id when it stops the run, so the 
retry resubmits instead of inheriting a `STOPPED` it caused itself?



##########
providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py:
##########
@@ -384,6 +363,146 @@ def on_kill(self):
             if not response["SuccessfulSubmissions"]:
                 self.log.error("Failed to stop AWS Glue Job: %s. Run Id: %s", 
self.job_name, self._job_run_id)
 
+    def _set_job_run_id(self, context: Context, job_run_id: str) -> None:
+        """
+        Record the run id and surface its console link.
+
+        Called from every path that learns a run id (fresh submit, bootstrap 
search, reconnect) so
+        ``on_kill`` can stop the run and the link is available before polling 
starts. Guarded on the
+        id changing, so a single attempt logs the link once.
+        """
+        if self._job_run_id == job_run_id:
+            return
+        self._job_run_id = job_run_id
+        GlueJobRunDetailsLink.persist(
+            context=context,
+            operator=self,
+            region_name=self.hook.conn_region_name,
+            aws_partition=self.hook.conn_partition,
+            job_name=urllib.parse.quote(self.job_name, safe=""),
+            job_run_id=job_run_id,
+        )
+        self.log.info(
+            "You can monitor this Glue Job run at: %s",
+            GlueJobRunDetailsLink.format_str.format(
+                
aws_domain=GlueJobRunDetailsLink.get_aws_domain(self.hook.conn_partition),
+                region_name=self.hook.conn_region_name,
+                job_name=urllib.parse.quote(self.job_name, safe=""),
+                job_run_id=job_run_id,
+            ),
+        )
+
+    def _build_script_args(self, context: Context) -> dict:
+        script_args = dict(self.script_args)
+        if self.openlineage_inject_parent_job_info:
+            self.log.info("Injecting OpenLineage parent job information into 
Glue job arguments.")
+            script_args = 
inject_parent_job_information_into_glue_arguments(script_args, context)
+        if self.openlineage_inject_transport_info:
+            self.log.info("Injecting OpenLineage transport information into 
Glue job arguments.")
+            script_args = 
inject_transport_information_into_glue_arguments(script_args, context)
+        if self.durable:
+            script_args, _ = self._prepare_script_args_with_task_uuid(context, 
base_args=script_args)
+        return script_args
+
+    def _find_previous_job_run(self, context: Context, task_uuid: str) -> str 
| None:
+        """
+        Look for a Glue job run this task instance already started.
+
+        Checks XCom for a cached run id first; falls back to a task-UUID scan 
of the job's run
+        history when XCom has nothing (e.g. the cached id expired, or this is 
the first retry to
+        run under a provider version that persists it).
+        """
+        ti = context["ti"]
+        previous_job_run_id = ti.xcom_pull(key="glue_job_run_id", 
task_ids=ti.task_id)
+        if previous_job_run_id:
+            try:
+                job_run = self.hook.conn.get_job_run(JobName=self.job_name, 
RunId=previous_job_run_id)
+                state = job_run.get("JobRun", {}).get("JobRunState")
+                self.log.info("Previous Glue job_run_id: %s, state: %s", 
previous_job_run_id, state)
+                if state in ("RUNNING", "STARTING"):
+                    return previous_job_run_id
+            except Exception:
+                self.log.warning("Failed to get previous Glue job run state", 
exc_info=True)
+        else:
+            try:
+                existing = self._find_job_run_id_by_task_uuid(task_uuid)
+                if existing:
+                    existing_job_run_id, existing_job_run_state = existing
+                    self.log.info(
+                        "Found Glue job_run_id by task UUID: %s, state: %s",
+                        existing_job_run_id,
+                        existing_job_run_state,
+                    )
+                    if existing_job_run_state in ("RUNNING", "STARTING"):
+                        ti.xcom_push(key="glue_job_run_id", 
value=existing_job_run_id)
+                        return existing_job_run_id
+            except Exception:
+                self.log.warning("Failed to find previous Glue job run by task 
UUID", exc_info=True)
+        return None
+
+    def submit_job(self, context: Context) -> str:
+        """Start a Glue job run and return its run id, or reconnect to one 
this task already started."""
+        script_args = self._build_script_args(context)
+        if self.durable:
+            existing_job_run_id = self._find_previous_job_run(context, 
script_args[self.TASK_UUID_ARG])

Review Comment:
   This runs on the first attempt too, not just retries, and with `durable` 
defaulting to `True` every Glue task now pays it. Because nothing pushes the 
XCom key any more, `_find_previous_job_run` always takes the scan branch, and 
`_find_job_run_id_by_task_uuid` paginates the entire run history at 50 per page 
when there's no match, so a busy job is dozens of `GetJobRuns` calls before 
every submission. It also runs on the deferrable path, since `execute` calls 
`submit_job` there.
   
   Can it be gated on the attempt actually being a retry 
(`context["ti"].try_number > 1`), or on the task state store being absent, 
rather than on `durable` alone? Line 174 of the docs describes it as the 
fallback for when the store is unavailable, which is what I expected the code 
to do.



##########
providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py:
##########
@@ -384,6 +363,146 @@ def on_kill(self):
             if not response["SuccessfulSubmissions"]:
                 self.log.error("Failed to stop AWS Glue Job: %s. Run Id: %s", 
self.job_name, self._job_run_id)
 
+    def _set_job_run_id(self, context: Context, job_run_id: str) -> None:
+        """
+        Record the run id and surface its console link.
+
+        Called from every path that learns a run id (fresh submit, bootstrap 
search, reconnect) so
+        ``on_kill`` can stop the run and the link is available before polling 
starts. Guarded on the
+        id changing, so a single attempt logs the link once.
+        """
+        if self._job_run_id == job_run_id:
+            return
+        self._job_run_id = job_run_id
+        GlueJobRunDetailsLink.persist(
+            context=context,
+            operator=self,
+            region_name=self.hook.conn_region_name,
+            aws_partition=self.hook.conn_partition,
+            job_name=urllib.parse.quote(self.job_name, safe=""),
+            job_run_id=job_run_id,
+        )
+        self.log.info(
+            "You can monitor this Glue Job run at: %s",
+            GlueJobRunDetailsLink.format_str.format(
+                
aws_domain=GlueJobRunDetailsLink.get_aws_domain(self.hook.conn_partition),
+                region_name=self.hook.conn_region_name,
+                job_name=urllib.parse.quote(self.job_name, safe=""),
+                job_run_id=job_run_id,
+            ),
+        )
+
+    def _build_script_args(self, context: Context) -> dict:
+        script_args = dict(self.script_args)
+        if self.openlineage_inject_parent_job_info:
+            self.log.info("Injecting OpenLineage parent job information into 
Glue job arguments.")
+            script_args = 
inject_parent_job_information_into_glue_arguments(script_args, context)
+        if self.openlineage_inject_transport_info:
+            self.log.info("Injecting OpenLineage transport information into 
Glue job arguments.")
+            script_args = 
inject_transport_information_into_glue_arguments(script_args, context)
+        if self.durable:
+            script_args, _ = self._prepare_script_args_with_task_uuid(context, 
base_args=script_args)
+        return script_args
+
+    def _find_previous_job_run(self, context: Context, task_uuid: str) -> str 
| None:
+        """
+        Look for a Glue job run this task instance already started.
+
+        Checks XCom for a cached run id first; falls back to a task-UUID scan 
of the job's run
+        history when XCom has nothing (e.g. the cached id expired, or this is 
the first retry to
+        run under a provider version that persists it).
+        """
+        ti = context["ti"]
+        previous_job_run_id = ti.xcom_pull(key="glue_job_run_id", 
task_ids=ti.task_id)
+        if previous_job_run_id:
+            try:
+                job_run = self.hook.conn.get_job_run(JobName=self.job_name, 
RunId=previous_job_run_id)
+                state = job_run.get("JobRun", {}).get("JobRunState")
+                self.log.info("Previous Glue job_run_id: %s, state: %s", 
previous_job_run_id, state)
+                if state in ("RUNNING", "STARTING"):
+                    return previous_job_run_id
+            except Exception:
+                self.log.warning("Failed to get previous Glue job run state", 
exc_info=True)
+        else:
+            try:
+                existing = self._find_job_run_id_by_task_uuid(task_uuid)
+                if existing:
+                    existing_job_run_id, existing_job_run_state = existing
+                    self.log.info(
+                        "Found Glue job_run_id by task UUID: %s, state: %s",
+                        existing_job_run_id,
+                        existing_job_run_state,
+                    )
+                    if existing_job_run_state in ("RUNNING", "STARTING"):
+                        ti.xcom_push(key="glue_job_run_id", 
value=existing_job_run_id)
+                        return existing_job_run_id
+            except Exception:
+                self.log.warning("Failed to find previous Glue job run by task 
UUID", exc_info=True)
+        return None
+
+    def submit_job(self, context: Context) -> str:
+        """Start a Glue job run and return its run id, or reconnect to one 
this task already started."""
+        script_args = self._build_script_args(context)
+        if self.durable:
+            existing_job_run_id = self._find_previous_job_run(context, 
script_args[self.TASK_UUID_ARG])
+            if existing_job_run_id:
+                self._set_job_run_id(context, existing_job_run_id)
+                return existing_job_run_id
+        self.log.info(
+            "Initializing AWS Glue Job: %s. Wait for completion: %s",
+            self.job_name,
+            self.wait_for_completion,
+        )
+        glue_job_run = self.hook.initialize_job(script_args, 
self.run_job_kwargs)
+        # Set immediately (before any polling) so on_kill can stop the run 
even if the worker dies
+        # before poll_until_complete runs.
+        self._set_job_run_id(context, glue_job_run["JobRunId"])

Review Comment:
   The `context["ti"].xcom_push(key="glue_job_run_id", ...)` that main does 
right after `initialize_job` is gone, and `_set_job_run_id` doesn't replace it 
(it sets `self._job_run_id` and persists the link; the SDK mixin writes only to 
`task_state_store`). Two consequences.
   
   Downstream tasks doing `xcom_pull(task_ids="glue_job", 
key="glue_job_run_id")` stop getting a value. The `return_value` XCom still 
carries the id, but the explicit key is what people wrote against.
   
   And the XCom tier in `_find_previous_job_run` can never hit now: the push 
inside the scan branch is the only remaining writer, so the `xcom_pull` above 
it is always `None` and every durable attempt falls through to the history 
scan. The description says the XCom-then-scan mechanism is kept exactly as it 
was, but the write half was dropped. The assertions that guarded the push 
(`assert len(xcom_calls) == 1, "Should push new glue_job_run_id"`) were deleted 
from `test_check_previous_job_id_run_new_on_finished` in this PR, which is why 
nothing caught it.



##########
providers/amazon/docs/operators/glue.rst:
##########
@@ -139,6 +139,70 @@ To submit a new AWS Glue job you can use 
:class:`~airflow.providers.amazon.aws.o
   The same AWS IAM role used for the crawler can be used here as well, but it 
will need
   policies to provide access to the output location for result data.
 
+Durable execution
+==================
+
+``GlueJobOperator`` submits a job run and then polls it to completion on the 
worker. By default
+the operator runs in a *durable* mode that makes this crash-safe: the Glue job 
run id is
+persisted to :doc:`task state store 
<apache-airflow:core-concepts/task-state-store>` before
+polling begins, so if the worker crashes or is preempted and the task is 
retried, the operator
+reconnects to the run that is already executing in Glue instead of starting a 
new one.
+
+This matters more for Glue because a Glue job's ``concurrent_run_limit`` 
defaults to ``1``, so
+submitting a second run while the first is still active does not create a 
harmless duplicate, it
+fails outright with ``ConcurrentRunsExceededException`` and the task keeps 
retrying against a run
+it can never see. Durable execution turns that retry into a normal reconnect.
+
+On retry the operator checks the prior run's state:
+
+* if it is still starting, running, waiting for capacity, or being stopped, 
the operator
+  reconnects and continues polling
+* if it already succeeded, or was stopped outside Airflow, the operator 
returns immediately
+  without resubmitting
+* if it failed terminally, or its id has expired and is no longer found, the 
operator submits the
+  job fresh
+
+A run that was stopped outside Airflow (for example, cancelled manually in the 
AWS console) is
+treated as a success rather than resubmitted, since the work is genuinely 
finished, just not the
+way the task expected - the operator logs a warning when this happens.
+
+This protection also applies when ``wait_for_completion=False`` -- even though 
that task attempt
+never polls at all, a retry after a successful submission still reconnects 
rather than
+resubmitting, since the run id is persisted immediately after submission 
regardless of whether the
+task waits for it to finish.
+
+Durable execution requires Airflow 3.3 or newer for the task state store 
lookup above. On earlier
+Airflow versions, or if the task state store is unavailable at runtime, 
``durable=True`` still
+recovers a prior run, just via an older mechanism: the operator checks XCom 
for a cached run id
+first, then falls back to scanning the job's run history for a run tagged with 
this task
+instance's identity, and reconnects if it finds one that is still active.
+
+Like the persisted state itself, the stored run id isn't deleted 
automatically, that only happens
+when someone runs ``airflow state-store clean``. If a task's ``retry_delay`` 
is longer than
+``[state_store] default_retention_days`` (30 days by default) and cleanup runs 
in between, the run
+id won't be there for the next retry, and the operator falls back to the 
XCom/scan mechanism
+above rather than reconnecting via task state store. Avoid running cleanup on 
a schedule shorter
+than your longest ``retry_delay``.
+
+To opt out and always start a fresh run on retry, set ``durable=False``:
+
+.. code-block:: python
+
+  glue_job = GlueJobOperator(
+      task_id="glue_job",
+      job_name="my_glue_job",
+      script_location="s3://glue-examples/glue-scripts/sample_aws_glue_job.py",
+      durable=False,
+  )
+
+Durable execution applies to the synchronous path. When ``deferrable=True`` is 
set, the Triggerer
+already tracks the run across the wait, so deferrable mode takes precedence 
and ``durable`` has no

Review Comment:
   `durable` isn't fully inert on the deferrable path. `execute` calls 
`submit_job`, which for `durable=True` (the new default) injects 
`--airflow_task_uuid` into every run's arguments and runs the reconnect lookup 
before submitting. Deferrable Glue runs get an extra script arg and a history 
scan they didn't have before. Either gate 
`_build_script_args`/`_find_previous_job_run` on `not self.deferrable`, or say 
here what deferrable plus durable actually does.



##########
providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py:
##########
@@ -384,6 +363,146 @@ def on_kill(self):
             if not response["SuccessfulSubmissions"]:
                 self.log.error("Failed to stop AWS Glue Job: %s. Run Id: %s", 
self.job_name, self._job_run_id)
 
+    def _set_job_run_id(self, context: Context, job_run_id: str) -> None:
+        """
+        Record the run id and surface its console link.
+
+        Called from every path that learns a run id (fresh submit, bootstrap 
search, reconnect) so
+        ``on_kill`` can stop the run and the link is available before polling 
starts. Guarded on the
+        id changing, so a single attempt logs the link once.
+        """
+        if self._job_run_id == job_run_id:
+            return
+        self._job_run_id = job_run_id
+        GlueJobRunDetailsLink.persist(
+            context=context,
+            operator=self,
+            region_name=self.hook.conn_region_name,
+            aws_partition=self.hook.conn_partition,
+            job_name=urllib.parse.quote(self.job_name, safe=""),
+            job_run_id=job_run_id,
+        )
+        self.log.info(
+            "You can monitor this Glue Job run at: %s",
+            GlueJobRunDetailsLink.format_str.format(
+                
aws_domain=GlueJobRunDetailsLink.get_aws_domain(self.hook.conn_partition),
+                region_name=self.hook.conn_region_name,
+                job_name=urllib.parse.quote(self.job_name, safe=""),
+                job_run_id=job_run_id,
+            ),
+        )
+
+    def _build_script_args(self, context: Context) -> dict:
+        script_args = dict(self.script_args)
+        if self.openlineage_inject_parent_job_info:
+            self.log.info("Injecting OpenLineage parent job information into 
Glue job arguments.")
+            script_args = 
inject_parent_job_information_into_glue_arguments(script_args, context)
+        if self.openlineage_inject_transport_info:
+            self.log.info("Injecting OpenLineage transport information into 
Glue job arguments.")
+            script_args = 
inject_transport_information_into_glue_arguments(script_args, context)
+        if self.durable:
+            script_args, _ = self._prepare_script_args_with_task_uuid(context, 
base_args=script_args)
+        return script_args
+
+    def _find_previous_job_run(self, context: Context, task_uuid: str) -> str 
| None:
+        """
+        Look for a Glue job run this task instance already started.
+
+        Checks XCom for a cached run id first; falls back to a task-UUID scan 
of the job's run
+        history when XCom has nothing (e.g. the cached id expired, or this is 
the first retry to
+        run under a provider version that persists it).
+        """
+        ti = context["ti"]
+        previous_job_run_id = ti.xcom_pull(key="glue_job_run_id", 
task_ids=ti.task_id)
+        if previous_job_run_id:
+            try:
+                job_run = self.hook.conn.get_job_run(JobName=self.job_name, 
RunId=previous_job_run_id)
+                state = job_run.get("JobRun", {}).get("JobRunState")
+                self.log.info("Previous Glue job_run_id: %s, state: %s", 
previous_job_run_id, state)
+                if state in ("RUNNING", "STARTING"):
+                    return previous_job_run_id
+            except Exception:
+                self.log.warning("Failed to get previous Glue job run state", 
exc_info=True)
+        else:
+            try:
+                existing = self._find_job_run_id_by_task_uuid(task_uuid)
+                if existing:
+                    existing_job_run_id, existing_job_run_state = existing
+                    self.log.info(
+                        "Found Glue job_run_id by task UUID: %s, state: %s",
+                        existing_job_run_id,
+                        existing_job_run_state,
+                    )
+                    if existing_job_run_state in ("RUNNING", "STARTING"):
+                        ti.xcom_push(key="glue_job_run_id", 
value=existing_job_run_id)
+                        return existing_job_run_id
+            except Exception:
+                self.log.warning("Failed to find previous Glue job run by task 
UUID", exc_info=True)
+        return None
+
+    def submit_job(self, context: Context) -> str:
+        """Start a Glue job run and return its run id, or reconnect to one 
this task already started."""
+        script_args = self._build_script_args(context)
+        if self.durable:
+            existing_job_run_id = self._find_previous_job_run(context, 
script_args[self.TASK_UUID_ARG])
+            if existing_job_run_id:
+                self._set_job_run_id(context, existing_job_run_id)
+                return existing_job_run_id
+        self.log.info(
+            "Initializing AWS Glue Job: %s. Wait for completion: %s",
+            self.job_name,
+            self.wait_for_completion,
+        )
+        glue_job_run = self.hook.initialize_job(script_args, 
self.run_job_kwargs)
+        # Set immediately (before any polling) so on_kill can stop the run 
even if the worker dies
+        # before poll_until_complete runs.
+        self._set_job_run_id(context, glue_job_run["JobRunId"])
+        return glue_job_run["JobRunId"]
+
+    def get_job_status(self, external_id: JsonValue, context: Context) -> str:
+        """Query the raw job run state; a run id Glue no longer knows about 
degrades to NOT_FOUND."""
+        job_run_id = cast("str", external_id)
+        # This is the first place a reconnecting attempt learns the run id.
+        self._set_job_run_id(context, job_run_id)
+        try:
+            return self.hook.get_job_state(self.job_name, job_run_id)
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "EntityNotFoundException":
+                return "NOT_FOUND"

Review Comment:
   `NOT_FOUND_STATE` is defined above and used in `is_job_active`, but this 
returns the literal. Worth using the constant so the two can't drift, since a 
mismatch makes `is_job_active` treat a deleted run as active.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to