renamed connection_uuid as session_id_hex

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

Author: saishreeeee

src/databricks/sql/client.py

@@ -309,13 +309,13 @@ def read(self) -> Optional[OAuthToken]:
 
         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 = get_telemetry_client(
-            connection_uuid=self.get_session_id_hex()
+            session_id_hex=self.get_session_id_hex()
         )
 
         driver_connection_params = DriverConnectionParameters(
@@ -427,7 +427,7 @@ def cursor(
         if not self.open:
             raise InterfaceError(
                 "Cannot create cursor from closed connection",
-                connection_uuid=self.get_session_id_hex(),
+                session_id_hex=self.get_session_id_hex(),
             )
 
         cursor = Cursor(
@@ -480,7 +480,7 @@ def commit(self):
     def rollback(self):
         raise NotSupportedError(
             "Transactions are not supported on Databricks",
-            connection_uuid=self.get_session_id_hex(),
+            session_id_hex=self.get_session_id_hex(),
         )
 
 
@@ -535,7 +535,7 @@ def __iter__(self):
         else:
             raise ProgrammingError(
                 "There is no active result set",
-                connection_uuid=self.connection.get_session_id_hex(),
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def _determine_parameter_approach(
@@ -675,7 +675,7 @@ def _check_not_closed(self):
         if not self.open:
             raise InterfaceError(
                 "Attempting operation on closed cursor",
-                connection_uuid=self.connection.get_session_id_hex(),
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def _handle_staging_operation(
@@ -695,7 +695,7 @@ def _handle_staging_operation(
         else:
             raise ProgrammingError(
                 "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(),
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
         abs_staging_allowed_local_paths = [
@@ -725,7 +725,7 @@ def _handle_staging_operation(
             if not allow_operation:
                 raise ProgrammingError(
                     "Local file operations are restricted to paths within the configured staging_allowed_local_path",
-                    connection_uuid=self.connection.get_session_id_hex(),
+                    session_id_hex=self.connection.get_session_id_hex(),
                 )
 
         # May be real headers, or could be json string
@@ -756,7 +756,7 @@ def _handle_staging_operation(
             raise ProgrammingError(
                 f"Operation {row.operation} is not supported. "
                 + "Supported operations are GET, PUT, and REMOVE",
-                connection_uuid=self.connection.get_session_id_hex(),
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def _handle_staging_put(
@@ -770,7 +770,7 @@ def _handle_staging_put(
         if local_file is None:
             raise ProgrammingError(
                 "Cannot perform PUT without specifying a local_file",
-                connection_uuid=self.connection.get_session_id_hex(),
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
         with open(local_file, "rb") as fh:
@@ -789,7 +789,7 @@ def _handle_staging_put(
         if r.status_code not in [OK, CREATED, NO_CONTENT, ACCEPTED]:
             raise OperationalError(
                 f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
-                connection_uuid=self.connection.get_session_id_hex(),
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
         if r.status_code == ACCEPTED:
@@ -809,7 +809,7 @@ def _handle_staging_get(
         if local_file is None:
             raise ProgrammingError(
                 "Cannot perform GET without specifying a local_file",
-                connection_uuid=self.connection.get_session_id_hex(),
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
         r = requests.get(url=presigned_url, headers=headers)
@@ -819,7 +819,7 @@ def _handle_staging_get(
         if not r.ok:
             raise OperationalError(
                 f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
-                connection_uuid=self.connection.get_session_id_hex(),
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
         with open(local_file, "wb") as fp:
@@ -835,7 +835,7 @@ def _handle_staging_remove(
         if not r.ok:
             raise OperationalError(
                 f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
-                connection_uuid=self.connection.get_session_id_hex(),
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def execute(
@@ -1035,7 +1035,7 @@ def get_async_execution_result(self):
         else:
             raise OperationalError(
                 f"get_execution_result failed with Operation status {operation_state}",
-                connection_uuid=self.connection.get_session_id_hex(),
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def executemany(self, operation, seq_of_parameters):
@@ -1187,7 +1187,7 @@ def fetchall(self) -> List[Row]:
         else:
             raise ProgrammingError(
                 "There is no active result set",
-                connection_uuid=self.connection.get_session_id_hex(),
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def fetchone(self) -> Optional[Row]:
@@ -1204,7 +1204,7 @@ def fetchone(self) -> Optional[Row]:
         else:
             raise ProgrammingError(
                 "There is no active result set",
-                connection_uuid=self.connection.get_session_id_hex(),
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def fetchmany(self, size: int) -> List[Row]:
@@ -1229,7 +1229,7 @@ def fetchmany(self, size: int) -> List[Row]:
         else:
             raise ProgrammingError(
                 "There is no active result set",
-                connection_uuid=self.connection.get_session_id_hex(),
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def fetchall_arrow(self) -> "pyarrow.Table":
@@ -1239,7 +1239,7 @@ def fetchall_arrow(self) -> "pyarrow.Table":
         else:
             raise ProgrammingError(
                 "There is no active result set",
-                connection_uuid=self.connection.get_session_id_hex(),
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def fetchmany_arrow(self, size) -> "pyarrow.Table":
@@ -1249,7 +1249,7 @@ def fetchmany_arrow(self, size) -> "pyarrow.Table":
         else:
             raise ProgrammingError(
                 "There is no active result set",
-                connection_uuid=self.connection.get_session_id_hex(),
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def cancel(self) -> None:

src/databricks/sql/exc.py

@@ -14,15 +14,15 @@ class Error(Exception):
     """
 
     def __init__(
-        self, message=None, context=None, connection_uuid=None, *args, **kwargs
+        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 connection_uuid:
-            telemetry_client = get_telemetry_client(connection_uuid)
+        if session_id_hex:
+            telemetry_client = get_telemetry_client(session_id_hex)
             telemetry_client.export_failure_log(error_name, self.message)
 
     def __str__(self):

src/databricks/sql/telemetry/telemetry_client.py

@@ -140,15 +140,15 @@ class TelemetryClient(BaseTelemetryClient):
     def __init__(
         self,
         telemetry_enabled,
-        connection_uuid,
+        session_id_hex,
         auth_provider,
         host_url,
         executor,
     ):
-        logger.debug("Initializing TelemetryClient for connection: %s", connection_uuid)
+        logger.debug("Initializing TelemetryClient for connection: %s", session_id_hex)
         self._telemetry_enabled = telemetry_enabled
         self._batch_size = 10  # TODO: Decide on batch size
-        self._connection_uuid = connection_uuid
+        self._session_id_hex = session_id_hex
         self._auth_provider = auth_provider
         self._user_agent = None
         self._events_batch = []
@@ -159,7 +159,7 @@ def __init__(
 
     def _export_event(self, event):
         """Add an event to the batch queue and flush if batch is full"""
-        logger.debug("Exporting event for connection %s", self._connection_uuid)
+        logger.debug("Exporting event for connection %s", self._session_id_hex)
         with self._lock:
             self._events_batch.append(event)
         if len(self._events_batch) >= self._batch_size:
@@ -230,7 +230,7 @@ def _telemetry_request_callback(self, future):
 
     def export_initial_telemetry_log(self, driver_connection_params, user_agent):
         logger.debug(
-            "Exporting initial telemetry log for connection %s", self._connection_uuid
+            "Exporting initial telemetry log for connection %s", self._session_id_hex
         )
 
         try:
@@ -247,7 +247,7 @@ def export_initial_telemetry_log(self, driver_connection_params, user_agent):
                 ),
                 entry=FrontendLogEntry(
                     sql_driver_log=TelemetryEvent(
-                        session_id=self._connection_uuid,
+                        session_id=self._session_id_hex,
                         system_configuration=TelemetryHelper.get_driver_system_configuration(),
                         driver_connection_params=self._driver_connection_params,
                     )
@@ -260,7 +260,7 @@ def export_initial_telemetry_log(self, driver_connection_params, user_agent):
             logger.debug("Failed to export initial telemetry log: %s", e)
 
     def export_failure_log(self, error_name, error_message):
-        logger.debug("Exporting failure log for connection %s", self._connection_uuid)
+        logger.debug("Exporting failure log for connection %s", self._session_id_hex)
         try:
             error_info = DriverErrorInfo(
                 error_name=error_name, stack_trace=error_message
@@ -275,7 +275,7 @@ def export_failure_log(self, error_name, error_message):
                 ),
                 entry=FrontendLogEntry(
                     sql_driver_log=TelemetryEvent(
-                        session_id=self._connection_uuid,
+                        session_id=self._session_id_hex,
                         system_configuration=TelemetryHelper.get_driver_system_configuration(),
                         driver_connection_params=self._driver_connection_params,
                         error_info=error_info,
@@ -288,9 +288,9 @@ def export_failure_log(self, error_name, error_message):
 
     def close(self):
         """Flush remaining events before closing"""
-        logger.debug("Closing TelemetryClient for connection %s", self._connection_uuid)
+        logger.debug("Closing TelemetryClient for connection %s", self._session_id_hex)
         self._flush()
-        _remove_telemetry_client(self._connection_uuid)
+        _remove_telemetry_client(self._session_id_hex)
 
 
 # Module-level state
@@ -340,55 +340,55 @@ def _handle_unhandled_exception(exc_type, exc_value, exc_traceback):
 
 
 def initialize_telemetry_client(
-    telemetry_enabled, connection_uuid, auth_provider, host_url
+    telemetry_enabled, session_id_hex, auth_provider, host_url
 ):
     """Initialize a telemetry client for a specific connection if telemetry is enabled"""
     try:
         _initialize()
 
         with _lock:
-            if connection_uuid not in _clients:
+            if session_id_hex not in _clients:
                 logger.debug(
-                    "Creating new TelemetryClient for connection %s", connection_uuid
+                    "Creating new TelemetryClient for connection %s", session_id_hex
                 )
                 if telemetry_enabled:
-                    _clients[connection_uuid] = TelemetryClient(
+                    _clients[session_id_hex] = TelemetryClient(
                         telemetry_enabled=telemetry_enabled,
-                        connection_uuid=connection_uuid,
+                        session_id_hex=session_id_hex,
                         auth_provider=auth_provider,
                         host_url=host_url,
                         executor=_executor,
                     )
                 else:
-                    _clients[connection_uuid] = NOOP_TELEMETRY_CLIENT
+                    _clients[session_id_hex] = NOOP_TELEMETRY_CLIENT
     except Exception as e:
         logger.debug("Failed to initialize telemetry client: %s", e)
         # Fallback to NoopTelemetryClient to ensure connection doesn't fail
-        _clients[connection_uuid] = NOOP_TELEMETRY_CLIENT
+        _clients[session_id_hex] = NOOP_TELEMETRY_CLIENT
 
 
-def get_telemetry_client(connection_uuid):
+def get_telemetry_client(session_id_hex):
     """Get the telemetry client for a specific connection"""
     try:
-        if connection_uuid in _clients:
-            return _clients[connection_uuid]
+        if session_id_hex in _clients:
+            return _clients[session_id_hex]
         else:
             logger.error(
-                "Telemetry client not initialized for connection %s", connection_uuid
+                "Telemetry client not initialized for connection %s", session_id_hex
             )
             return NOOP_TELEMETRY_CLIENT
     except Exception as e:
         logger.debug("Failed to get telemetry client: %s", e)
         return NOOP_TELEMETRY_CLIENT
 
 
-def _remove_telemetry_client(connection_uuid):
+def _remove_telemetry_client(session_id_hex):
     """Remove the telemetry client for a specific connection"""
     global _initialized, _executor
     with _lock:
-        if connection_uuid in _clients:
-            logger.debug("Removing telemetry client for connection %s", connection_uuid)
-            _clients.pop(connection_uuid, None)
+        if session_id_hex in _clients:
+            logger.debug("Removing telemetry client for connection %s", session_id_hex)
+            _clients.pop(session_id_hex, None)
 
         # Shutdown executor if no more clients
         if not _clients and _executor: