databricks-sql-connector 4.2.4__tar.gz → 4.2.6__tar.gz

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/CHANGELOG.md

@@ -1,5 +1,20 @@
 # Release History
 
+# 4.2.6 (2026-04-22)
+- Add SPOG routing support for account-level vanity URLs (databricks/databricks-sql-python#767 by @msrathore-db)
+- Fix dependency_manager: handle PEP 440 ~= compatible release syntax (databricks/databricks-sql-python#776 by @vikrantpuppala)
+- Bump thrift to fix deprecation warning (databricks/databricks-sql-python#733 by @Korijn)
+- Add AI coding agent detection to User-Agent header (databricks/databricks-sql-python#740 by @vikrantpuppala)
+- Add statement-level query_tags support for SEA backend (databricks/databricks-sql-python#754 by @sreekanth-db)
+- Update PyArrow concatenation of tables to use promote_options as default (databricks/databricks-sql-python#751 by @jprakash-db)
+- Fix float inference to use DoubleParameter (64-bit) instead of FloatParameter (databricks/databricks-sql-python#742 by @Shubhambhusate)
+- Allow specifying query_tags as a dict upon connection creation (databricks/databricks-sql-python#749 by @jiabin-hu)
+- Add query_tags parameter support for execute methods (databricks/databricks-sql-python#736 by @jiabin-hu)
+
+# 4.2.5 (2026-02-09)
+- Fix feature-flag endpoint retries in gov region (databricks/databricks-sql-python#735 by @samikshya-db)
+- Improve telemetry lifecycle management (databricks/databricks-sql-python#734 by @msrathore-db)
+
 # 4.2.4 (2026-01-07)
 - Fixed the exception handler close() on _TelemetryClientHolder (databricks/databricks-sql-python#723 by @msrathore-db)
 - Created util method to normalise http protocol in http path (databricks/databricks-sql-python#724 by @nikhilsuri-db)

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: databricks-sql-connector
-Version: 4.2.4
+Version: 4.2.6
 Summary: Databricks SQL Connector for Python
 License: Apache-2.0
 License-File: LICENSE
@@ -30,7 +30,7 @@ Requires-Dist: pybreaker (>=1.0.0,<2.0.0)
 Requires-Dist: pyjwt (>=2.0.0,<3.0.0)
 Requires-Dist: python-dateutil (>=2.8.0,<3.0.0)
 Requires-Dist: requests (>=2.18.1,<3.0.0)
-Requires-Dist: thrift (>=0.16.0,<0.21.0)
+Requires-Dist: thrift (>=0.22.0,<0.23.0)
 Requires-Dist: urllib3 (>=1.26)
 Project-URL: Bug Tracker, https://github.com/databricks/databricks-sql-python/issues
 Project-URL: Homepage, https://github.com/databricks/databricks-sql-python

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databricks-sql-connector"
-version = "4.2.4"
+version = "4.2.6"
 description = "Databricks SQL Connector for Python"
 authors = ["Databricks <databricks-sql-connector-maintainers@databricks.com>"]
 license = "Apache-2.0"
@@ -10,7 +10,7 @@ include = ["CHANGELOG.md"]
 
 [tool.poetry.dependencies]
 python = "^3.8.0"
-thrift = ">=0.16.0,<0.21.0"
+thrift = "~=0.22.0"
 pandas = [
     { version = ">=1.2.5,<2.4.0", python = ">=3.8,<3.13" },
     { version = ">=2.2.3,<2.4.0", python = ">=3.13" }
@@ -92,4 +92,4 @@ show_missing = true
 skip_covered = false
 
 [tool.coverage.xml]
-output = "coverage.xml"
+output = "coverage.xml"

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/src/databricks/sql/__init__.py

@@ -71,7 +71,7 @@ DATETIME = DBAPITypeObject("timestamp")
 DATE = DBAPITypeObject("date")
 ROWID = DBAPITypeObject()
 
-__version__ = "4.2.4"
+__version__ = "4.2.6"
 USER_AGENT_NAME = "PyDatabricksSqlConnector"
 
 # These two functions are pyhive legacy

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/src/databricks/sql/auth/retry.py

@@ -373,6 +373,13 @@ class DatabricksRetryPolicy(Retry):
         if status_code == 403:
             return False, "403 codes are not retried"
 
+        # Request failed with 404. Don't retry for any command type.
+        if status_code == 404:
+            return (
+                False,
+                "Received 404 - NOT_FOUND. The requested resource does not exist.",
+            )
+
         # Request failed and server said NotImplemented. This isn't recoverable. Don't retry.
         if status_code == 501:
             return False, "Received code 501 from server."
@@ -381,33 +388,6 @@ class DatabricksRetryPolicy(Retry):
         if not self._is_method_retryable(method):
             return False, "Only POST requests are retried"
 
-        # Request failed with 404 and was a GetOperationStatus. This is not recoverable. Don't retry.
-        if status_code == 404 and self.command_type == CommandType.GET_OPERATION_STATUS:
-            return (
-                False,
-                "GetOperationStatus received 404 code from Databricks. Operation was canceled.",
-            )
-
-        # Request failed with 404 because CloseSession returns 404 if you repeat the request.
-        if (
-            status_code == 404
-            and self.command_type == CommandType.CLOSE_SESSION
-            and len(self.history) > 0
-        ):
-            raise SessionAlreadyClosedError(
-                "CloseSession received 404 code from Databricks. Session is already closed."
-            )
-
-        # Request failed with 404 because CloseOperation returns 404 if you repeat the request.
-        if (
-            status_code == 404
-            and self.command_type == CommandType.CLOSE_OPERATION
-            and len(self.history) > 0
-        ):
-            raise CursorAlreadyClosedError(
-                "CloseOperation received 404 code from Databricks. Cursor is already closed."
-            )
-
         # Request failed, was an ExecuteStatement and the command may have reached the server
         if (
             self.command_type == CommandType.EXECUTE_STATEMENT

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/src/databricks/sql/backend/databricks_client.py

@@ -83,6 +83,7 @@ class DatabricksClient(ABC):
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
         row_limit: Optional[int] = None,
+        query_tags: Optional[Dict[str, Optional[str]]] = None,
     ) -> Union[ResultSet, None]:
         """
         Executes a SQL command or query within the specified session.
@@ -102,6 +103,7 @@ class DatabricksClient(ABC):
             async_op: Whether to execute the command asynchronously
             enforce_embedded_schema_correctness: Whether to enforce schema correctness
             row_limit: Maximum number of rows in the response.
+            query_tags: Optional dictionary of query tags to apply for this query only.
 
         Returns:
             If async_op is False, returns a ResultSet object containing the

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/src/databricks/sql/backend/sea/backend.py

@@ -188,8 +188,9 @@ class SeaDatabricksClient(DatabricksClient):
             ValueError: If the warehouse ID cannot be extracted from the path
         """
 
-        warehouse_pattern = re.compile(r".*/warehouses/(.+)")
-        endpoint_pattern = re.compile(r".*/endpoints/(.+)")
+        # [^?&]+ stops at query params (e.g. ?o= for SPOG routing)
+        warehouse_pattern = re.compile(r".*/warehouses/([^?&]+)")
+        endpoint_pattern = re.compile(r".*/endpoints/([^?&]+)")
 
         for pattern in [warehouse_pattern, endpoint_pattern]:
             match = pattern.match(http_path)
@@ -463,6 +464,7 @@ class SeaDatabricksClient(DatabricksClient):
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
         row_limit: Optional[int] = None,
+        query_tags: Optional[Dict[str, Optional[str]]] = None,
     ) -> Union[SeaResultSet, None]:
         """
         Execute a SQL command using the SEA backend.
@@ -529,6 +531,7 @@ class SeaDatabricksClient(DatabricksClient):
             row_limit=row_limit,
             parameters=sea_parameters if sea_parameters else None,
             result_compression=result_compression,
+            query_tags=query_tags,
         )
 
         response_data = self._http_client._make_request(

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/src/databricks/sql/backend/sea/models/requests.py

@@ -31,6 +31,7 @@ class ExecuteStatementRequest:
     wait_timeout: str = "10s"
     on_wait_timeout: str = "CONTINUE"
     row_limit: Optional[int] = None
+    query_tags: Optional[Dict[str, Optional[str]]] = None
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert the request to a dictionary for JSON serialization."""
@@ -60,6 +61,13 @@ class ExecuteStatementRequest:
                 for param in self.parameters
             ]
 
+        # SEA API expects query_tags as an array of {key, value} objects.
+        # None/empty values are left to the server to handle as key-only tags.
+        if self.query_tags:
+            result["query_tags"] = [
+                {"key": k, "value": v} for k, v in self.query_tags.items()
+            ]
+
         return result
 
 

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/src/databricks/sql/backend/thrift_backend.py

@@ -5,7 +5,7 @@ import logging
 import math
 import time
 import threading
-from typing import List, Optional, Union, Any, TYPE_CHECKING
+from typing import Dict, List, Optional, Union, Any, TYPE_CHECKING
 from uuid import UUID
 
 from databricks.sql.common.unified_http_client import UnifiedHttpClient
@@ -53,6 +53,7 @@ from databricks.sql.utils import (
     convert_arrow_based_set_to_arrow_table,
     convert_decimals_in_arrow_table,
     convert_column_based_set_to_arrow_table,
+    serialize_query_tags,
 )
 from databricks.sql.types import SSLOptions
 from databricks.sql.backend.databricks_client import DatabricksClient
@@ -1003,6 +1004,7 @@ class ThriftDatabricksClient(DatabricksClient):
         async_op=False,
         enforce_embedded_schema_correctness=False,
         row_limit: Optional[int] = None,
+        query_tags: Optional[Dict[str, Optional[str]]] = None,
     ) -> Union["ResultSet", None]:
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
@@ -1022,6 +1024,19 @@ class ThriftDatabricksClient(DatabricksClient):
             # DBR should be changed to use month_day_nano_interval
             intervalTypesAsArrow=False,
         )
+
+        # Build confOverlay with default configs and query_tags
+        merged_conf_overlay = {
+            # We want to receive proper Timestamp arrow types.
+            "spark.thriftserver.arrowBasedRowSet.timestampAsString": "false"
+        }
+
+        # Serialize and add query_tags to confOverlay if provided
+        if query_tags:
+            serialized_tags = serialize_query_tags(query_tags)
+            if serialized_tags:
+                merged_conf_overlay["query_tags"] = serialized_tags
+
         req = ttypes.TExecuteStatementReq(
             sessionHandle=thrift_handle,
             statement=operation,
@@ -1036,10 +1051,7 @@ class ThriftDatabricksClient(DatabricksClient):
             canReadArrowResult=True if pyarrow else False,
             canDecompressLZ4Result=lz4_compression,
             canDownloadResult=use_cloud_fetch,
-            confOverlay={
-                # We want to receive proper Timestamp arrow types.
-                "spark.thriftserver.arrowBasedRowSet.timestampAsString": "false"
-            },
+            confOverlay=merged_conf_overlay,
             useArrowNativeTypes=spark_arrow_types,
             parameters=parameters,
             enforceEmbeddedSchemaCorrectness=enforce_embedded_schema_correctness,

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/src/databricks/sql/client.py

@@ -35,6 +35,8 @@ from databricks.sql.utils import (
     ColumnTable,
     ColumnQueue,
     build_client_context,
+    get_session_config_value,
+    serialize_query_tags,
 )
 from databricks.sql.parameters.native import (
     DbsqlParameterBase,
@@ -105,6 +107,7 @@ class Connection:
         schema: Optional[str] = None,
         _use_arrow_native_complex_types: Optional[bool] = True,
         ignore_transactions: bool = True,
+        query_tags: Optional[Dict[str, Optional[str]]] = None,
         **kwargs,
     ) -> None:
         """
@@ -280,6 +283,15 @@ class Connection:
                 "spark.sql.thriftserver.metadata.metricview.enabled"
             ] = "true"
 
+        if query_tags is not None:
+            if session_configuration is None:
+                session_configuration = {}
+            serialized = serialize_query_tags(query_tags)
+            if serialized:
+                session_configuration["QUERY_TAGS"] = serialized
+            else:
+                session_configuration.pop("QUERY_TAGS", None)
+
         self.disable_pandas = kwargs.get("_disable_pandas", False)
         self.lz4_compression = kwargs.get("enable_query_result_lz4_compression", True)
         self.use_cloud_fetch = kwargs.get("use_cloud_fetch", True)
@@ -305,6 +317,8 @@ class Connection:
             )
             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),
@@ -315,6 +329,7 @@ class Connection:
                 user_agent=self.session.useragent_header
                 if hasattr(self, "session")
                 else None,
+                enable_telemetry=enable_telemetry,
             )
             raise e
 
@@ -338,6 +353,7 @@ class Connection:
             host_url=self.session.host,
             batch_size=self.telemetry_batch_size,
             client_context=client_context,
+            extra_headers=self.session.get_spog_headers(),
         )
 
         self._telemetry_client = TelemetryClientFactory.get_telemetry_client(
@@ -386,6 +402,7 @@ class Connection:
             support_many_parameters=True,  # Native parameters supported
             enable_complex_datatype_support=_use_arrow_native_complex_types,
             allowed_volume_ingestion_paths=self.staging_allowed_local_path,
+            query_tags=get_session_config_value(session_configuration, "query_tags"),
         )
 
         self._telemetry_client.export_initial_telemetry_log(
@@ -1258,6 +1275,7 @@ class Cursor:
         parameters: Optional[TParameterCollection] = None,
         enforce_embedded_schema_correctness=False,
         input_stream: Optional[BinaryIO] = None,
+        query_tags: Optional[Dict[str, Optional[str]]] = None,
     ) -> "Cursor":
         """
         Execute a query and wait for execution to complete.
@@ -1288,6 +1306,10 @@ class Cursor:
         Both will result in the query equivalent to "SELECT * FROM table WHERE field = 'foo'
         being sent to the server
 
+        :param query_tags: Optional dictionary of query tags to apply for this query only.
+            Tags are key-value pairs that can be used to identify and categorize queries.
+            Example: {"team": "data-eng", "application": "etl"}
+
         :returns self
         """
 
@@ -1328,6 +1350,7 @@ class Cursor:
             async_op=False,
             enforce_embedded_schema_correctness=enforce_embedded_schema_correctness,
             row_limit=self.row_limit,
+            query_tags=query_tags,
         )
 
         if self.active_result_set and self.active_result_set.is_staging_operation:
@@ -1344,6 +1367,7 @@ class Cursor:
         operation: str,
         parameters: Optional[TParameterCollection] = None,
         enforce_embedded_schema_correctness=False,
+        query_tags: Optional[Dict[str, Optional[str]]] = None,
     ) -> "Cursor":
         """
 
@@ -1351,6 +1375,9 @@ class Cursor:
 
         :param operation:
         :param parameters:
+        :param query_tags: Optional dictionary of query tags to apply for this query only.
+            Tags are key-value pairs that can be used to identify and categorize queries.
+            Example: {"team": "data-eng", "application": "etl"}
         :return:
         """
 
@@ -1387,6 +1414,7 @@ class Cursor:
             async_op=True,
             enforce_embedded_schema_correctness=enforce_embedded_schema_correctness,
             row_limit=self.row_limit,
+            query_tags=query_tags,
         )
 
         return self
@@ -1443,7 +1471,12 @@ class Cursor:
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
-    def executemany(self, operation, seq_of_parameters):
+    def executemany(
+        self,
+        operation,
+        seq_of_parameters,
+        query_tags: Optional[Dict[str, Optional[str]]] = None,
+    ):
         """
         Execute the operation once for every set of passed in parameters.
 
@@ -1452,10 +1485,14 @@ class Cursor:
 
         Only the final result set is retained.
 
+        :param query_tags: Optional dictionary of query tags to apply for all queries in this batch.
+            Tags are key-value pairs that can be used to identify and categorize queries.
+            Example: {"team": "data-eng", "application": "etl"}
+
         :returns self
         """
         for parameters in seq_of_parameters:
-            self.execute(operation, parameters)
+            self.execute(operation, parameters, query_tags=query_tags)
         return self
 
     @log_latency(StatementType.METADATA)

databricks_sql_connector-4.2.6/src/databricks/sql/common/agent.py

@@ -0,0 +1,52 @@
+"""
+Detects whether the Python SQL connector is being invoked by an AI coding agent
+by checking for well-known environment variables that agents set in their spawned
+shell processes.
+
+Detection only succeeds when exactly one agent environment variable is present,
+to avoid ambiguous attribution when multiple agent environments overlap.
+
+Adding a new agent requires only a new entry in KNOWN_AGENTS.
+
+References for each environment variable:
+  - ANTIGRAVITY_AGENT: Closed source. Google Antigravity sets this variable.
+  - CLAUDECODE: https://github.com/anthropics/claude-code (sets CLAUDECODE=1)
+  - CLINE_ACTIVE: https://github.com/cline/cline (shipped in v3.24.0)
+  - CODEX_CI: https://github.com/openai/codex (part of UNIFIED_EXEC_ENV array in codex-rs)
+  - CURSOR_AGENT: Closed source. Referenced in a gist by johnlindquist.
+  - GEMINI_CLI: https://google-gemini.github.io/gemini-cli/docs/tools/shell.html (sets GEMINI_CLI=1)
+  - OPENCODE: https://github.com/opencode-ai/opencode (sets OPENCODE=1)
+"""
+
+import os
+
+KNOWN_AGENTS = [
+    ("ANTIGRAVITY_AGENT", "antigravity"),
+    ("CLAUDECODE", "claude-code"),
+    ("CLINE_ACTIVE", "cline"),
+    ("CODEX_CI", "codex"),
+    ("CURSOR_AGENT", "cursor"),
+    ("GEMINI_CLI", "gemini-cli"),
+    ("OPENCODE", "opencode"),
+]
+
+
+def detect(env=None):
+    """Detect which AI coding agent (if any) is driving the current process.
+
+    Args:
+        env: Optional dict-like object for environment variable lookup.
+             Defaults to os.environ. Exists for testability.
+
+    Returns:
+        The agent product string if exactly one agent is detected,
+        or an empty string otherwise.
+    """
+    if env is None:
+        env = os.environ
+
+    detected = [product for var, product in KNOWN_AGENTS if env.get(var)]
+
+    if len(detected) == 1:
+        return detected[0]
+    return ""

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/src/databricks/sql/common/feature_flag.py

@@ -113,6 +113,7 @@ class FeatureFlagsContext:
             # Authenticate the request
             self._connection.session.auth_provider.add_headers(headers)
             headers["User-Agent"] = self._connection.session.useragent_header
+            headers.update(self._connection.session.get_spog_headers())
 
             response = self._http_client.request(
                 HttpMethod.GET, self._feature_flag_endpoint, headers=headers, timeout=30

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/src/databricks/sql/common/unified_http_client.py

@@ -217,7 +217,7 @@ class UnifiedHttpClient:
             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.
 
@@ -225,7 +225,7 @@ class UnifiedHttpClient:
             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
@@ -291,6 +291,14 @@ class UnifiedHttpClient:
         # 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:

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/src/databricks/sql/parameters/native.py

@@ -659,7 +659,7 @@ def dbsql_parameter_from_primitive(
     elif isinstance(value, str):
         return StringParameter(value=value, name=name)
     elif isinstance(value, float):
-        return FloatParameter(value=value, name=name)
+        return DoubleParameter(value=value, name=name)
     elif isinstance(value, datetime.datetime):
         return TimestampParameter(value=value, name=name)
     elif isinstance(value, datetime.date):

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/src/databricks/sql/session.py

@@ -13,6 +13,7 @@ from databricks.sql.backend.sea.backend import SeaDatabricksClient
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.backend.types import SessionId, BackendType
 from databricks.sql.common.unified_http_client import UnifiedHttpClient
+from databricks.sql.common.agent import detect as detect_agent
 
 logger = logging.getLogger(__name__)
 
@@ -64,9 +65,21 @@ class Session:
         else:
             self.useragent_header = "{}/{}".format(USER_AGENT_NAME, __version__)
 
+        agent_product = detect_agent()
+        if agent_product:
+            self.useragent_header += " agent/{}".format(agent_product)
+
         base_headers = [("User-Agent", self.useragent_header)]
         all_headers = (http_headers or []) + base_headers
 
+        # Extract ?o=<workspaceId> from http_path for SPOG routing.
+        # On SPOG hosts, the httpPath contains ?o=<workspaceId> which routes Thrift
+        # requests via the URL. For SEA, telemetry, and feature flags (which use
+        # separate endpoints), we inject x-databricks-org-id as an HTTP header.
+        self._spog_headers = self._extract_spog_headers(http_path, all_headers)
+        if self._spog_headers:
+            all_headers = all_headers + list(self._spog_headers.items())
+
         self.ssl_options = SSLOptions(
             # Double negation is generally a bad thing, but we have to keep backward compatibility
             tls_verify=not kwargs.get(
@@ -131,6 +144,44 @@ class Session:
         }
         return databricks_client_class(**common_args)
 
+    @staticmethod
+    def _extract_spog_headers(http_path, existing_headers):
+        """Extract ?o=<workspaceId> from http_path and return as a header dict for SPOG routing."""
+        if not http_path or "?" not in http_path:
+            return {}
+
+        from urllib.parse import parse_qs
+
+        query_string = http_path.split("?", 1)[1]
+        params = parse_qs(query_string)
+        org_id = params.get("o", [None])[0]
+        if not org_id:
+            logger.debug(
+                "SPOG header extraction: http_path has query string but no ?o= param, "
+                "skipping x-databricks-org-id injection"
+            )
+            return {}
+
+        # Don't override if explicitly set
+        if any(k == "x-databricks-org-id" for k, _ in existing_headers):
+            logger.debug(
+                "SPOG header extraction: x-databricks-org-id already set by caller, "
+                "not overriding with ?o=%s from http_path",
+                org_id,
+            )
+            return {}
+
+        logger.debug(
+            "SPOG header extraction: injecting x-databricks-org-id=%s "
+            "(extracted from ?o= in http_path)",
+            org_id,
+        )
+        return {"x-databricks-org-id": org_id}
+
+    def get_spog_headers(self):
+        """Returns SPOG routing headers (x-databricks-org-id) if ?o= was in http_path."""
+        return dict(self._spog_headers)
+
     def open(self):
         self._session_id = self.backend.open_session(
             session_configuration=self.session_configuration,

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/src/databricks/sql/telemetry/models/event.py

@@ -57,6 +57,7 @@ class DriverConnectionParameters(JsonSerializableMixin):
         support_many_parameters (bool): Whether many parameters are supported
         enable_complex_datatype_support (bool): Whether complex datatypes are supported
         allowed_volume_ingestion_paths (str): Allowed paths for volume ingestion
+        query_tags (str): Query tags for tracking and attribution
     """
 
     http_path: str
@@ -84,6 +85,7 @@ class DriverConnectionParameters(JsonSerializableMixin):
     support_many_parameters: Optional[bool] = None
     enable_complex_datatype_support: Optional[bool] = None
     allowed_volume_ingestion_paths: Optional[str] = None
+    query_tags: Optional[str] = None
 
 
 @dataclass

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/src/databricks/sql/telemetry/telemetry_client.py

@@ -42,6 +42,7 @@ from databricks.sql.telemetry.utils import BaseTelemetryClient
 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,
@@ -187,6 +188,7 @@ class TelemetryClient(BaseTelemetryClient):
         executor,
         batch_size: int,
         client_context,
+        extra_headers: Optional[Dict[str, str]] = None,
     ) -> None:
         logger.debug("Initializing TelemetryClient for connection: %s", session_id_hex)
         self._telemetry_enabled = telemetry_enabled
@@ -194,6 +196,7 @@ class TelemetryClient(BaseTelemetryClient):
         self._session_id_hex = session_id_hex
         self._auth_provider = auth_provider
         self._user_agent = None
+        self._extra_headers = extra_headers or {}
 
         # OPTIMIZATION: Use lock-free Queue instead of list + lock
         # Queue is thread-safe internally and has better performance under concurrency
@@ -286,6 +289,8 @@ class TelemetryClient(BaseTelemetryClient):
         if self._auth_provider:
             self._auth_provider.add_headers(headers)
 
+        headers.update(self._extra_headers)
+
         try:
             logger.debug("Submitting telemetry request to thread pool")
 
@@ -417,10 +422,38 @@ class TelemetryClient(BaseTelemetryClient):
         )
 
     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:
     """
@@ -558,6 +591,7 @@ class TelemetryClientFactory:
         host_url,
         batch_size,
         client_context,
+        extra_headers=None,
     ):
         """
         Initialize a telemetry client for a specific connection if telemetry is enabled.
@@ -598,6 +632,7 @@ class TelemetryClientFactory:
                             executor=TelemetryClientFactory._executor,
                             batch_size=batch_size,
                             client_context=client_context,
+                            extra_headers=extra_headers,
                         )
                         TelemetryClientFactory._clients[
                             host_url
@@ -674,7 +709,8 @@ class TelemetryClientFactory:
                 )
                 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
@@ -689,9 +725,15 @@ class TelemetryClientFactory:
         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(

{databricks_sql_connector-4.2.4 → databricks_sql_connector-4.2.6}/src/databricks/sql/utils.py

@@ -38,6 +38,21 @@ DEFAULT_ERROR_CONTEXT = "Unknown error"
 logger = logging.getLogger(__name__)
 
 
+def get_session_config_value(
+    session_configuration: Optional[Dict[str, Any]], key: str
+) -> Optional[str]:
+    """Get a session configuration value with case-insensitive key matching"""
+    if not session_configuration:
+        return None
+
+    key_upper = key.upper()
+    for k, v in session_configuration.items():
+        if k.upper() == key_upper:
+            return str(v) if v is not None else None
+
+    return None
+
+
 class ResultSetQueue(ABC):
     @abstractmethod
     def next_n_rows(self, num_rows: int):
@@ -880,7 +895,50 @@ def concat_table_chunks(
                 result_table[j].extend(table_chunks[i].column_table[j])
         return ColumnTable(result_table, table_chunks[0].column_names)
     else:
-        return pyarrow.concat_tables(table_chunks)
+        return pyarrow.concat_tables(table_chunks, promote_options="default")
+
+
+def serialize_query_tags(
+    query_tags: Optional[Dict[str, Optional[str]]]
+) -> Optional[str]:
+    """
+    Serialize query_tags dictionary to a string format.
+
+    Format: "key1:value1,key2:value2"
+    Special cases:
+    - If value is None, omit the colon and value (e.g., "key1:value1,key2,key3:value3")
+    - Escape special characters (:, ,, \\) in values with a leading backslash
+    - Backslashes in keys are escaped; other special characters in keys are not escaped
+
+    Args:
+        query_tags: Dictionary of query tags where keys are strings and values are optional strings
+
+    Returns:
+        Serialized string or None if query_tags is None or empty
+    """
+    if not query_tags:
+        return None
+
+    def escape_value(value: str) -> str:
+        """Escape special characters in tag values."""
+        # Escape backslash first to avoid double-escaping
+        value = value.replace("\\", r"\\")
+        # Escape colon and comma
+        value = value.replace(":", r"\:")
+        value = value.replace(",", r"\,")
+        return value
+
+    serialized_parts = []
+    for key, value in query_tags.items():
+        escaped_key = key.replace("\\", r"\\")
+        if value is None:
+            # No colon or value when value is None
+            serialized_parts.append(escaped_key)
+        else:
+            escaped_value = escape_value(value)
+            serialized_parts.append(f"{escaped_key}:{escaped_value}")
+
+    return ",".join(serialized_parts)
 
 
 def build_client_context(server_hostname: str, version: str, **kwargs):
@@ -899,12 +957,18 @@ def build_client_context(server_hostname: str, version: str, **kwargs):
     )
 
     # Build user agent
+    from databricks.sql.common.agent import detect as detect_agent
+
     user_agent_entry = kwargs.get("user_agent_entry", "")
     if user_agent_entry:
         user_agent = f"PyDatabricksSqlConnector/{version} ({user_agent_entry})"
     else:
         user_agent = f"PyDatabricksSqlConnector/{version}"
 
+    agent_product = detect_agent()
+    if agent_product:
+        user_agent += f" agent/{agent_product}"
+
     # Explicitly construct ClientContext with proper types
     return ClientContext(
         hostname=server_hostname,