amoghrajesh commented on code in PR #71211:
URL: https://github.com/apache/airflow/pull/71211#discussion_r3756133975
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/glue.py:
##########
@@ -384,6 +366,150 @@ 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 and not self.deferrable:
+ 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)
+ # The --airflow_task_uuid tag is still written on every nondeferred
attempt,
+ # including the first: a later retry's scan can only find it if it was
there from the
+ # start. The lookup itself is skipped on the first attempt, since
nothing tagged with this
+ # task's UUID could possibly exist yet -- paying for a full
run-history scan there would
+ # be pure waste.
+ if self.durable and not self.deferrable and context["ti"].try_number >
1:
+ 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"])
+ # Pushed unconditionally, matching pre-durable-execution behavior:
downstream tasks read
+ # this key directly, and it's also the only writer that makes
_find_previous_job_run's
+ # XCom tier reachable on the next retry.
+ context["ti"].xcom_push(key="glue_job_run_id",
value=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_STATE
+ raise
+
+ def is_job_active(self, status: str) -> bool:
+ return status not in (*JOB_RUN_TERMINAL_STATES, NOT_FOUND_STATE)
Review Comment:
Thanks, detailed review. I confirmed by reproducing your exact scenario.
Went with option 2 - polling now checks the state the hook's poll loop actually
finished on, and raises (new `GlueJobRunStoppedError`) if it's
not a real success.
STOPPING still correctly counts as active - nothing changes there, this only
intervenes at the moment the hook's own finished_states disagrees with ours.
Added a test that reproduces your log sequence exactly (STOPPING -> STOPPED)
and confirmed it fails without the fix.
--
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]
