Skip to content

Commit b5bf165

Browse files
committed
added multi threaded tests
Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>
1 parent 152e0da commit b5bf165

File tree

2 files changed

+406
-8
lines changed

2 files changed

+406
-8
lines changed

src/databricks/sql/telemetry/telemetry_client.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,7 @@ class TelemetryClient(BaseTelemetryClient):
158158
# Telemetry endpoint paths
159159
TELEMETRY_AUTHENTICATED_PATH = "/telemetry-ext"
160160
TELEMETRY_UNAUTHENTICATED_PATH = "/telemetry-unauth"
161+
DEFAULT_BATCH_SIZE = 10
161162

162163
def __init__(
163164
self,
@@ -169,7 +170,7 @@ def __init__(
169170
):
170171
logger.debug("Initializing TelemetryClient for connection: %s", session_id_hex)
171172
self._telemetry_enabled = telemetry_enabled
172-
self._batch_size = 10 # TODO: Decide on batch size
173+
self._batch_size = self.DEFAULT_BATCH_SIZE # TODO: Decide on batch size
173174
self._session_id_hex = session_id_hex
174175
self._auth_provider = auth_provider
175176
self._user_agent = None
@@ -466,6 +467,9 @@ def close(session_id_hex):
466467
logger.debug(
467468
"No more telemetry clients, shutting down thread pool executor"
468469
)
469-
TelemetryClientFactory._executor.shutdown(wait=True)
470+
try:
471+
TelemetryClientFactory._executor.shutdown(wait=True)
472+
except Exception as e:
473+
logger.debug("Failed to shutdown thread pool executor: %s", e)
470474
TelemetryClientFactory._executor = None
471475
TelemetryClientFactory._initialized = False

0 commit comments

Comments
 (0)