chinwobble commented on a change in pull request #19736:
URL: https://github.com/apache/airflow/pull/19736#discussion_r755019613
##########
File path: airflow/providers/databricks/operators/databricks.py
##########
@@ -343,6 +315,45 @@ def on_kill(self):
self.log.info('Task: %s with run_id: %s was requested to be
cancelled.', self.task_id, self.run_id)
+class DatabricksSubmitRunDeferrableOperator(DatabricksSubmitRunOperator):
+ """Deferrable version of ``DatabricksSubmitRunOperator``"""
+
+ def execute(self, context: Optional[dict]):
+ hook = self._get_hook()
+
+ self.run_id = hook.submit_run(self.json)
+ if self.do_xcom_push:
+ context['ti'].xcom_push(key=XCOM_RUN_ID_KEY, value=self.run_id)
+ self.log.info(f'Run submitted with run_id: {self.run_id}')
+
+ run_page_url = hook.get_run_page_url(self.run_id)
+ if self.do_xcom_push:
+ context['ti'].xcom_push(key=XCOM_RUN_PAGE_URL_KEY,
value=run_page_url)
+ self.log.info(f'View run status, Spark UI, and logs at {run_page_url}')
+
+ self.defer(
+ trigger=DatabricksExecutionTrigger(
+ run_id=self.run_id,
+ databricks_conn_id=self.databricks_conn_id,
+ polling_period_seconds=self.polling_period_seconds,
+ ),
+ method_name=DEFER_METHOD_NAME,
+ )
+
+ def execute_complete(self, context: Optional[dict], event: dict):
+ validate_trigger_event(event)
+ run_state = RunState.from_json(event['run_state'])
+ run_page_url = event['run_page_url']
+ self.log.info('View run status, Spark UI, and logs at %s',
run_page_url)
+
+ if run_state.is_successful:
+ self.log.info('Job run completed successfully.')
+ return
+ else:
+ error_message = f'Job run failed with terminal state: {run_state}'
+ raise AirflowException(error_message)
+
+
class DatabricksRunNowOperator(BaseOperator):
Review comment:
maybe slightly outside the scope of the particular pull request.
It would be nice if implement on the `on_kill` method too.
When a process is cancelled or marked as failed, it should also cancel the
running databricks job run.
To do this, you will need to push the run_id to xcom and retrieve it in the
on_kill method. I'm not sure if this is possible.
--
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]
