Pass session_id parameter to telemetry export methods

With host-level telemetry batching, multiple connections share one
TelemetryClient. Each client stores session_id_hex from the first connection
that created it. This caused all subsequent connections' telemetry events
to use the wrong session ID.

Changes:
- Modified telemetry export method signatures to accept optional session_id
- Updated Connection.export_initial_telemetry_log() to pass session_id
- Updated latency_logger.py export_latency_log() to pass session_id
- Updated Error.__init__() to accept optional session_id_hex and pass it
- Updated all error raises in Connection and Cursor to pass session_id_hex

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: samikshya-db

src/databricks/sql/client.py

@@ -391,6 +391,7 @@ def read(self) -> Optional[OAuthToken]:
         self._telemetry_client.export_initial_telemetry_log(
             driver_connection_params=driver_connection_params,
             user_agent=self.session.useragent_header,
+            session_id=self.get_session_id_hex(),
         )
 
     def _set_use_inline_params_with_warning(self, value: Union[bool, str]):
@@ -495,6 +496,7 @@ def cursor(
             raise InterfaceError(
                 "Cannot create cursor from closed connection",
                 host_url=self.session.host,
+                session_id_hex=self.get_session_id_hex(),
             )
 
         cursor = Cursor(
@@ -547,6 +549,7 @@ def autocommit(self) -> bool:
             raise InterfaceError(
                 "Cannot get autocommit on closed connection",
                 host_url=self.session.host,
+                session_id_hex=self.get_session_id_hex(),
             )
 
         if self._fetch_autocommit_from_server:
@@ -579,6 +582,7 @@ def autocommit(self, value: bool) -> None:
             raise InterfaceError(
                 "Cannot set autocommit on closed connection",
                 host_url=self.session.host,
+                session_id_hex=self.get_session_id_hex(),
             )
 
         # Create internal cursor for transaction control
@@ -601,6 +605,7 @@ def autocommit(self, value: bool) -> None:
                     "autocommit_value": value,
                 },
                 host_url=self.session.host,
+                session_id_hex=self.get_session_id_hex(),
             ) from e
         finally:
             if cursor:
@@ -628,6 +633,7 @@ def _fetch_autocommit_state_from_server(self) -> bool:
                     "No result returned from SET AUTOCOMMIT query",
                     context={"operation": "fetch_autocommit"},
                     host_url=self.session.host,
+                    session_id_hex=self.get_session_id_hex(),
                 )
 
             # Parse value (first column should be "true" or "false")
@@ -648,6 +654,7 @@ def _fetch_autocommit_state_from_server(self) -> bool:
                 f"Failed to fetch autocommit state from server: {e.message}",
                 context={**e.context, "operation": "fetch_autocommit"},
                 host_url=self.session.host,
+                session_id_hex=self.get_session_id_hex(),
             ) from e
         finally:
             if cursor:
@@ -681,6 +688,7 @@ def commit(self) -> None:
             raise InterfaceError(
                 "Cannot commit on closed connection",
                 host_url=self.session.host,
+                session_id_hex=self.get_session_id_hex(),
             )
 
         cursor = None
@@ -693,6 +701,7 @@ def commit(self) -> None:
                 f"Failed to commit transaction: {e.message}",
                 context={**e.context, "operation": "commit"},
                 host_url=self.session.host,
+                session_id_hex=self.get_session_id_hex(),
             ) from e
         finally:
             if cursor:
@@ -726,12 +735,14 @@ def rollback(self) -> None:
             raise NotSupportedError(
                 "Transactions are not supported on Databricks",
                 host_url=self.session.host,
+                session_id_hex=self.get_session_id_hex(),
             )
 
         if not self.open:
             raise InterfaceError(
                 "Cannot rollback on closed connection",
                 host_url=self.session.host,
+                session_id_hex=self.get_session_id_hex(),
             )
 
         cursor = None
@@ -744,6 +755,7 @@ def rollback(self) -> None:
                 f"Failed to rollback transaction: {e.message}",
                 context={**e.context, "operation": "rollback"},
                 host_url=self.session.host,
+                session_id_hex=self.get_session_id_hex(),
             ) from e
         finally:
             if cursor:
@@ -768,6 +780,7 @@ def get_transaction_isolation(self) -> str:
             raise InterfaceError(
                 "Cannot get transaction isolation on closed connection",
                 host_url=self.session.host,
+                session_id_hex=self.get_session_id_hex(),
             )
 
         return TRANSACTION_ISOLATION_LEVEL_REPEATABLE_READ
@@ -794,6 +807,7 @@ def set_transaction_isolation(self, level: str) -> None:
             raise InterfaceError(
                 "Cannot set transaction isolation on closed connection",
                 host_url=self.session.host,
+                session_id_hex=self.get_session_id_hex(),
             )
 
         # Normalize and validate isolation level
@@ -806,6 +820,7 @@ def set_transaction_isolation(self, level: str) -> None:
                 f"Setting transaction isolation level '{level}' is not supported. "
                 f"Only {TRANSACTION_ISOLATION_LEVEL_REPEATABLE_READ} is supported.",
                 host_url=self.session.host,
+                session_id_hex=self.get_session_id_hex(),
             )
 
 
@@ -858,6 +873,7 @@ def __iter__(self):
             raise ProgrammingError(
                 "There is no active result set",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def _determine_parameter_approach(
@@ -998,6 +1014,7 @@ def _check_not_closed(self):
             raise InterfaceError(
                 "Attempting operation on closed cursor",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def _handle_staging_operation(
@@ -1042,6 +1059,7 @@ def _handle_staging_operation(
             raise ProgrammingError(
                 "You must provide at least one staging_allowed_local_path when initialising a connection to perform ingestion commands",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
         abs_staging_allowed_local_paths = [
@@ -1068,6 +1086,7 @@ def _handle_staging_operation(
                 raise ProgrammingError(
                     "Local file operations are restricted to paths within the configured staging_allowed_local_path",
                     host_url=self.connection.session.host,
+                    session_id_hex=self.connection.get_session_id_hex(),
                 )
 
         handler_args = {
@@ -1096,6 +1115,7 @@ def _handle_staging_operation(
                 f"Operation {row.operation} is not supported. "
                 + "Supported operations are GET, PUT, and REMOVE",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     @log_latency(StatementType.SQL)
@@ -1111,6 +1131,7 @@ def _handle_staging_put(
             raise ProgrammingError(
                 "Cannot perform PUT without specifying a local_file",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
         with open(local_file, "rb") as fh:
@@ -1136,6 +1157,7 @@ def _handle_staging_http_response(self, r):
             raise OperationalError(
                 f"Staging operation over HTTP was unsuccessful: {r.status}-{error_text}",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
         if r.status == ACCEPTED:
@@ -1167,6 +1189,7 @@ def _handle_staging_put_stream(
             raise ProgrammingError(
                 "No input stream provided for streaming operation",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
         r = self.connection.http_client.request(
@@ -1188,6 +1211,7 @@ def _handle_staging_get(
             raise ProgrammingError(
                 "Cannot perform GET without specifying a local_file",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
         r = self.connection.http_client.request(
@@ -1202,6 +1226,7 @@ def _handle_staging_get(
             raise OperationalError(
                 f"Staging operation over HTTP was unsuccessful: {r.status}-{error_text}",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
         with open(local_file, "wb") as fp:
@@ -1223,6 +1248,7 @@ def _handle_staging_remove(
             raise OperationalError(
                 f"Staging operation over HTTP was unsuccessful: {r.status}-{error_text}",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     @log_latency(StatementType.QUERY)
@@ -1414,6 +1440,7 @@ def get_async_execution_result(self):
             raise OperationalError(
                 f"get_execution_result failed with Operation status {operation_state}",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def executemany(self, operation, seq_of_parameters):
@@ -1542,6 +1569,7 @@ def fetchall(self) -> List[Row]:
             raise ProgrammingError(
                 "There is no active result set",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def fetchone(self) -> Optional[Row]:
@@ -1559,6 +1587,7 @@ def fetchone(self) -> Optional[Row]:
             raise ProgrammingError(
                 "There is no active result set",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def fetchmany(self, size: int) -> List[Row]:
@@ -1584,6 +1613,7 @@ def fetchmany(self, size: int) -> List[Row]:
             raise ProgrammingError(
                 "There is no active result set",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def fetchall_arrow(self) -> "pyarrow.Table":
@@ -1594,6 +1624,7 @@ def fetchall_arrow(self) -> "pyarrow.Table":
             raise ProgrammingError(
                 "There is no active result set",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def fetchmany_arrow(self, size) -> "pyarrow.Table":
@@ -1604,6 +1635,7 @@ def fetchmany_arrow(self, size) -> "pyarrow.Table":
             raise ProgrammingError(
                 "There is no active result set",
                 host_url=self.connection.session.host,
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
     def cancel(self) -> None:

src/databricks/sql/exc.py

@@ -11,7 +11,15 @@ class Error(Exception):
     `context`: Optional extra context about the error. MUST be JSON serializable
     """
 
-    def __init__(self, message=None, context=None, host_url=None, *args, **kwargs):
+    def __init__(
+        self,
+        message=None,
+        context=None,
+        host_url=None,
+        session_id_hex=None,
+        *args,
+        **kwargs,
+    ):
         super().__init__(message, *args, **kwargs)
         self.message = message
         self.context = context or {}
@@ -23,7 +31,9 @@ def __init__(self, message=None, context=None, host_url=None, *args, **kwargs):
             telemetry_client = TelemetryClientFactory.get_telemetry_client(
                 host_url=host_url
             )
-            telemetry_client.export_failure_log(error_name, self.message)
+            telemetry_client.export_failure_log(
+                error_name, self.message, session_id=session_id_hex
+            )
 
     def __str__(self):
         return self.message

src/databricks/sql/telemetry/latency_logger.py

@@ -212,6 +212,7 @@ def wrapper(self, *args, **kwargs):
                                 latency_ms=duration_ms,
                                 sql_execution_event=sql_exec_event,
                                 sql_statement_id=telemetry_data.get("statement_id"),
+                                session_id=session_id_hex,
                             )
 
         return wrapper

src/databricks/sql/telemetry/telemetry_client.py

@@ -147,13 +147,13 @@ def __new__(cls):
                     cls._instance = super(NoopTelemetryClient, cls).__new__(cls)
         return cls._instance
 
-    def export_initial_telemetry_log(self, driver_connection_params, user_agent):
+    def export_initial_telemetry_log(self, driver_connection_params, user_agent, session_id=None):
         pass
 
-    def export_failure_log(self, error_name, error_message):
+    def export_failure_log(self, error_name, error_message, session_id=None):
         pass
 
-    def export_latency_log(self, latency_ms, sql_execution_event, sql_statement_id):
+    def export_latency_log(self, latency_ms, sql_execution_event, sql_statement_id, session_id=None):
         pass
 
     def close(self):
@@ -352,19 +352,22 @@ def _telemetry_request_callback(self, future, sent_count: int):
         except Exception as e:
             logger.debug("Telemetry request failed with exception: %s", e)
 
-    def _export_telemetry_log(self, **telemetry_event_kwargs):
+    def _export_telemetry_log(self, session_id=None, **telemetry_event_kwargs):
         """
         Common helper method for exporting telemetry logs.
 
         Args:
+            session_id: Optional session ID for this event. If not provided, uses the client's session ID.
             **telemetry_event_kwargs: Keyword arguments to pass to TelemetryEvent constructor
         """
-        logger.debug("Exporting telemetry log for connection %s", self._session_id_hex)
+        # Use provided session_id or fall back to client's session_id
+        actual_session_id = session_id or self._session_id_hex
+        logger.debug("Exporting telemetry log for connection %s", actual_session_id)
 
         try:
             # Set common fields for all telemetry events
             event_kwargs = {
-                "session_id": self._session_id_hex,
+                "session_id": actual_session_id,
                 "system_configuration": TelemetryHelper.get_driver_system_configuration(),
                 "driver_connection_params": self._driver_connection_params,
             }
@@ -387,17 +390,18 @@ def _export_telemetry_log(self, **telemetry_event_kwargs):
         except Exception as e:
             logger.debug("Failed to export telemetry log: %s", e)
 
-    def export_initial_telemetry_log(self, driver_connection_params, user_agent):
+    def export_initial_telemetry_log(self, driver_connection_params, user_agent, session_id=None):
         self._driver_connection_params = driver_connection_params
         self._user_agent = user_agent
-        self._export_telemetry_log()
+        self._export_telemetry_log(session_id=session_id)
 
-    def export_failure_log(self, error_name, error_message):
+    def export_failure_log(self, error_name, error_message, session_id=None):
         error_info = DriverErrorInfo(error_name=error_name, stack_trace=error_message)
-        self._export_telemetry_log(error_info=error_info)
+        self._export_telemetry_log(session_id=session_id, error_info=error_info)
 
-    def export_latency_log(self, latency_ms, sql_execution_event, sql_statement_id):
+    def export_latency_log(self, latency_ms, sql_execution_event, sql_statement_id, session_id=None):
         self._export_telemetry_log(
+            session_id=session_id,
             sql_statement_id=sql_statement_id,
             sql_operation=sql_execution_event,
             operation_latency_ms=latency_ms,