Skip to content

Fix #729 and #731: Telemetry lifecycle management #734

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
msrathore-db wants to merge 1 commit into
base: main
Choose a base branch
from
+51 −6
Open

Fix #729 and #731: Telemetry lifecycle management #734

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/databricks/sql/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -306,6 +306,8 @@ def read(self) -> Optional[OAuthToken]:
)
self.session.open()
except Exception as e:
# Respect user's telemetry preference even during connection failure
enable_telemetry = kwargs.get("enable_telemetry", True)
TelemetryClientFactory.connection_failure_log(
error_name="Exception",
error_message=str(e),
Expand All @@ -316,6 +318,7 @@ def read(self) -> Optional[OAuthToken]:
user_agent=self.session.useragent_header
if hasattr(self, "session")
else None,
enable_telemetry=enable_telemetry,
)
raise e

Expand Down
10 changes: 8 additions & 2 deletions src/databricks/sql/common/unified_http_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,15 +217,15 @@ def _should_use_proxy(self, target_host: str) -> bool:
logger.debug("Error checking proxy bypass for host %s: %s", target_host, e)
return True

def _get_pool_manager_for_url(self, url: str) -> urllib3.PoolManager:
def _get_pool_manager_for_url(self, url: str) -> Optional[urllib3.PoolManager]:
"""
Get the appropriate pool manager for the given URL.

Args:
url: The target URL

Returns:
PoolManager instance (either direct or proxy)
PoolManager instance (either direct or proxy), or None if client is closed
"""
parsed_url = urllib.parse.urlparse(url)
target_host = parsed_url.hostname
Expand Down Expand Up @@ -291,6 +291,12 @@ def request_context(
# Select appropriate pool manager based on target URL
pool_manager = self._get_pool_manager_for_url(url)

# DEFENSIVE: Check if pool_manager is None (client closing/closed)
# This prevents AttributeError race condition when telemetry cleanup happens
if pool_manager is None:
logger.debug("HTTP client closing or closed, cannot make request to %s", url)
raise RequestError("HTTP client is closing or has been closed")

response = None

try:
Expand Down
44 changes: 40 additions & 4 deletions src/databricks/sql/telemetry/telemetry_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@
from databricks.sql.common.feature_flag import FeatureFlagsContextFactory
from databricks.sql.common.unified_http_client import UnifiedHttpClient
from databricks.sql.common.http import HttpMethod
from databricks.sql.exc import RequestError
from databricks.sql.telemetry.telemetry_push_client import (
ITelemetryPushClient,
TelemetryPushClient,
Expand Down Expand Up @@ -295,7 +296,7 @@ def _send_telemetry(self, events):
url,
data=request.to_json(),
headers=headers,
timeout=900,
timeout=30,
Copy link
Contributor

@samikshya-db samikshya-db Feb 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are retries etc that incorporate backoff, the push is anyway async, we should not be changing this in my opinion.

)

future.add_done_callback(
Expand Down Expand Up @@ -417,10 +418,38 @@ def export_latency_log(
)

def close(self):
"""Flush remaining events before closing"""
"""Flush remaining events before closing

IMPORTANT: This method does NOT close self._http_client.

Rationale:
- _flush() submits async work to the executor that uses _http_client
- If we closed _http_client here, async callbacks would fail with AttributeError
- Instead, we let _http_client live as long as needed:
* Pending futures hold references to self (via bound methods)
* This keeps self alive, which keeps self._http_client alive
* When all futures complete, Python GC will clean up naturally
- The __del__ method ensures eventual cleanup during garbage collection

This design prevents race conditions while keeping telemetry truly async.
"""
logger.debug("Closing TelemetryClient for connection %s", self._session_id_hex)
self._flush()

def __del__(self):
"""Cleanup when TelemetryClient is garbage collected

This ensures _http_client is eventually closed when the TelemetryClient
object is destroyed. By this point, all async work should be complete
(since the futures held references keeping us alive), so it's safe to
close the http client.
"""
try:
if hasattr(self, '_http_client') and self._http_client:
self._http_client.close()
except Exception:
pass


class _TelemetryClientHolder:
"""
Expand Down Expand Up @@ -674,7 +703,8 @@ def close(host_url):
)
try:
TelemetryClientFactory._stop_flush_thread()
TelemetryClientFactory._executor.shutdown(wait=True)
# Use wait=False to allow process to exit immediately
TelemetryClientFactory._executor.shutdown(wait=False)
except Exception as e:
logger.debug("Failed to shutdown thread pool executor: %s", e)
TelemetryClientFactory._executor = None
Expand All @@ -689,13 +719,19 @@ def connection_failure_log(
port: int,
client_context,
user_agent: Optional[str] = None,
enable_telemetry: bool = True,
):
"""Send error telemetry when connection creation fails, using provided client context"""

# Respect user's telemetry preference - don't force-enable
if not enable_telemetry:
logger.debug("Telemetry disabled, skipping connection failure log")
return

UNAUTH_DUMMY_SESSION_ID = "unauth_session_id"

TelemetryClientFactory.initialize_telemetry_client(
telemetry_enabled=True,
telemetry_enabled=enable_telemetry,
Copy link
Contributor

@samikshya-db samikshya-db Feb 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit : this change is not required.

session_id_hex=UNAUTH_DUMMY_SESSION_ID,
auth_provider=None,
host_url=host_url,
Expand Down
Loading