Add missing execute_complete method for DatabricksRunNowOperator (#…

…32689)
apache · Jul 24, 2023 · 6313e52 · 6313e52
1 parent 0339f94
commit 6313e52
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/airflow/providers/databricks/operators/databricks.py b/airflow/providers/databricks/operators/databricks.py
@@ -671,6 +671,10 @@ def execute(self, context: Context):
         else:
             _handle_databricks_operator_execution(self, hook, self.log, context)
 
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+        if event:
+            _handle_deferrable_databricks_operator_completion(event, self.log)
+
     def on_kill(self):
         if self.run_id:
             self._hook.cancel_run(self.run_id)
@@ -698,6 +702,3 @@ def execute(self, context):
         hook = self._get_hook(caller="DatabricksRunNowDeferrableOperator")
         self.run_id = hook.run_now(self.json)
         _handle_deferrable_databricks_operator_execution(self, hook, self.log, context)
-
-    def execute_complete(self, context: dict | None, event: dict):
-        _handle_deferrable_databricks_operator_completion(event, self.log)
diff --git a/tests/providers/databricks/operators/test_databricks.py b/tests/providers/databricks/operators/test_databricks.py
@@ -517,7 +517,7 @@ def test_execute_complete_failure(self, db_mock_class):
             op.execute_complete(context=None, event=event)
 
     def test_execute_complete_incorrect_event_validation_failure(self):
-        event = {}
+        event = {"event_id": "no such column"}
         op = DatabricksSubmitRunDeferrableOperator(task_id=TASK_ID)
         with pytest.raises(AirflowException):
             op.execute_complete(context=None, event=event)
@@ -951,7 +951,7 @@ def test_execute_complete_failure(self, db_mock_class):
             op.execute_complete(context=None, event=event)
 
     def test_execute_complete_incorrect_event_validation_failure(self):
-        event = {}
+        event = {"event_id": "no such column"}
         op = DatabricksRunNowDeferrableOperator(task_id=TASK_ID, job_id=JOB_ID)
         with pytest.raises(AirflowException):
             op.execute_complete(context=None, event=event)