Added functionality for export of failure logs

src/databricks/sql/client.py

@@ -18,6 +18,9 @@
     OperationalError,
     SessionAlreadyClosedError,
     CursorAlreadyClosedError,
+    InterfaceError,
+    NotSupportedError,
+    ProgrammingError,
 )
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.thrift_backend import ThriftBackend
@@ -50,8 +53,8 @@
     TOperationState,
 )
 from databricks.sql.telemetry.telemetry_client import (
-    TelemetryClientFactory,
     TelemetryHelper,
+    TelemetryClientFactory,
 )
 from databricks.sql.telemetry.models.enums import DatabricksClientType
 from databricks.sql.telemetry.models.event import (
@@ -305,13 +308,13 @@ def read(self) -> Optional[OAuthToken]:
 
         TelemetryClientFactory.initialize_telemetry_client(
             telemetry_enabled=self.telemetry_enabled,
-            connection_uuid=self.get_session_id_hex(),
+            session_id_hex=self.get_session_id_hex(),
             auth_provider=auth_provider,
             host_url=self.host,
         )
 
         self._telemetry_client = TelemetryClientFactory.get_telemetry_client(
-            connection_uuid=self.get_session_id_hex()
+            session_id_hex=self.get_session_id_hex()
         )
 
         driver_connection_params = DriverConnectionParameters(
@@ -421,7 +424,10 @@ def cursor(
         Will throw an Error if the connection has been closed.
         """
         if not self.open:
-            raise Error("Cannot create cursor from closed connection")
+            raise InterfaceError(
+                "Cannot create cursor from closed connection",
+                session_id_hex=self.get_session_id_hex(),
+            )
 
         cursor = Cursor(
             self,
@@ -464,14 +470,17 @@ def _close(self, close_cursors=True) -> None:
 
         self.open = False
 
-        self._telemetry_client.close()
+        TelemetryClientFactory.close(self.get_session_id_hex())
 
     def commit(self):
         """No-op because Databricks does not support transactions"""
         pass
 
     def rollback(self):
-        raise NotSupportedError("Transactions are not supported on Databricks")
+        raise NotSupportedError(
+            "Transactions are not supported on Databricks",
+            session_id_hex=self.get_session_id_hex(),
+        )
 
 
 class Cursor:
@@ -523,7 +532,10 @@ def __iter__(self):
             for row in self.active_result_set:
                 yield row
         else:
-            raise Error("There is no active result set")
+            raise ProgrammingError(
+                "There is no active result set",
+                session_id_hex=self.connection.get_session_id_hex(),
+            )
 
     def _determine_parameter_approach(
         self, params: Optional[TParameterCollection]
@@ -660,7 +672,10 @@ def _close_and_clear_active_result_set(self):
 
     def _check_not_closed(self):
         if not self.open:
-            raise Error("Attempting operation on closed cursor")
+            raise InterfaceError(
+                "Attempting operation on closed cursor",
+                session_id_hex=self.connection.get_session_id_hex(),
+            )
 
     def _handle_staging_operation(
         self, staging_allowed_local_path: Union[None, str, List[str]]
@@ -677,8 +692,9 @@ def _handle_staging_operation(
         elif isinstance(staging_allowed_local_path, type(list())):
             _staging_allowed_local_paths = staging_allowed_local_path
         else:
-            raise Error(
-                "You must provide at least one staging_allowed_local_path when initialising a connection to perform ingestion commands"
+            raise ProgrammingError(
+                "You must provide at least one staging_allowed_local_path when initialising a connection to perform ingestion commands",
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
         abs_staging_allowed_local_paths = [
@@ -706,8 +722,9 @@ def _handle_staging_operation(
                 else:
                     continue
             if not allow_operation:
-                raise Error(
-                    "Local file operations are restricted to paths within the configured staging_allowed_local_path"
+                raise ProgrammingError(
+                    "Local file operations are restricted to paths within the configured staging_allowed_local_path",
+                    session_id_hex=self.connection.get_session_id_hex(),
                 )
 
         # May be real headers, or could be json string
@@ -735,9 +752,10 @@ def _handle_staging_operation(
             handler_args.pop("local_file")
             return self._handle_staging_remove(**handler_args)
         else:
-            raise Error(
+            raise ProgrammingError(
                 f"Operation {row.operation} is not supported. "
-                + "Supported operations are GET, PUT, and REMOVE"
+                + "Supported operations are GET, PUT, and REMOVE",
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def _handle_staging_put(
@@ -749,7 +767,10 @@ def _handle_staging_put(
         """
 
         if local_file is None:
-            raise Error("Cannot perform PUT without specifying a local_file")
+            raise ProgrammingError(
+                "Cannot perform PUT without specifying a local_file",
+                session_id_hex=self.connection.get_session_id_hex(),
+            )
 
         with open(local_file, "rb") as fh:
             r = requests.put(url=presigned_url, data=fh, headers=headers)
@@ -765,8 +786,9 @@ def _handle_staging_put(
         # fmt: on
 
         if r.status_code not in [OK, CREATED, NO_CONTENT, ACCEPTED]:
-            raise Error(
-                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}"
+            raise OperationalError(
+                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
         if r.status_code == ACCEPTED:
@@ -784,15 +806,19 @@ def _handle_staging_get(
         """
 
         if local_file is None:
-            raise Error("Cannot perform GET without specifying a local_file")
+            raise ProgrammingError(
+                "Cannot perform GET without specifying a local_file",
+                session_id_hex=self.connection.get_session_id_hex(),
+            )
 
         r = requests.get(url=presigned_url, headers=headers)
 
         # response.ok verifies the status code is not between 400-600.
         # Any 2xx or 3xx will evaluate r.ok == True
         if not r.ok:
-            raise Error(
-                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}"
+            raise OperationalError(
+                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
         with open(local_file, "wb") as fp:
@@ -806,8 +832,9 @@ def _handle_staging_remove(
         r = requests.delete(url=presigned_url, headers=headers)
 
         if not r.ok:
-            raise Error(
-                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}"
+            raise OperationalError(
+                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def execute(
@@ -1005,8 +1032,9 @@ def get_async_execution_result(self):
 
             return self
         else:
-            raise Error(
-                f"get_execution_result failed with Operation status {operation_state}"
+            raise OperationalError(
+                f"get_execution_result failed with Operation status {operation_state}",
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def executemany(self, operation, seq_of_parameters):
@@ -1156,7 +1184,10 @@ def fetchall(self) -> List[Row]:
         if self.active_result_set:
             return self.active_result_set.fetchall()
         else:
-            raise Error("There is no active result set")
+            raise ProgrammingError(
+                "There is no active result set",
+                session_id_hex=self.connection.get_session_id_hex(),
+            )
 
     def fetchone(self) -> Optional[Row]:
         """
@@ -1170,7 +1201,10 @@ def fetchone(self) -> Optional[Row]:
         if self.active_result_set:
             return self.active_result_set.fetchone()
         else:
-            raise Error("There is no active result set")
+            raise ProgrammingError(
+                "There is no active result set",
+                session_id_hex=self.connection.get_session_id_hex(),
+            )
 
     def fetchmany(self, size: int) -> List[Row]:
         """
@@ -1192,21 +1226,30 @@ def fetchmany(self, size: int) -> List[Row]:
         if self.active_result_set:
             return self.active_result_set.fetchmany(size)
         else:
-            raise Error("There is no active result set")
+            raise ProgrammingError(
+                "There is no active result set",
+                session_id_hex=self.connection.get_session_id_hex(),
+            )
 
     def fetchall_arrow(self) -> "pyarrow.Table":
         self._check_not_closed()
         if self.active_result_set:
             return self.active_result_set.fetchall_arrow()
         else:
-            raise Error("There is no active result set")
+            raise ProgrammingError(
+                "There is no active result set",
+                session_id_hex=self.connection.get_session_id_hex(),
+            )
 
     def fetchmany_arrow(self, size) -> "pyarrow.Table":
         self._check_not_closed()
         if self.active_result_set:
             return self.active_result_set.fetchmany_arrow(size)
         else:
-            raise Error("There is no active result set")
+            raise ProgrammingError(
+                "There is no active result set",
+                session_id_hex=self.connection.get_session_id_hex(),
+            )
 
     def cancel(self) -> None:
         """

src/databricks/sql/exc.py

@@ -1,21 +1,32 @@
 import json
 import logging
 
-logger = logging.getLogger(__name__)
+from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 
+logger = logging.getLogger(__name__)
 
 ### PEP-249 Mandated ###
+# https://peps.python.org/pep-0249/#exceptions
 class Error(Exception):
     """Base class for DB-API2.0 exceptions.
     `message`: An optional user-friendly error message. It should be short, actionable and stable
     `context`: Optional extra context about the error. MUST be JSON serializable
     """
 
-    def __init__(self, message=None, context=None, *args, **kwargs):
+    def __init__(
+        self, message=None, context=None, session_id_hex=None, *args, **kwargs
+    ):
         super().__init__(message, *args, **kwargs)
         self.message = message
         self.context = context or {}
 
+        error_name = self.__class__.__name__
+        if session_id_hex:
+            telemetry_client = TelemetryClientFactory.get_telemetry_client(
+                session_id_hex
+            )
+            telemetry_client.export_failure_log(error_name, self.message)
+
     def __str__(self):
         return self.message