feat: Implement host-level telemetry batching to reduce rate limiting (#718)

* feat: Implement host-level telemetry batching to reduce rate limiting

Changes telemetry client architecture from per-session to per-host batching,
matching the JDBC driver implementation. This reduces the number of HTTP
requests to the telemetry endpoint and prevents rate limiting in test
environments.

Key changes:
- Add _TelemetryClientHolder with reference counting for shared clients
- Change TelemetryClientFactory to key clients by host_url instead of session_id
- Add getHostUrlSafely() helper for defensive null handling
- Update all callers (client.py, exc.py, latency_logger.py) to pass host_url

Before: 100 connections to same host = 100 separate TelemetryClients
After:  100 connections to same host = 1 shared TelemetryClient (refcount=100)

This fixes rate limiting issues seen in e2e tests where 300+ parallel
connections were overwhelming the telemetry endpoint with 429 errors.

* chore: Change all telemetry logging to DEBUG level

Reduces log noise by changing all telemetry-related log statements
(info, warning, error) to debug level. Telemetry operations are
background tasks and should not clutter logs with operational messages.

Changes:
- Circuit breaker state changes: info/warning -> debug
- Telemetry send failures: error -> debug
- All telemetry operations now consistently use debug level

* chore: Fix remaining telemetry warning log to debug

Changes remaining logger.warning in telemetry_push_client.py to debug level
for consistency with other telemetry logging.

* fix: Update tests to use host_url instead of session_id_hex

- Update circuit breaker test to check logger.debug instead of logger.info
- Replace all session_id_hex test parameters with host_url
- Apply Black formatting to exc.py and telemetry_client.py

This fixes test failures caused by the signature change from session_id_hex
to host_url in the Error class and TelemetryClientFactory.

* fix: Revert session_id_hex in tests for functions that still use it

Only Error classes changed from session_id_hex to host_url.
Other classes (TelemetryClient, ResultSetDownloadHandler, etc.) still use session_id_hex.

Reverted:
- test_telemetry.py: TelemetryClient and initialize_telemetry_client
- test_downloader.py: ResultSetDownloadHandler
- test_download_manager.py: ResultFileDownloadManager

Kept as host_url:
- test_client.py: Error class instantiation

* fix: Update all Error raises and test calls to use host_url

Changes:
1. client.py: Changed all error raises from session_id_hex to host_url
   - Connection class: session_id_hex=self.get_session_id_hex() -> host_url=self.session.host
   - Cursor class: session_id_hex=self.connection.get_session_id_hex() -> host_url=self.connection.session.host

2. test_telemetry.py: Updated get_telemetry_client() and close() calls
   - get_telemetry_client(session_id) -> get_telemetry_client(host_url)
   - close(session_id) -> close(host_url=host_url)

3. test_telemetry_push_client.py: Changed logger.warning to logger.debug
   - Updated test assertion to match debug logging level

These changes complete the migration from session-level to host-level
telemetry client management.

* fix: Update thrift_backend.py to use host_url instead of session_id_hex

Changes:
1. Added self._host attribute to store server_hostname
2. Updated all error raises to use host_url=self._host
3. Changed method signatures from session_id_hex to host_url:
   - _check_response_for_error
   - _hive_schema_to_arrow_schema
   - _col_to_description
   - _hive_schema_to_description
   - _check_direct_results_for_error
4. Updated all method calls to pass self._host instead of self._session_id_hex

This completes the migration from session-level to host-level error reporting.

* Fix Black formatting by adjusting fmt directive placement

Moved the `# fmt: on` directive to the except block level instead
of inside the if statement to resolve Black parsing confusion.

* Fix telemetry feature flag tests to set mock session host

The tests were failing because they called get_telemetry_client("test")
but the mock session didn't have .host set, so the telemetry client was
registered under a different key (likely None or MagicMock). This caused
the factory to return NoopTelemetryClient instead of the expected client.

Fixed by setting mock_session_instance.host = "test" in all three tests.

* Add teardown_method to clear telemetry factory state between tests

Without this cleanup, tests were sharing telemetry clients because they
all used the same host key ("test"), causing test pollution. The first
test would create an enabled client, and subsequent tests would reuse it
even when they expected a disabled client.

* Clear feature flag context cache in teardown to fix test pollution

The FeatureFlagsContextFactory caches feature flag contexts per session,
causing tests to share the same feature flag state. This resulted in the
first test creating a context with telemetry enabled, and subsequent tests
incorrectly reusing that enabled state even when they expected disabled.

* fix: Access actual client from holder in flush worker

The flush worker was calling _flush() on _TelemetryClientHolder objects
instead of the actual TelemetryClient. Fixed by accessing holder.client
before calling _flush().

Fixes AttributeError in e2e tests: '_TelemetryClientHolder' object has
no attribute '_flush'

* Clear telemetry client cache in e2e test teardown

Added _clients.clear() to the teardown fixture to prevent telemetry
clients from persisting across e2e tests, which was causing session ID
pollution in test_concurrent_queries_sends_telemetry.

* 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>

* Fix Black formatting in telemetry_client.py

* Use 'test-host' instead of 'test' for mock host in telemetry tests

* Replace test-session-id with test-host in test_client.py

* Fix telemetry client lookup to use test-host in tests

* Make session_id_hex keyword-only parameter in Error.__init__

---------

Co-authored-by: Claude <noreply@anthropic.com>

Author: samikshya-db

src/databricks/sql/backend/thrift_backend.py

@@ -163,6 +163,7 @@ def __init__(
         else:
             raise ValueError("No valid connection settings.")
 
+        self._host = server_hostname
         self._initialize_retry_args(kwargs)
         self._use_arrow_native_complex_types = kwargs.get(
             "_use_arrow_native_complex_types", True
@@ -279,14 +280,14 @@ def _initialize_retry_args(self, kwargs):
             )
 
     @staticmethod
-    def _check_response_for_error(response, session_id_hex=None):
+    def _check_response_for_error(response, host_url=None):
         if response.status and response.status.statusCode in [
             ttypes.TStatusCode.ERROR_STATUS,
             ttypes.TStatusCode.INVALID_HANDLE_STATUS,
         ]:
             raise DatabaseError(
                 response.status.errorMessage,
-                session_id_hex=session_id_hex,
+                host_url=host_url,
             )
 
     @staticmethod
@@ -340,7 +341,7 @@ def _handle_request_error(self, error_info, attempt, elapsed):
             network_request_error = RequestError(
                 user_friendly_error_message,
                 full_error_info_context,
-                self._session_id_hex,
+                self._host,
                 error_info.error,
             )
             logger.info(network_request_error.message_with_context())
@@ -461,13 +462,12 @@ def attempt_request(attempt):
                     errno.ECONNRESET,   # |   104  |   54   |
                     errno.ETIMEDOUT,    # |   110  |   60   |
                 ]
+                # fmt: on
 
                 gos_name = TCLIServiceClient.GetOperationStatus.__name__
                 # retry on timeout. Happens a lot in Azure and it is safe as data has not been sent to server yet
                 if method.__name__ == gos_name or err.errno == errno.ETIMEDOUT:
                     retry_delay = bound_retry_delay(attempt, self._retry_delay_default)
-
-                    # fmt: on
                     log_string = f"{gos_name} failed with code {err.errno} and will attempt to retry"
                     if err.errno in info_errs:
                         logger.info(log_string)
@@ -516,9 +516,7 @@ def attempt_request(attempt):
             if not isinstance(response_or_error_info, RequestErrorInfo):
                 # log nothing here, presume that main request logging covers
                 response = response_or_error_info
-                ThriftDatabricksClient._check_response_for_error(
-                    response, self._session_id_hex
-                )
+                ThriftDatabricksClient._check_response_for_error(response, self._host)
                 return response
 
             error_info = response_or_error_info
@@ -533,7 +531,7 @@ def _check_protocol_version(self, t_open_session_resp):
                 "Error: expected server to use a protocol version >= "
                 "SPARK_CLI_SERVICE_PROTOCOL_V2, "
                 "instead got: {}".format(protocol_version),
-                session_id_hex=self._session_id_hex,
+                host_url=self._host,
             )
 
     def _check_initial_namespace(self, catalog, schema, response):
@@ -547,15 +545,15 @@ def _check_initial_namespace(self, catalog, schema, response):
             raise InvalidServerResponseError(
                 "Setting initial namespace not supported by the DBR version, "
                 "Please use a Databricks SQL endpoint or a cluster with DBR >= 9.0.",
-                session_id_hex=self._session_id_hex,
+                host_url=self._host,
             )
 
         if catalog:
             if not response.canUseMultipleCatalogs:
                 raise InvalidServerResponseError(
                     "Unexpected response from server: Trying to set initial catalog to {}, "
                     + "but server does not support multiple catalogs.".format(catalog),  # type: ignore
-                    session_id_hex=self._session_id_hex,
+                    host_url=self._host,
                 )
 
     def _check_session_configuration(self, session_configuration):
@@ -570,7 +568,7 @@ def _check_session_configuration(self, session_configuration):
                     TIMESTAMP_AS_STRING_CONFIG,
                     session_configuration[TIMESTAMP_AS_STRING_CONFIG],
                 ),
-                session_id_hex=self._session_id_hex,
+                host_url=self._host,
             )
 
     def open_session(self, session_configuration, catalog, schema) -> SessionId:
@@ -639,7 +637,7 @@ def _check_command_not_in_error_or_closed_state(
                         and guid_to_hex_id(op_handle.operationId.guid),
                         "diagnostic-info": get_operations_resp.diagnosticInfo,
                     },
-                    session_id_hex=self._session_id_hex,
+                    host_url=self._host,
                 )
             else:
                 raise ServerOperationError(
@@ -649,7 +647,7 @@ def _check_command_not_in_error_or_closed_state(
                         and guid_to_hex_id(op_handle.operationId.guid),
                         "diagnostic-info": None,
                     },
-                    session_id_hex=self._session_id_hex,
+                    host_url=self._host,
                 )
         elif get_operations_resp.operationState == ttypes.TOperationState.CLOSED_STATE:
             raise DatabaseError(
@@ -660,7 +658,7 @@ def _check_command_not_in_error_or_closed_state(
                     "operation-id": op_handle
                     and guid_to_hex_id(op_handle.operationId.guid)
                 },
-                session_id_hex=self._session_id_hex,
+                host_url=self._host,
             )
 
     def _poll_for_status(self, op_handle):
@@ -683,7 +681,7 @@ def _create_arrow_table(self, t_row_set, lz4_compressed, schema_bytes, descripti
         else:
             raise OperationalError(
                 "Unsupported TRowSet instance {}".format(t_row_set),
-                session_id_hex=self._session_id_hex,
+                host_url=self._host,
             )
         return convert_decimals_in_arrow_table(arrow_table, description), num_rows
 
@@ -692,7 +690,7 @@ def _get_metadata_resp(self, op_handle):
         return self.make_request(self._client.GetResultSetMetadata, req)
 
     @staticmethod
-    def _hive_schema_to_arrow_schema(t_table_schema, session_id_hex=None):
+    def _hive_schema_to_arrow_schema(t_table_schema, host_url=None):
         def map_type(t_type_entry):
             if t_type_entry.primitiveEntry:
                 return {
@@ -724,7 +722,7 @@ def map_type(t_type_entry):
                 # even for complex types
                 raise OperationalError(
                     "Thrift protocol error: t_type_entry not a primitiveEntry",
-                    session_id_hex=session_id_hex,
+                    host_url=host_url,
                 )
 
         def convert_col(t_column_desc):
@@ -735,7 +733,7 @@ def convert_col(t_column_desc):
         return pyarrow.schema([convert_col(col) for col in t_table_schema.columns])
 
     @staticmethod
-    def _col_to_description(col, field=None, session_id_hex=None):
+    def _col_to_description(col, field=None, host_url=None):
         type_entry = col.typeDesc.types[0]
 
         if type_entry.primitiveEntry:
@@ -745,7 +743,7 @@ def _col_to_description(col, field=None, session_id_hex=None):
         else:
             raise OperationalError(
                 "Thrift protocol error: t_type_entry not a primitiveEntry",
-                session_id_hex=session_id_hex,
+                host_url=host_url,
             )
 
         if type_entry.primitiveEntry.type == ttypes.TTypeId.DECIMAL_TYPE:
@@ -759,7 +757,7 @@ def _col_to_description(col, field=None, session_id_hex=None):
                 raise OperationalError(
                     "Decimal type did not provide typeQualifier precision, scale in "
                     "primitiveEntry {}".format(type_entry.primitiveEntry),
-                    session_id_hex=session_id_hex,
+                    host_url=host_url,
                 )
         else:
             precision, scale = None, None
@@ -778,9 +776,7 @@ def _col_to_description(col, field=None, session_id_hex=None):
         return col.columnName, cleaned_type, None, None, precision, scale, None
 
     @staticmethod
-    def _hive_schema_to_description(
-        t_table_schema, schema_bytes=None, session_id_hex=None
-    ):
+    def _hive_schema_to_description(t_table_schema, schema_bytes=None, host_url=None):
         field_dict = {}
         if pyarrow and schema_bytes:
             try:
@@ -795,7 +791,7 @@ def _hive_schema_to_description(
             ThriftDatabricksClient._col_to_description(
                 col,
                 field_dict.get(col.columnName) if field_dict else None,
-                session_id_hex,
+                host_url,
             )
             for col in t_table_schema.columns
         ]
@@ -818,7 +814,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
                         t_result_set_metadata_resp.resultFormat
                     ]
                 ),
-                session_id_hex=self._session_id_hex,
+                host_url=self._host,
             )
         direct_results = resp.directResults
         has_been_closed_server_side = direct_results and direct_results.closeOperation
@@ -833,7 +829,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
             schema_bytes = (
                 t_result_set_metadata_resp.arrowSchema
                 or self._hive_schema_to_arrow_schema(
-                    t_result_set_metadata_resp.schema, self._session_id_hex
+                    t_result_set_metadata_resp.schema, self._host
                 )
                 .serialize()
                 .to_pybytes()
@@ -844,7 +840,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
         description = self._hive_schema_to_description(
             t_result_set_metadata_resp.schema,
             schema_bytes,
-            self._session_id_hex,
+            self._host,
         )
 
         lz4_compressed = t_result_set_metadata_resp.lz4Compressed
@@ -895,7 +891,7 @@ def get_execution_result(
             schema_bytes = (
                 t_result_set_metadata_resp.arrowSchema
                 or self._hive_schema_to_arrow_schema(
-                    t_result_set_metadata_resp.schema, self._session_id_hex
+                    t_result_set_metadata_resp.schema, self._host
                 )
                 .serialize()
                 .to_pybytes()
@@ -906,7 +902,7 @@ def get_execution_result(
         description = self._hive_schema_to_description(
             t_result_set_metadata_resp.schema,
             schema_bytes,
-            self._session_id_hex,
+            self._host,
         )
 
         lz4_compressed = t_result_set_metadata_resp.lz4Compressed
@@ -971,27 +967,27 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
         return state
 
     @staticmethod
-    def _check_direct_results_for_error(t_spark_direct_results, session_id_hex=None):
+    def _check_direct_results_for_error(t_spark_direct_results, host_url=None):
         if t_spark_direct_results:
             if t_spark_direct_results.operationStatus:
                 ThriftDatabricksClient._check_response_for_error(
                     t_spark_direct_results.operationStatus,
-                    session_id_hex,
+                    host_url,
                 )
             if t_spark_direct_results.resultSetMetadata:
                 ThriftDatabricksClient._check_response_for_error(
                     t_spark_direct_results.resultSetMetadata,
-                    session_id_hex,
+                    host_url,
                 )
             if t_spark_direct_results.resultSet:
                 ThriftDatabricksClient._check_response_for_error(
                     t_spark_direct_results.resultSet,
-                    session_id_hex,
+                    host_url,
                 )
             if t_spark_direct_results.closeOperation:
                 ThriftDatabricksClient._check_response_for_error(
                     t_spark_direct_results.closeOperation,
-                    session_id_hex,
+                    host_url,
                 )
 
     def execute_command(
@@ -1260,7 +1256,7 @@ def _handle_execute_response(self, resp, cursor):
             raise ValueError(f"Invalid Thrift handle: {resp.operationHandle}")
 
         cursor.active_command_id = command_id
-        self._check_direct_results_for_error(resp.directResults, self._session_id_hex)
+        self._check_direct_results_for_error(resp.directResults, self._host)
 
         final_operation_state = self._wait_until_command_done(
             resp.operationHandle,
@@ -1275,7 +1271,7 @@ def _handle_execute_response_async(self, resp, cursor):
             raise ValueError(f"Invalid Thrift handle: {resp.operationHandle}")
 
         cursor.active_command_id = command_id
-        self._check_direct_results_for_error(resp.directResults, self._session_id_hex)
+        self._check_direct_results_for_error(resp.directResults, self._host)
 
     def fetch_results(
         self,
@@ -1313,7 +1309,7 @@ def fetch_results(
                 "fetch_results failed due to inconsistency in the state between the client and the server. Expected results to start from {} but they instead start at {}, some result batches must have been skipped".format(
                     expected_row_start_offset, resp.results.startRowOffset
                 ),
-                session_id_hex=self._session_id_hex,
+                host_url=self._host,
             )
 
         queue = ThriftResultSetQueueFactory.build_queue(