amoghrajesh commented on code in PR #76295:
URL: https://github.com/apache/airflow/pull/76295#discussion_r4453271725
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -2275,6 +2286,51 @@ def _monitor_submit_run(self) -> None:
errors = extract_failed_task_errors(self._hook, run, run_state)
self._handle_terminal_run_state(run_state, errors)
+ def submit_job(self, context: Context) -> int | None:
+ return self._launch_job(context=context)
+
+ def get_job_status(self, external_id: JsonValue, context: Context) -> str:
+ self.databricks_run_id = cast("int", external_id)
+ try:
+ run = self._hook.get_run(external_id)
+ except DatabricksApiError as e:
+ # Run history expired or was deleted: neither active nor
succeeded, so the mixin
+ # resubmits fresh instead of failing the task.
+ if e.http_status_code == 404:
+ return "NOT_FOUND:"
+ raise
+ # Databricks splits run state across life_cycle_state and
result_state; the mixin's status
+ # interface is a single string, so encode both and decode them in the
predicates below.
+ state = RunState(**run["state"])
+ return f"{state.life_cycle_state}:{state.result_state}"
+
+ def is_job_active(self, status: str) -> bool:
+ return status.split(":", 1)[0] in (
+ "PENDING",
+ "QUEUED",
+ "RUNNING",
+ "TERMINATING",
+ "BLOCKED",
+ "WAITING_FOR_RETRY",
+ )
+
+ def is_job_succeeded(self, status: str) -> bool:
+ return status == "TERMINATED:SUCCESS"
+
+ def poll_until_complete(self, external_id: JsonValue, context: Context) ->
None:
Review Comment:
Good catch, `poll_until_complete` now calls a shared `_record_run_details`
before `monitor_databricks_job`, so the reconnect path pushes run_id and
run_page_url too. `get_job_result` calls the same helper, and
`_run_xcoms_pushed` keeps it to one push per attempt. The reconnect test now
asserts both XComs.
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -2275,6 +2286,51 @@ def _monitor_submit_run(self) -> None:
errors = extract_failed_task_errors(self._hook, run, run_state)
self._handle_terminal_run_state(run_state, errors)
+ def submit_job(self, context: Context) -> int | None:
Review Comment:
`_launch_job` returns right after submit_run now. The page URL lookup and
the pushes moved into `_record_run_details`, called from the poll hook, so
nothing runs between the run existing and the mixin persisting its id.
--
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]
