jplauri commented on a change in pull request #19736:
URL: https://github.com/apache/airflow/pull/19736#discussion_r754526756
##########
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)
Review comment:
For consistency, should this also be an f-string?
--
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]
