telemetry client and logs

Signed-off-by: Sai Shree Pradhan <[email protected]>

Author: saishreeeee

src/databricks/sql/client.py

@@ -43,13 +43,16 @@
 
 from databricks.sql.types import Row, SSLOptions
 from databricks.sql.auth.auth import get_python_sql_connector_auth_provider
-from databricks.sql.experimental.oauth_persistence import OAuthPersistence
-
+from databricks.sql.telemetry.telemetry_client import (
+    telemetry_client,
+    NoopTelemetryClient,
+)
 from databricks.sql.thrift_api.TCLIService.ttypes import (
     TSparkParameter,
     TOperationState,
 )
-
+from databricks.sql.telemetry.latency_logger import log_latency
+from databricks.sql.telemetry.models.enums import DriverVolumeOperationType
 
 logger = logging.getLogger(__name__)
 
@@ -238,6 +241,9 @@ def read(self) -> Optional[OAuthToken]:
         self.telemetry_enabled = (
             self.client_telemetry_enabled and self.server_telemetry_enabled
         )
+        telemetry_batch_size = kwargs.get(
+            "telemetry_batch_size", 100
+        )  # TODO: Decide on batch size
 
         user_agent_entry = kwargs.get("user_agent_entry")
         if user_agent_entry is None:
@@ -294,6 +300,25 @@ def read(self) -> Optional[OAuthToken]:
             kwargs.get("use_inline_params", False)
         )
 
+        if self.telemetry_enabled:
+            telemetry_client.initialize(
+                host=self.host,
+                connection_uuid=self.get_session_id_hex(),
+                auth_provider=auth_provider,
+                is_authenticated=True,  # TODO: Add authentication logic later
+                batch_size=telemetry_batch_size,
+                user_agent=useragent_header,
+            )
+
+            telemetry_client.export_initial_telemetry_log(
+                http_path,
+                self.port,
+                kwargs.get("_socket_timeout", None),
+                self.get_session_id_hex(),
+            )
+        else:
+            self.telemetry_client = NoopTelemetryClient()
+
     def _set_use_inline_params_with_warning(self, value: Union[bool, str]):
         """Valid values are True, False, and "silent"
 
@@ -430,6 +455,9 @@ def _close(self, close_cursors=True) -> None:
 
         self.open = False
 
+        if self.telemetry_enabled:
+            telemetry_client.close(self.get_session_id_hex())
+
     def commit(self):
         """No-op because Databricks does not support transactions"""
         pass
@@ -487,7 +515,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]
@@ -624,7 +655,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]]
@@ -642,7 +676,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 = [
@@ -671,7 +706,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
@@ -701,9 +737,11 @@ 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(),
             )
 
+    @log_latency()
     def _handle_staging_put(
         self, presigned_url: str, local_file: str, headers: Optional[dict] = None
     ):
@@ -713,7 +751,13 @@ 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(),
+            )
+
+        self.volume_operation_type = DriverVolumeOperationType.PUT
+        self.volume_path = local_file
 
         with open(local_file, "rb") as fh:
             r = requests.put(url=presigned_url, data=fh, headers=headers)
@@ -730,7 +774,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:
@@ -739,6 +784,7 @@ def _handle_staging_put(
                 + "but not yet applied on the server. It's possible this command may fail later."
             )
 
+    @log_latency()
     def _handle_staging_get(
         self, local_file: str, presigned_url: str, headers: Optional[dict] = None
     ):
@@ -748,25 +794,38 @@ 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(),
+            )
+
+        self.volume_operation_type = DriverVolumeOperationType.GET
+        self.volume_path = local_file
 
         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:
             fp.write(r.content)
 
+    @log_latency()
     def _handle_staging_remove(
         self, presigned_url: str, headers: Optional[dict] = None
     ):
         """Make an HTTP DELETE request to the presigned_url"""
 
+        self.volume_operation_type = DriverVolumeOperationType.DELETE
+        self.volume_path = (
+            presigned_url  # Using presigned URL as path since there's no local file
+        )
+
         r = requests.delete(url=presigned_url, headers=headers)
 
         if not r.ok:
@@ -970,7 +1029,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):
@@ -1120,7 +1180,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]:
         """
@@ -1134,7 +1197,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]:
         """
@@ -1156,21 +1222,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,5 +1,6 @@
 import json
 import logging
+from databricks.sql.telemetry.telemetry_client import telemetry_client
 
 logger = logging.getLogger(__name__)
 
@@ -11,10 +12,22 @@ 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
+
+        try:
+            error_name = self.__class__.__name__
+            error_message = self.message_with_context()
+            telemetry_client.export_failure_log(
+                error_name, error_message, self.connection_uuid
+            )
+        except Exception as e:
+            logger.warning(f"Failed to log error to telemetry: {str(e)}")
 
     def __str__(self):
         return self.message