Allow Databricks SQL hook to cancel timed out queries (#42668)

providers/src/airflow/providers/databricks/exceptions.py

@@ -0,0 +1,32 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# Note: Any AirflowException raised is expected to cause the TaskInstance
+#       to be marked in an ERROR state
+"""Exceptions used by Databricks Provider."""
+
+from __future__ import annotations
+
+from airflow.exceptions import AirflowException
+
+
+class DatabricksSqlExecutionError(AirflowException):
+    """Raised when there is an error in sql execution."""
+
+
+class DatabricksSqlExecutionTimeout(DatabricksSqlExecutionError):
+    """Raised when a sql execution times out."""

providers/src/airflow/providers/databricks/hooks/databricks_sql.py

@@ -16,10 +16,12 @@
 # under the License.
 from __future__ import annotations
 
+import threading
 import warnings
 from collections import namedtuple
 from contextlib import closing
 from copy import copy
+from datetime import timedelta
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -35,8 +37,12 @@
 
 from databricks import sql  # type: ignore[attr-defined]
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import (
+    AirflowException,
+    AirflowProviderDeprecationWarning,
+)
 from airflow.providers.common.sql.hooks.sql import DbApiHook, return_single_query_results
+from airflow.providers.databricks.exceptions import DatabricksSqlExecutionError, DatabricksSqlExecutionTimeout
 from airflow.providers.databricks.hooks.databricks_base import BaseDatabricksHook
 
 if TYPE_CHECKING:
@@ -49,6 +55,16 @@
 T = TypeVar("T")
 
 
+def create_timeout_thread(cur, execution_timeout: timedelta | None) -> threading.Timer | None:
+    if execution_timeout is not None:
+        seconds_to_timeout = execution_timeout.total_seconds()
+        t = threading.Timer(seconds_to_timeout, cur.connection.cancel)
+    else:
+        t = None
+
+    return t
+
+
 class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
     """
     Hook to interact with Databricks SQL.
@@ -184,6 +200,7 @@ def run(
         handler: None = ...,
         split_statements: bool = ...,
         return_last: bool = ...,
+        execution_timeout: timedelta | None = None,
     ) -> None: ...
 
     @overload
@@ -195,6 +212,7 @@ def run(
         handler: Callable[[Any], T] = ...,
         split_statements: bool = ...,
         return_last: bool = ...,
+        execution_timeout: timedelta | None = None,
     ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None: ...
 
     def run(
@@ -205,6 +223,7 @@ def run(
         handler: Callable[[Any], T] | None = None,
         split_statements: bool = True,
         return_last: bool = True,
+        execution_timeout: timedelta | None = None,
     ) -> tuple | list[tuple] | list[list[tuple] | tuple] | None:
         """
         Run a command or a list of commands.
@@ -224,6 +243,8 @@ def run(
         :param return_last: Whether to return result for only last statement or for all after split
         :return: return only result of the LAST SQL expression if handler was provided unless return_last
             is set to False.
+        :param execution_timeout: max time allowed for the execution of this task instance, if it goes beyond
+            it will raise and fail.
         """
         self.descriptions = []
         if isinstance(sql, str):
@@ -248,7 +269,23 @@ def run(
                 self.set_autocommit(conn, autocommit)
 
                 with closing(conn.cursor()) as cur:
-                    self._run_command(cur, sql_statement, parameters)  # type: ignore[attr-defined]
+                    t = create_timeout_thread(cur, execution_timeout)
+
+                    # TODO: adjust this to make testing easier
+                    try:
+                        self._run_command(cur, sql_statement, parameters)  # type: ignore[attr-defined]
+                    except Exception as e:
+                        if t is None or t.is_alive():
+                            raise DatabricksSqlExecutionError(
+                                f"Error running SQL statement: {sql_statement}. {str(e)}"
+                            )
+                        raise DatabricksSqlExecutionTimeout(
+                            f"Timeout threshold exceeded for SQL statement: {sql_statement} was cancelled."
+                        )
+                    finally:
+                        if t is not None:
+                            t.cancel()
+
                     if handler is not None:
                         raw_result = handler(cur)
                         if self.return_tuple:

providers/src/airflow/providers/databricks/operators/databricks_sql.py

@@ -353,3 +353,8 @@ def execute(self, context: Context) -> Any:
         self.log.info("Executing: %s", sql)
         hook = self._get_hook()
         hook.run(sql)
+
+    def on_kill(self) -> None:
+        # NB: on_kill isn't required for this operator since query cancelling gets
+        # handled in `DatabricksSqlHook.run()` method which is called in `execute()`
+        ...