databricks-sql-connector 4.2.3__tar.gz → 4.2.5__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (69) hide show
  1. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/CHANGELOG.md +8 -0
  2. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/PKG-INFO +1 -1
  3. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/pyproject.toml +2 -2
  4. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/__init__.py +1 -1
  5. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/auth/auth_utils.py +0 -17
  6. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/auth/retry.py +7 -27
  7. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/auth/token_federation.py +3 -3
  8. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/sea/utils/http_client.py +4 -2
  9. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/client.py +5 -0
  10. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/common/feature_flag.py +3 -1
  11. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/common/unified_http_client.py +10 -2
  12. databricks_sql_connector-4.2.5/src/databricks/sql/common/url_utils.py +45 -0
  13. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/telemetry/models/event.py +2 -0
  14. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/telemetry/telemetry_client.py +42 -5
  15. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/utils.py +15 -0
  16. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/LICENSE +0 -0
  17. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/README.md +0 -0
  18. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/__init__.py +0 -0
  19. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/auth/__init__.py +0 -0
  20. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/auth/auth.py +0 -0
  21. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/auth/authenticators.py +0 -0
  22. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/auth/common.py +0 -0
  23. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/auth/endpoint.py +0 -0
  24. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/auth/oauth.py +0 -0
  25. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/auth/oauth_http_handler.py +0 -0
  26. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/auth/thrift_http_client.py +0 -0
  27. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/databricks_client.py +0 -0
  28. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/sea/backend.py +0 -0
  29. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/sea/models/__init__.py +0 -0
  30. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/sea/models/base.py +0 -0
  31. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/sea/models/requests.py +0 -0
  32. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/sea/models/responses.py +0 -0
  33. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/sea/queue.py +0 -0
  34. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/sea/result_set.py +0 -0
  35. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/sea/utils/constants.py +0 -0
  36. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/sea/utils/conversion.py +0 -0
  37. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/sea/utils/filters.py +0 -0
  38. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/sea/utils/normalize.py +0 -0
  39. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/thrift_backend.py +0 -0
  40. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/types.py +0 -0
  41. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/utils/__init__.py +0 -0
  42. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/backend/utils/guid_utils.py +0 -0
  43. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/cloudfetch/download_manager.py +0 -0
  44. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/cloudfetch/downloader.py +0 -0
  45. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/common/http.py +0 -0
  46. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/common/http_utils.py +0 -0
  47. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/exc.py +0 -0
  48. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/experimental/__init__.py +0 -0
  49. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/experimental/oauth_persistence.py +0 -0
  50. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/parameters/__init__.py +0 -0
  51. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/parameters/native.py +0 -0
  52. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/parameters/py.typed +0 -0
  53. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/py.typed +0 -0
  54. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/result_set.py +0 -0
  55. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/session.py +0 -0
  56. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/telemetry/circuit_breaker_manager.py +0 -0
  57. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/telemetry/latency_logger.py +0 -0
  58. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/telemetry/models/endpoint_models.py +0 -0
  59. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/telemetry/models/enums.py +0 -0
  60. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/telemetry/models/frontend_logs.py +0 -0
  61. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/telemetry/telemetry_push_client.py +0 -0
  62. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/telemetry/utils.py +0 -0
  63. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/thrift_api/TCLIService/TCLIService-remote +0 -0
  64. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/thrift_api/TCLIService/TCLIService.py +0 -0
  65. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/thrift_api/TCLIService/__init__.py +0 -0
  66. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/thrift_api/TCLIService/constants.py +0 -0
  67. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/thrift_api/TCLIService/ttypes.py +0 -0
  68. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/thrift_api/__init__.py +0 -0
  69. {databricks_sql_connector-4.2.3 → databricks_sql_connector-4.2.5}/src/databricks/sql/types.py +0 -0
@@ -1,5 +1,13 @@
1
1
  # Release History
2
2
 
3
+ # 4.2.5 (2026-02-09)
4
+ - Fix feature-flag endpoint retries in gov region (databricks/databricks-sql-python#735 by @samikshya-db)
5
+ - Improve telemetry lifecycle management (databricks/databricks-sql-python#734 by @msrathore-db)
6
+
7
+ # 4.2.4 (2026-01-07)
8
+ - Fixed the exception handler close() on _TelemetryClientHolder (databricks/databricks-sql-python#723 by @msrathore-db)
9
+ - Created util method to normalise http protocol in http path (databricks/databricks-sql-python#724 by @nikhilsuri-db)
10
+
3
11
  # 4.2.3 (2025-12-18)
4
12
  - added pandas < 2.4.0 support and tests for py 3.14 (databricks/databricks-sql-python#720 by @sreekanth-db)
5
13
  - pandas 2.3.3 support for py < 3.14 (databricks/databricks-sql-python#721 by @sreekanth-db)
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: databricks-sql-connector
3
- Version: 4.2.3
3
+ Version: 4.2.5
4
4
  Summary: Databricks SQL Connector for Python
5
5
  License: Apache-2.0
6
6
  License-File: LICENSE
@@ -1,6 +1,6 @@
1
1
  [tool.poetry]
2
2
  name = "databricks-sql-connector"
3
- version = "4.2.3"
3
+ version = "4.2.5"
4
4
  description = "Databricks SQL Connector for Python"
5
5
  authors = ["Databricks <databricks-sql-connector-maintainers@databricks.com>"]
6
6
  license = "Apache-2.0"
@@ -92,4 +92,4 @@ show_missing = true
92
92
  skip_covered = false
93
93
 
94
94
  [tool.coverage.xml]
95
- output = "coverage.xml"
95
+ output = "coverage.xml"
@@ -71,7 +71,7 @@ DATETIME = DBAPITypeObject("timestamp")
71
71
  DATE = DBAPITypeObject("date")
72
72
  ROWID = DBAPITypeObject()
73
73
 
74
- __version__ = "4.2.3"
74
+ __version__ = "4.2.5"
75
75
  USER_AGENT_NAME = "PyDatabricksSqlConnector"
76
76
 
77
77
  # These two functions are pyhive legacy
@@ -7,23 +7,6 @@ from urllib.parse import urlparse
7
7
  logger = logging.getLogger(__name__)
8
8
 
9
9
 
10
- def parse_hostname(hostname: str) -> str:
11
- """
12
- Normalize the hostname to include scheme and trailing slash.
13
-
14
- Args:
15
- hostname: The hostname to normalize
16
-
17
- Returns:
18
- Normalized hostname with scheme and trailing slash
19
- """
20
- if not hostname.startswith("http://") and not hostname.startswith("https://"):
21
- hostname = f"https://{hostname}"
22
- if not hostname.endswith("/"):
23
- hostname = f"{hostname}/"
24
- return hostname
25
-
26
-
27
10
  def decode_token(access_token: str) -> Optional[Dict]:
28
11
  """
29
12
  Decode a JWT token without verification to extract claims.
@@ -373,6 +373,13 @@ class DatabricksRetryPolicy(Retry):
373
373
  if status_code == 403:
374
374
  return False, "403 codes are not retried"
375
375
 
376
+ # Request failed with 404. Don't retry for any command type.
377
+ if status_code == 404:
378
+ return (
379
+ False,
380
+ "Received 404 - NOT_FOUND. The requested resource does not exist.",
381
+ )
382
+
376
383
  # Request failed and server said NotImplemented. This isn't recoverable. Don't retry.
377
384
  if status_code == 501:
378
385
  return False, "Received code 501 from server."
@@ -381,33 +388,6 @@ class DatabricksRetryPolicy(Retry):
381
388
  if not self._is_method_retryable(method):
382
389
  return False, "Only POST requests are retried"
383
390
 
384
- # Request failed with 404 and was a GetOperationStatus. This is not recoverable. Don't retry.
385
- if status_code == 404 and self.command_type == CommandType.GET_OPERATION_STATUS:
386
- return (
387
- False,
388
- "GetOperationStatus received 404 code from Databricks. Operation was canceled.",
389
- )
390
-
391
- # Request failed with 404 because CloseSession returns 404 if you repeat the request.
392
- if (
393
- status_code == 404
394
- and self.command_type == CommandType.CLOSE_SESSION
395
- and len(self.history) > 0
396
- ):
397
- raise SessionAlreadyClosedError(
398
- "CloseSession received 404 code from Databricks. Session is already closed."
399
- )
400
-
401
- # Request failed with 404 because CloseOperation returns 404 if you repeat the request.
402
- if (
403
- status_code == 404
404
- and self.command_type == CommandType.CLOSE_OPERATION
405
- and len(self.history) > 0
406
- ):
407
- raise CursorAlreadyClosedError(
408
- "CloseOperation received 404 code from Databricks. Cursor is already closed."
409
- )
410
-
411
391
  # Request failed, was an ExecuteStatement and the command may have reached the server
412
392
  if (
413
393
  self.command_type == CommandType.EXECUTE_STATEMENT
@@ -6,10 +6,10 @@ from urllib.parse import urlencode
6
6
 
7
7
  from databricks.sql.auth.authenticators import AuthProvider
8
8
  from databricks.sql.auth.auth_utils import (
9
- parse_hostname,
10
9
  decode_token,
11
10
  is_same_host,
12
11
  )
12
+ from databricks.sql.common.url_utils import normalize_host_with_protocol
13
13
  from databricks.sql.common.http import HttpMethod
14
14
 
15
15
  logger = logging.getLogger(__name__)
@@ -99,7 +99,7 @@ class TokenFederationProvider(AuthProvider):
99
99
  if not http_client:
100
100
  raise ValueError("http_client is required for TokenFederationProvider")
101
101
 
102
- self.hostname = parse_hostname(hostname)
102
+ self.hostname = normalize_host_with_protocol(hostname)
103
103
  self.external_provider = external_provider
104
104
  self.http_client = http_client
105
105
  self.identity_federation_client_id = identity_federation_client_id
@@ -164,7 +164,7 @@ class TokenFederationProvider(AuthProvider):
164
164
 
165
165
  def _exchange_token(self, access_token: str) -> Token:
166
166
  """Exchange the external token for a Databricks token."""
167
- token_url = f"{self.hostname.rstrip('/')}{self.TOKEN_EXCHANGE_ENDPOINT}"
167
+ token_url = f"{self.hostname}{self.TOKEN_EXCHANGE_ENDPOINT}"
168
168
 
169
169
  data = {
170
170
  "grant_type": self.TOKEN_EXCHANGE_GRANT_TYPE,
@@ -18,6 +18,7 @@ from databricks.sql.exc import (
18
18
  from databricks.sql.common.http_utils import (
19
19
  detect_and_parse_proxy,
20
20
  )
21
+ from databricks.sql.common.url_utils import normalize_host_with_protocol
21
22
 
22
23
  logger = logging.getLogger(__name__)
23
24
 
@@ -66,8 +67,9 @@ class SeaHttpClient:
66
67
  self.auth_provider = auth_provider
67
68
  self.ssl_options = ssl_options
68
69
 
69
- # Build base URL
70
- self.base_url = f"https://{server_hostname}:{self.port}"
70
+ # Build base URL using url_utils for consistent normalization
71
+ normalized_host = normalize_host_with_protocol(server_hostname)
72
+ self.base_url = f"{normalized_host}:{self.port}"
71
73
 
72
74
  # Parse URL for proxy handling
73
75
  parsed_url = urllib.parse.urlparse(self.base_url)
@@ -35,6 +35,7 @@ from databricks.sql.utils import (
35
35
  ColumnTable,
36
36
  ColumnQueue,
37
37
  build_client_context,
38
+ get_session_config_value,
38
39
  )
39
40
  from databricks.sql.parameters.native import (
40
41
  DbsqlParameterBase,
@@ -305,6 +306,8 @@ class Connection:
305
306
  )
306
307
  self.session.open()
307
308
  except Exception as e:
309
+ # Respect user's telemetry preference even during connection failure
310
+ enable_telemetry = kwargs.get("enable_telemetry", True)
308
311
  TelemetryClientFactory.connection_failure_log(
309
312
  error_name="Exception",
310
313
  error_message=str(e),
@@ -315,6 +318,7 @@ class Connection:
315
318
  user_agent=self.session.useragent_header
316
319
  if hasattr(self, "session")
317
320
  else None,
321
+ enable_telemetry=enable_telemetry,
318
322
  )
319
323
  raise e
320
324
 
@@ -386,6 +390,7 @@ class Connection:
386
390
  support_many_parameters=True, # Native parameters supported
387
391
  enable_complex_datatype_support=_use_arrow_native_complex_types,
388
392
  allowed_volume_ingestion_paths=self.staging_allowed_local_path,
393
+ query_tags=get_session_config_value(session_configuration, "query_tags"),
389
394
  )
390
395
 
391
396
  self._telemetry_client.export_initial_telemetry_log(
@@ -6,6 +6,7 @@ from concurrent.futures import ThreadPoolExecutor
6
6
  from typing import Dict, Optional, List, Any, TYPE_CHECKING
7
7
 
8
8
  from databricks.sql.common.http import HttpMethod
9
+ from databricks.sql.common.url_utils import normalize_host_with_protocol
9
10
 
10
11
  if TYPE_CHECKING:
11
12
  from databricks.sql.client import Connection
@@ -67,7 +68,8 @@ class FeatureFlagsContext:
67
68
 
68
69
  endpoint_suffix = FEATURE_FLAGS_ENDPOINT_SUFFIX_FORMAT.format(__version__)
69
70
  self._feature_flag_endpoint = (
70
- f"https://{self._connection.session.host}{endpoint_suffix}"
71
+ normalize_host_with_protocol(self._connection.session.host)
72
+ + endpoint_suffix
71
73
  )
72
74
 
73
75
  # Use the provided HTTP client
@@ -217,7 +217,7 @@ class UnifiedHttpClient:
217
217
  logger.debug("Error checking proxy bypass for host %s: %s", target_host, e)
218
218
  return True
219
219
 
220
- def _get_pool_manager_for_url(self, url: str) -> urllib3.PoolManager:
220
+ def _get_pool_manager_for_url(self, url: str) -> Optional[urllib3.PoolManager]:
221
221
  """
222
222
  Get the appropriate pool manager for the given URL.
223
223
 
@@ -225,7 +225,7 @@ class UnifiedHttpClient:
225
225
  url: The target URL
226
226
 
227
227
  Returns:
228
- PoolManager instance (either direct or proxy)
228
+ PoolManager instance (either direct or proxy), or None if client is closed
229
229
  """
230
230
  parsed_url = urllib.parse.urlparse(url)
231
231
  target_host = parsed_url.hostname
@@ -291,6 +291,14 @@ class UnifiedHttpClient:
291
291
  # Select appropriate pool manager based on target URL
292
292
  pool_manager = self._get_pool_manager_for_url(url)
293
293
 
294
+ # DEFENSIVE: Check if pool_manager is None (client closing/closed)
295
+ # This prevents AttributeError race condition when telemetry cleanup happens
296
+ if pool_manager is None:
297
+ logger.debug(
298
+ "HTTP client closing or closed, cannot make request to %s", url
299
+ )
300
+ raise RequestError("HTTP client is closing or has been closed")
301
+
294
302
  response = None
295
303
 
296
304
  try:
@@ -0,0 +1,45 @@
1
+ """
2
+ URL utility functions for the Databricks SQL connector.
3
+ """
4
+
5
+
6
+ def normalize_host_with_protocol(host: str) -> str:
7
+ """
8
+ Normalize a connection hostname by ensuring it has a protocol.
9
+
10
+ This is useful for handling cases where users may provide hostnames with or without protocols
11
+ (common with dbt-databricks users copying URLs from their browser).
12
+
13
+ Args:
14
+ host: Connection hostname which may or may not include a protocol prefix (https:// or http://)
15
+ and may or may not have a trailing slash
16
+
17
+ Returns:
18
+ Normalized hostname with protocol prefix and no trailing slashes
19
+
20
+ Examples:
21
+ normalize_host_with_protocol("myserver.com") -> "https://myserver.com"
22
+ normalize_host_with_protocol("https://myserver.com") -> "https://myserver.com"
23
+ normalize_host_with_protocol("HTTPS://myserver.com/") -> "https://myserver.com"
24
+ normalize_host_with_protocol("http://localhost:8080/") -> "http://localhost:8080"
25
+
26
+ Raises:
27
+ ValueError: If host is None or empty string
28
+ """
29
+ # Handle None or empty host
30
+ if not host or not host.strip():
31
+ raise ValueError("Host cannot be None or empty")
32
+
33
+ # Remove trailing slashes
34
+ host = host.rstrip("/")
35
+
36
+ # Add protocol if not present (case-insensitive check)
37
+ host_lower = host.lower()
38
+ if not host_lower.startswith("https://") and not host_lower.startswith("http://"):
39
+ host = f"https://{host}"
40
+ elif host_lower.startswith("https://") or host_lower.startswith("http://"):
41
+ # Normalize protocol to lowercase
42
+ protocol_end = host.index("://") + 3
43
+ host = host[:protocol_end].lower() + host[protocol_end:]
44
+
45
+ return host
@@ -57,6 +57,7 @@ class DriverConnectionParameters(JsonSerializableMixin):
57
57
  support_many_parameters (bool): Whether many parameters are supported
58
58
  enable_complex_datatype_support (bool): Whether complex datatypes are supported
59
59
  allowed_volume_ingestion_paths (str): Allowed paths for volume ingestion
60
+ query_tags (str): Query tags for tracking and attribution
60
61
  """
61
62
 
62
63
  http_path: str
@@ -84,6 +85,7 @@ class DriverConnectionParameters(JsonSerializableMixin):
84
85
  support_many_parameters: Optional[bool] = None
85
86
  enable_complex_datatype_support: Optional[bool] = None
86
87
  allowed_volume_ingestion_paths: Optional[str] = None
88
+ query_tags: Optional[str] = None
87
89
 
88
90
 
89
91
  @dataclass
@@ -42,11 +42,13 @@ from databricks.sql.telemetry.utils import BaseTelemetryClient
42
42
  from databricks.sql.common.feature_flag import FeatureFlagsContextFactory
43
43
  from databricks.sql.common.unified_http_client import UnifiedHttpClient
44
44
  from databricks.sql.common.http import HttpMethod
45
+ from databricks.sql.exc import RequestError
45
46
  from databricks.sql.telemetry.telemetry_push_client import (
46
47
  ITelemetryPushClient,
47
48
  TelemetryPushClient,
48
49
  CircuitBreakerTelemetryPushClient,
49
50
  )
51
+ from databricks.sql.common.url_utils import normalize_host_with_protocol
50
52
 
51
53
  if TYPE_CHECKING:
52
54
  from databricks.sql.client import Connection
@@ -278,7 +280,7 @@ class TelemetryClient(BaseTelemetryClient):
278
280
  if self._auth_provider
279
281
  else self.TELEMETRY_UNAUTHENTICATED_PATH
280
282
  )
281
- url = f"https://{self._host_url}{path}"
283
+ url = normalize_host_with_protocol(self._host_url) + path
282
284
 
283
285
  headers = {"Accept": "application/json", "Content-Type": "application/json"}
284
286
 
@@ -416,10 +418,38 @@ class TelemetryClient(BaseTelemetryClient):
416
418
  )
417
419
 
418
420
  def close(self):
419
- """Flush remaining events before closing"""
421
+ """Flush remaining events before closing
422
+
423
+ IMPORTANT: This method does NOT close self._http_client.
424
+
425
+ Rationale:
426
+ - _flush() submits async work to the executor that uses _http_client
427
+ - If we closed _http_client here, async callbacks would fail with AttributeError
428
+ - Instead, we let _http_client live as long as needed:
429
+ * Pending futures hold references to self (via bound methods)
430
+ * This keeps self alive, which keeps self._http_client alive
431
+ * When all futures complete, Python GC will clean up naturally
432
+ - The __del__ method ensures eventual cleanup during garbage collection
433
+
434
+ This design prevents race conditions while keeping telemetry truly async.
435
+ """
420
436
  logger.debug("Closing TelemetryClient for connection %s", self._session_id_hex)
421
437
  self._flush()
422
438
 
439
+ def __del__(self):
440
+ """Cleanup when TelemetryClient is garbage collected
441
+
442
+ This ensures _http_client is eventually closed when the TelemetryClient
443
+ object is destroyed. By this point, all async work should be complete
444
+ (since the futures held references keeping us alive), so it's safe to
445
+ close the http client.
446
+ """
447
+ try:
448
+ if hasattr(self, "_http_client") and self._http_client:
449
+ self._http_client.close()
450
+ except Exception:
451
+ pass
452
+
423
453
 
424
454
  class _TelemetryClientHolder:
425
455
  """
@@ -542,8 +572,8 @@ class TelemetryClientFactory:
542
572
  logger.debug("Handling unhandled exception: %s", exc_type.__name__)
543
573
 
544
574
  clients_to_close = list(cls._clients.values())
545
- for client in clients_to_close:
546
- client.close()
575
+ for holder in clients_to_close:
576
+ holder.client.close()
547
577
 
548
578
  # Call the original exception handler to maintain normal behavior
549
579
  if cls._original_excepthook:
@@ -673,7 +703,8 @@ class TelemetryClientFactory:
673
703
  )
674
704
  try:
675
705
  TelemetryClientFactory._stop_flush_thread()
676
- TelemetryClientFactory._executor.shutdown(wait=True)
706
+ # Use wait=False to allow process to exit immediately
707
+ TelemetryClientFactory._executor.shutdown(wait=False)
677
708
  except Exception as e:
678
709
  logger.debug("Failed to shutdown thread pool executor: %s", e)
679
710
  TelemetryClientFactory._executor = None
@@ -688,9 +719,15 @@ class TelemetryClientFactory:
688
719
  port: int,
689
720
  client_context,
690
721
  user_agent: Optional[str] = None,
722
+ enable_telemetry: bool = True,
691
723
  ):
692
724
  """Send error telemetry when connection creation fails, using provided client context"""
693
725
 
726
+ # Respect user's telemetry preference - don't force-enable
727
+ if not enable_telemetry:
728
+ logger.debug("Telemetry disabled, skipping connection failure log")
729
+ return
730
+
694
731
  UNAUTH_DUMMY_SESSION_ID = "unauth_session_id"
695
732
 
696
733
  TelemetryClientFactory.initialize_telemetry_client(
@@ -38,6 +38,21 @@ DEFAULT_ERROR_CONTEXT = "Unknown error"
38
38
  logger = logging.getLogger(__name__)
39
39
 
40
40
 
41
+ def get_session_config_value(
42
+ session_configuration: Optional[Dict[str, Any]], key: str
43
+ ) -> Optional[str]:
44
+ """Get a session configuration value with case-insensitive key matching"""
45
+ if not session_configuration:
46
+ return None
47
+
48
+ key_upper = key.upper()
49
+ for k, v in session_configuration.items():
50
+ if k.upper() == key_upper:
51
+ return str(v) if v is not None else None
52
+
53
+ return None
54
+
55
+
41
56
  class ResultSetQueue(ABC):
42
57
  @abstractmethod
43
58
  def next_n_rows(self, num_rows: int):