eskarimov commented on a change in pull request #19736:
URL: https://github.com/apache/airflow/pull/19736#discussion_r759627482
##########
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:
I've just tried to play with the deferrable operator's behaviour when a
task is killed, and it might be there's a bug with it.
If a task was killed while being executed by `Triggerer`, there's no log
available for this task after kill. Later on, if the same task is started
again, it finishes immediately, like it was continued after being deferred.
Will raise an issue with detailed description how to reproduce it.
--
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]
