chinwobble commented on code in PR #19736:
URL: https://github.com/apache/airflow/pull/19736#discussion_r878773432
##########
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:
@andrewgodwin @alexott
Any thoughts on how this should be handled?
--
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: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
