add deferrable support to DatabricksNotebookOperator

airflow/providers/databricks/operators/databricks.py

@@ -926,6 +926,7 @@ class DatabricksNotebookOperator(BaseOperator):
     """
 
     template_fields = ("notebook_params",)
+    CALLER = "DatabricksNotebookOperator"
 
     def __init__(
         self,
@@ -942,6 +943,7 @@ def __init__(
         databricks_retry_args: dict[Any, Any] | None = None,
         wait_for_termination: bool = True,
         databricks_conn_id: str = "databricks_default",
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs: Any,
     ):
         self.notebook_path = notebook_path
@@ -958,11 +960,12 @@ def __init__(
         self.wait_for_termination = wait_for_termination
         self.databricks_conn_id = databricks_conn_id
         self.databricks_run_id: int | None = None
+        self.deferrable = deferrable
         super().__init__(**kwargs)
 
     @cached_property
     def _hook(self) -> DatabricksHook:
-        return self._get_hook(caller="DatabricksNotebookOperator")
+        return self._get_hook(caller=self.CALLER)
 
     def _get_hook(self, caller: str) -> DatabricksHook:
         return DatabricksHook(
@@ -1041,7 +1044,21 @@ def monitor_databricks_job(self) -> None:
         run = self._hook.get_run(self.databricks_run_id)
         run_state = RunState(**run["state"])
         self.log.info("Current state of the job: %s", run_state.life_cycle_state)
+
         while not run_state.is_terminal:
+            if self.deferrable:
+                return self.defer(
+                    trigger=DatabricksExecutionTrigger(
+                        run_id=self.databricks_run_id,
+                        databricks_conn_id=self.databricks_conn_id,
+                        polling_period_seconds=self.polling_period_seconds,
+                        retry_limit=self.databricks_retry_limit,
+                        retry_delay=self.databricks_retry_delay,
+                        retry_args=self.databricks_retry_args,
+                        caller=self.CALLER,
+                    ),
+                    method_name=DEFER_METHOD_NAME,
+                )
             time.sleep(self.polling_period_seconds)
             run = self._hook.get_run(self.databricks_run_id)
             run_state = RunState(**run["state"])
@@ -1066,3 +1083,6 @@ def execute(self, context: Context) -> None:
         self.launch_notebook_job()
         if self.wait_for_termination:
             self.monitor_databricks_job()
+
+    def execute_complete(self, context: dict | None, event: dict):
+        _handle_deferrable_databricks_operator_completion(event, self.log)

airflow/providers/databricks/triggers/databricks.py

@@ -48,6 +48,7 @@ def __init__(
         retry_args: dict[Any, Any] | None = None,
         run_page_url: str | None = None,
         repair_run: bool = False,
+        caller: str = "DatabricksExecutionTrigger",
     ) -> None:
         super().__init__()
         self.run_id = run_id
@@ -63,6 +64,7 @@ def __init__(
             retry_limit=self.retry_limit,
             retry_delay=self.retry_delay,
             retry_args=retry_args,
+            caller=caller,
         )
 
     def serialize(self) -> tuple[str, dict[str, Any]]: