andrewgodwin commented on a change in pull request #19736:
URL: https://github.com/apache/airflow/pull/19736#discussion_r759631853
##########
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:
Ah yes, I can see how that might leave things in a strange state. We may
need to add code to do better cleanup when a task is killed-while-deferred.
--
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]
