amoghrajesh commented on code in PR #76295:
URL: https://github.com/apache/airflow/pull/76295#discussion_r4453415395
##########
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:
Handled in d983bdfd8ba
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -1995,15 +2005,15 @@ def _get_run_json(self) -> dict[str, Any]:
"""Get run json to be used for task submissions."""
if self.new_cluster and self.existing_cluster_id:
raise ValueError("Both new_cluster and existing_cluster_id are
set. Only one should be set.")
- cluster: dict[str, Any]
+ # LOCAL MANUAL-TEST HACK -- NOT PART OF THIS PR. Allows no compute
spec so the run
Review Comment:
Handled in d983bdfd8ba
--
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]
