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

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

@@ -1,5 +1,16 @@
 # 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)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: databricks-sql-connector
-Version: 4.2.5
+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.5 → databricks_sql_connector-4.2.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databricks-sql-connector"
-version = "4.2.5"
+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" }

{databricks_sql_connector-4.2.5 → 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.5"
+__version__ = "4.2.6"
 USER_AGENT_NAME = "PyDatabricksSqlConnector"
 
 # These two functions are pyhive legacy

{databricks_sql_connector-4.2.5 → 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.5 → 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.5 → 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.5 → 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.5 → databricks_sql_connector-4.2.6}/src/databricks/sql/client.py

@@ -36,6 +36,7 @@ from databricks.sql.utils import (
     ColumnQueue,
     build_client_context,
     get_session_config_value,
+    serialize_query_tags,
 )
 from databricks.sql.parameters.native import (
     DbsqlParameterBase,
@@ -106,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:
         """
@@ -281,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)
@@ -342,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(
@@ -1263,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.
@@ -1293,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
         """
 
@@ -1333,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:
@@ -1349,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":
         """
 
@@ -1356,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:
         """
 
@@ -1392,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
@@ -1448,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.
 
@@ -1457,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.5 → 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.5 → 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.5 → 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.5 → databricks_sql_connector-4.2.6}/src/databricks/sql/telemetry/telemetry_client.py

@@ -188,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
@@ -195,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
@@ -287,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")
 
@@ -587,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.
@@ -627,6 +632,7 @@ class TelemetryClientFactory:
                             executor=TelemetryClientFactory._executor,
                             batch_size=batch_size,
                             client_context=client_context,
+                            extra_headers=extra_headers,
                         )
                         TelemetryClientFactory._clients[
                             host_url

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

@@ -895,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):
@@ -914,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,