added functionality for export of failure logs

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

src/databricks/sql/client.py

@@ -421,7 +421,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 Error(
+                "Cannot create cursor from closed connection",
+                connection_uuid=self.get_session_id_hex(),
+            )
 
         cursor = Cursor(
             self,
@@ -471,7 +474,10 @@ def commit(self):
         pass
 
     def rollback(self):
-        raise NotSupportedError("Transactions are not supported on Databricks")
+        raise NotSupportedError(
+            "Transactions are not supported on Databricks",
+            connection_uuid=self.get_session_id_hex(),
+        )
 
 
 class Cursor:
@@ -523,7 +529,10 @@ def __iter__(self):
             for row in self.active_result_set:
                 yield row
         else:
-            raise Error("There is no active result set")
+            raise Error(
+                "There is no active result set",
+                connection_uuid=self.connection.get_session_id_hex(),
+            )
 
     def _determine_parameter_approach(
         self, params: Optional[TParameterCollection]
@@ -660,7 +669,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 Error(
+                "Attempting operation on closed cursor",
+                connection_uuid=self.connection.get_session_id_hex(),
+            )
 
     def _handle_staging_operation(
         self, staging_allowed_local_path: Union[None, str, List[str]]
@@ -678,7 +690,8 @@ def _handle_staging_operation(
             _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"
+                "You must provide at least one staging_allowed_local_path when initialising a connection to perform ingestion commands",
+                connection_uuid=self.connection.get_session_id_hex(),
             )
 
         abs_staging_allowed_local_paths = [
@@ -707,7 +720,8 @@ def _handle_staging_operation(
                     continue
             if not allow_operation:
                 raise Error(
-                    "Local file operations are restricted to paths within the configured staging_allowed_local_path"
+                    "Local file operations are restricted to paths within the configured staging_allowed_local_path",
+                    connection_uuid=self.connection.get_session_id_hex(),
                 )
 
         # May be real headers, or could be json string
@@ -737,7 +751,8 @@ def _handle_staging_operation(
         else:
             raise Error(
                 f"Operation {row.operation} is not supported. "
-                + "Supported operations are GET, PUT, and REMOVE"
+                + "Supported operations are GET, PUT, and REMOVE",
+                connection_uuid=self.connection.get_session_id_hex(),
             )
 
     def _handle_staging_put(
@@ -749,7 +764,10 @@ def _handle_staging_put(
         """
 
         if local_file is None:
-            raise Error("Cannot perform PUT without specifying a local_file")
+            raise Error(
+                "Cannot perform PUT without specifying a local_file",
+                connection_uuid=self.connection.get_session_id_hex(),
+            )
 
         with open(local_file, "rb") as fh:
             r = requests.put(url=presigned_url, data=fh, headers=headers)
@@ -766,7 +784,8 @@ def _handle_staging_put(
 
         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}"
+                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
+                connection_uuid=self.connection.get_session_id_hex(),
             )
 
         if r.status_code == ACCEPTED:
@@ -784,15 +803,19 @@ def _handle_staging_get(
         """
 
         if local_file is None:
-            raise Error("Cannot perform GET without specifying a local_file")
+            raise Error(
+                "Cannot perform GET without specifying a local_file",
+                connection_uuid=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}"
+                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
+                connection_uuid=self.connection.get_session_id_hex(),
             )
 
         with open(local_file, "wb") as fp:
@@ -807,7 +830,8 @@ def _handle_staging_remove(
 
         if not r.ok:
             raise Error(
-                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}"
+                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
+                connection_uuid=self.connection.get_session_id_hex(),
             )
 
     def execute(
@@ -1006,7 +1030,8 @@ def get_async_execution_result(self):
             return self
         else:
             raise Error(
-                f"get_execution_result failed with Operation status {operation_state}"
+                f"get_execution_result failed with Operation status {operation_state}",
+                connection_uuid=self.connection.get_session_id_hex(),
             )
 
     def executemany(self, operation, seq_of_parameters):
@@ -1156,7 +1181,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 Error(
+                "There is no active result set",
+                connection_uuid=self.connection.get_session_id_hex(),
+            )
 
     def fetchone(self) -> Optional[Row]:
         """
@@ -1170,7 +1198,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 Error(
+                "There is no active result set",
+                connection_uuid=self.connection.get_session_id_hex(),
+            )
 
     def fetchmany(self, size: int) -> List[Row]:
         """
@@ -1192,21 +1223,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 Error(
+                "There is no active result set",
+                connection_uuid=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 Error(
+                "There is no active result set",
+                connection_uuid=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 Error(
+                "There is no active result set",
+                connection_uuid=self.connection.get_session_id_hex(),
+            )
 
     def cancel(self) -> None:
         """

src/databricks/sql/exc.py

@@ -1,8 +1,10 @@
 import json
 import logging
+import traceback
 
-logger = logging.getLogger(__name__)
+from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 
+logger = logging.getLogger(__name__)
 
 ### PEP-249 Mandated ###
 class Error(Exception):
@@ -11,10 +13,23 @@ class Error(Exception):
     `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, connection_uuid=None, *args, **kwargs
+    ):
         super().__init__(message, *args, **kwargs)
         self.message = message
         self.context = context or {}
+        self.connection_uuid = connection_uuid
+
+        error_name = self.__class__.__name__
+        if self.connection_uuid:
+            try:
+                telemetry_client = TelemetryClientFactory.get_telemetry_client(
+                    self.connection_uuid
+                )
+                telemetry_client.export_failure_log(error_name, self.message)
+            except Exception as telemetry_error:
+                logger.error(f"Failed to send error to telemetry: {telemetry_error}")
 
     def __str__(self):
         return self.message

src/databricks/sql/telemetry/telemetry_client.py

@@ -8,6 +8,7 @@
 from databricks.sql.telemetry.models.event import (
     TelemetryEvent,
     DriverSystemConfiguration,
+    DriverErrorInfo,
 )
 from databricks.sql.telemetry.models.frontend_logs import (
     TelemetryFrontendLog,
@@ -26,30 +27,33 @@
 import uuid
 import locale
 from abc import ABC, abstractmethod
-from databricks.sql import __version__
 
 logger = logging.getLogger(__name__)
 
 
 class TelemetryHelper:
     """Helper class for getting telemetry related information."""
 
-    _DRIVER_SYSTEM_CONFIGURATION = DriverSystemConfiguration(
-        driver_name="Databricks SQL Python Connector",
-        driver_version=__version__,
-        runtime_name=f"Python {sys.version.split()[0]}",
-        runtime_vendor=platform.python_implementation(),
-        runtime_version=platform.python_version(),
-        os_name=platform.system(),
-        os_version=platform.release(),
-        os_arch=platform.machine(),
-        client_app_name=None,  # TODO: Add client app name
-        locale_name=locale.getlocale()[0] or locale.getdefaultlocale()[0],
-        char_set_encoding=sys.getdefaultencoding(),
-    )
+    _DRIVER_SYSTEM_CONFIGURATION = None
 
     @classmethod
     def getDriverSystemConfiguration(cls) -> DriverSystemConfiguration:
+        if cls._DRIVER_SYSTEM_CONFIGURATION is None:
+            from databricks.sql import __version__
+
+            cls._DRIVER_SYSTEM_CONFIGURATION = DriverSystemConfiguration(
+                driver_name="Databricks SQL Python Connector",
+                driver_version=__version__,
+                runtime_name=f"Python {sys.version.split()[0]}",
+                runtime_vendor=platform.python_implementation(),
+                runtime_version=platform.python_version(),
+                os_name=platform.system(),
+                os_version=platform.release(),
+                os_arch=platform.machine(),
+                client_app_name=None,  # TODO: Add client app name
+                locale_name=locale.getlocale()[0] or locale.getdefaultlocale()[0],
+                char_set_encoding=sys.getdefaultencoding(),
+            )
         return cls._DRIVER_SYSTEM_CONFIGURATION
 
     @staticmethod
@@ -99,7 +103,11 @@ class BaseTelemetryClient(ABC):
     """
 
     @abstractmethod
-    def export_initial_telemetry_log(self, **kwargs):
+    def export_initial_telemetry_log(self, driver_connection_params, user_agent):
+        pass
+
+    @abstractmethod
+    def export_failure_log(self, error_name, error_message):
         pass
 
     @abstractmethod
@@ -123,6 +131,9 @@ def __new__(cls):
     def export_initial_telemetry_log(self, driver_connection_params, user_agent):
         pass
 
+    def export_failure_log(self, error_name, error_message):
+        pass
+
     def close(self):
         pass
 
@@ -255,10 +266,33 @@ def export_initial_telemetry_log(self, driver_connection_params, user_agent):
 
         self.export_event(telemetry_frontend_log)
 
+    def export_failure_log(self, error_name, error_message):
+        logger.debug("Exporting failure log for connection %s", self._connection_uuid)
+        error_info = DriverErrorInfo(error_name=error_name, stack_trace=error_message)
+        telemetry_frontend_log = TelemetryFrontendLog(
+            frontend_log_event_id=str(uuid.uuid4()),
+            context=FrontendLogContext(
+                client_context=TelemetryClientContext(
+                    timestamp_millis=int(time.time() * 1000),
+                    user_agent=self._user_agent,
+                )
+            ),
+            entry=FrontendLogEntry(
+                sql_driver_log=TelemetryEvent(
+                    session_id=self._connection_uuid,
+                    system_configuration=TelemetryHelper.getDriverSystemConfiguration(),
+                    driver_connection_params=self._driver_connection_params,
+                    error_info=error_info,
+                )
+            ),
+        )
+        self.export_event(telemetry_frontend_log)
+
     def close(self):
         """Flush remaining events before closing"""
         logger.debug("Closing TelemetryClient for connection %s", self._connection_uuid)
         self.flush()
+
         TelemetryClientFactory.close(self._connection_uuid)