databricks-sql-connector 4.1.1__tar.gz → 4.1.3__tar.gz

{databricks_sql_connector-4.1.1 → databricks_sql_connector-4.1.3}/CHANGELOG.md

@@ -1,5 +1,13 @@
 # Release History
 
+# 4.1.3 (2025-09-17)
+- Query tags integration (databricks/databricks-sql-python#663 by @sreekanth-db)
+- Add variant support (databricks/databricks-sql-python#560 by @shivam2680)
+
+# 4.1.2 (2025-08-22)
+- Streaming ingestion support for PUT operation (databricks/databricks-sql-python#643 by @sreekanth-db)
+- Removed use_threads argument on concat_tables for compatibility with pyarrow<14 (databricks/databricks-sql-python#684 by @jprakash-db)
+
 # 4.1.1 (2025-08-21)
 - Add documentation for proxy support (databricks/databricks-sql-python#680 by @vikrantpuppala)
 - Fix compatibility with urllib3<2 and add CI actions to improve dependency checks (databricks/databricks-sql-python#678 by @vikrantpuppala)

{databricks_sql_connector-4.1.1 → databricks_sql_connector-4.1.3}/PKG-INFO

@@ -1,8 +1,9 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: databricks-sql-connector
-Version: 4.1.1
+Version: 4.1.3
 Summary: Databricks SQL Connector for Python
 License: Apache-2.0
+License-File: LICENSE
 Author: Databricks
 Author-email: databricks-sql-connector-maintainers@databricks.com
 Requires-Python: >=3.8.0,<4.0.0
@@ -14,6 +15,7 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
 Provides-Extra: pyarrow
 Requires-Dist: lz4 (>=4.0.2,<5.0.0)
 Requires-Dist: oauthlib (>=3.1.0,<4.0.0)

{databricks_sql_connector-4.1.1 → databricks_sql_connector-4.1.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databricks-sql-connector"
-version = "4.1.1"
+version = "4.1.3"
 description = "Databricks SQL Connector for Python"
 authors = ["Databricks <databricks-sql-connector-maintainers@databricks.com>"]
 license = "Apache-2.0"

{databricks_sql_connector-4.1.1 → databricks_sql_connector-4.1.3}/src/databricks/sql/__init__.py

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

{databricks_sql_connector-4.1.1 → databricks_sql_connector-4.1.3}/src/databricks/sql/backend/sea/utils/constants.py

@@ -15,6 +15,7 @@ ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP: Dict[str, str] = {
     "STATEMENT_TIMEOUT": "0",
     "TIMEZONE": "UTC",
     "USE_CACHED_RESULT": "true",
+    "QUERY_TAGS": "",
 }
 
 

{databricks_sql_connector-4.1.1 → databricks_sql_connector-4.1.3}/src/databricks/sql/backend/thrift_backend.py

@@ -735,7 +735,7 @@ class ThriftDatabricksClient(DatabricksClient):
         return pyarrow.schema([convert_col(col) for col in t_table_schema.columns])
 
     @staticmethod
-    def _col_to_description(col, session_id_hex=None):
+    def _col_to_description(col, field=None, session_id_hex=None):
         type_entry = col.typeDesc.types[0]
 
         if type_entry.primitiveEntry:
@@ -764,12 +764,39 @@ class ThriftDatabricksClient(DatabricksClient):
         else:
             precision, scale = None, None
 
+        # Extract variant type from field if available
+        if field is not None:
+            try:
+                # Check for variant type in metadata
+                if field.metadata and b"Spark:DataType:SqlName" in field.metadata:
+                    sql_type = field.metadata.get(b"Spark:DataType:SqlName")
+                    if sql_type == b"VARIANT":
+                        cleaned_type = "variant"
+            except Exception as e:
+                logger.debug(f"Could not extract variant type from field: {e}")
+
         return col.columnName, cleaned_type, None, None, precision, scale, None
 
     @staticmethod
-    def _hive_schema_to_description(t_table_schema, session_id_hex=None):
+    def _hive_schema_to_description(
+        t_table_schema, schema_bytes=None, session_id_hex=None
+    ):
+        field_dict = {}
+        if pyarrow and schema_bytes:
+            try:
+                arrow_schema = pyarrow.ipc.read_schema(pyarrow.py_buffer(schema_bytes))
+                # Build a dictionary mapping column names to fields
+                for field in arrow_schema:
+                    field_dict[field.name] = field
+            except Exception as e:
+                logger.debug(f"Could not parse arrow schema: {e}")
+
         return [
-            ThriftDatabricksClient._col_to_description(col, session_id_hex)
+            ThriftDatabricksClient._col_to_description(
+                col,
+                field_dict.get(col.columnName) if field_dict else None,
+                session_id_hex,
+            )
             for col in t_table_schema.columns
         ]
 
@@ -802,11 +829,6 @@ class ThriftDatabricksClient(DatabricksClient):
             or direct_results.resultSet.hasMoreRows
         )
 
-        description = self._hive_schema_to_description(
-            t_result_set_metadata_resp.schema,
-            self._session_id_hex,
-        )
-
         if pyarrow:
             schema_bytes = (
                 t_result_set_metadata_resp.arrowSchema
@@ -819,6 +841,12 @@ class ThriftDatabricksClient(DatabricksClient):
         else:
             schema_bytes = None
 
+        description = self._hive_schema_to_description(
+            t_result_set_metadata_resp.schema,
+            schema_bytes,
+            self._session_id_hex,
+        )
+
         lz4_compressed = t_result_set_metadata_resp.lz4Compressed
         command_id = CommandId.from_thrift_handle(resp.operationHandle)
 
@@ -863,11 +891,6 @@ class ThriftDatabricksClient(DatabricksClient):
 
         t_result_set_metadata_resp = resp.resultSetMetadata
 
-        description = self._hive_schema_to_description(
-            t_result_set_metadata_resp.schema,
-            self._session_id_hex,
-        )
-
         if pyarrow:
             schema_bytes = (
                 t_result_set_metadata_resp.arrowSchema
@@ -880,6 +903,12 @@ class ThriftDatabricksClient(DatabricksClient):
         else:
             schema_bytes = None
 
+        description = self._hive_schema_to_description(
+            t_result_set_metadata_resp.schema,
+            schema_bytes,
+            self._session_id_hex,
+        )
+
         lz4_compressed = t_result_set_metadata_resp.lz4Compressed
         is_staging_operation = t_result_set_metadata_resp.isStagingOperation
         has_more_rows = resp.hasMoreRows

{databricks_sql_connector-4.1.1 → databricks_sql_connector-4.1.3}/src/databricks/sql/client.py

@@ -1,5 +1,5 @@
 import time
-from typing import Dict, Tuple, List, Optional, Any, Union, Sequence
+from typing import Dict, Tuple, List, Optional, Any, Union, Sequence, BinaryIO
 import pandas
 
 try:
@@ -662,7 +662,9 @@ class Cursor:
             )
 
     def _handle_staging_operation(
-        self, staging_allowed_local_path: Union[None, str, List[str]]
+        self,
+        staging_allowed_local_path: Union[None, str, List[str]],
+        input_stream: Optional[BinaryIO] = None,
     ):
         """Fetch the HTTP request instruction from a staging ingestion command
         and call the designated handler.
@@ -671,6 +673,28 @@ class Cursor:
         is not descended from staging_allowed_local_path.
         """
 
+        assert self.active_result_set is not None
+        row = self.active_result_set.fetchone()
+        assert row is not None
+
+        # May be real headers, or could be json string
+        headers = (
+            json.loads(row.headers) if isinstance(row.headers, str) else row.headers
+        )
+        headers = dict(headers) if headers else {}
+
+        # Handle __input_stream__ token for PUT operations
+        if (
+            row.operation == "PUT"
+            and getattr(row, "localFile", None) == "__input_stream__"
+        ):
+            return self._handle_staging_put_stream(
+                presigned_url=row.presignedUrl,
+                stream=input_stream,
+                headers=headers,
+            )
+
+        # For non-streaming operations, validate staging_allowed_local_path
         if isinstance(staging_allowed_local_path, type(str())):
             _staging_allowed_local_paths = [staging_allowed_local_path]
         elif isinstance(staging_allowed_local_path, type(list())):
@@ -685,10 +709,6 @@ class Cursor:
             os.path.abspath(i) for i in _staging_allowed_local_paths
         ]
 
-        assert self.active_result_set is not None
-        row = self.active_result_set.fetchone()
-        assert row is not None
-
         # Must set to None in cases where server response does not include localFile
         abs_localFile = None
 
@@ -711,19 +731,16 @@ class Cursor:
                     session_id_hex=self.connection.get_session_id_hex(),
                 )
 
-        # May be real headers, or could be json string
-        headers = (
-            json.loads(row.headers) if isinstance(row.headers, str) else row.headers
-        )
-
         handler_args = {
             "presigned_url": row.presignedUrl,
             "local_file": abs_localFile,
-            "headers": dict(headers) or {},
+            "headers": headers,
         }
 
         logger.debug(
-            f"Attempting staging operation indicated by server: {row.operation} - {getattr(row, 'localFile', '')}"
+            "Attempting staging operation indicated by server: %s - %s",
+            row.operation,
+            getattr(row, "localFile", ""),
         )
 
         # TODO: Create a retry loop here to re-attempt if the request times out or fails
@@ -762,6 +779,10 @@ class Cursor:
                 HttpMethod.PUT, presigned_url, body=fh.read(), headers=headers
             )
 
+        self._handle_staging_http_response(r)
+
+    def _handle_staging_http_response(self, r):
+
         # fmt: off
         # HTTP status codes
         OK = 200
@@ -784,6 +805,37 @@ class Cursor:
                 + "but not yet applied on the server. It's possible this command may fail later."
             )
 
+    @log_latency(StatementType.SQL)
+    def _handle_staging_put_stream(
+        self,
+        presigned_url: str,
+        stream: BinaryIO,
+        headers: dict = {},
+    ) -> None:
+        """Handle PUT operation with streaming data.
+
+        Args:
+            presigned_url: The presigned URL for upload
+            stream: Binary stream to upload
+            headers: HTTP headers
+
+        Raises:
+            ProgrammingError: If no input stream is provided
+            OperationalError: If the upload fails
+        """
+
+        if not stream:
+            raise ProgrammingError(
+                "No input stream provided for streaming operation",
+                session_id_hex=self.connection.get_session_id_hex(),
+            )
+
+        r = self.connection.http_client.request(
+            HttpMethod.PUT, presigned_url, body=stream.read(), headers=headers
+        )
+
+        self._handle_staging_http_response(r)
+
     @log_latency(StatementType.SQL)
     def _handle_staging_get(
         self, local_file: str, presigned_url: str, headers: Optional[dict] = None
@@ -840,6 +892,7 @@ class Cursor:
         operation: str,
         parameters: Optional[TParameterCollection] = None,
         enforce_embedded_schema_correctness=False,
+        input_stream: Optional[BinaryIO] = None,
     ) -> "Cursor":
         """
         Execute a query and wait for execution to complete.
@@ -914,7 +967,8 @@ class Cursor:
 
         if self.active_result_set and self.active_result_set.is_staging_operation:
             self._handle_staging_operation(
-                staging_allowed_local_path=self.connection.staging_allowed_local_path
+                staging_allowed_local_path=self.connection.staging_allowed_local_path,
+                input_stream=input_stream,
             )
 
         return self

{databricks_sql_connector-4.1.1 → databricks_sql_connector-4.1.3}/src/databricks/sql/utils.py

@@ -298,7 +298,7 @@ class CloudFetchQueue(ResultSetQueue, ABC):
             num_rows -= table_slice.num_rows
 
         logger.debug("CloudFetchQueue: collected {} next rows".format(results.num_rows))
-        return pyarrow.concat_tables(partial_result_chunks, use_threads=True)
+        return concat_table_chunks(partial_result_chunks)
 
     def remaining_rows(self) -> "pyarrow.Table":
         """
@@ -321,7 +321,7 @@ class CloudFetchQueue(ResultSetQueue, ABC):
             self.table_row_index += table_slice.num_rows
             self.table = self._create_next_table()
             self.table_row_index = 0
-        return pyarrow.concat_tables(partial_result_chunks, use_threads=True)
+        return concat_table_chunks(partial_result_chunks)
 
     def _create_table_at_offset(self, offset: int) -> Union["pyarrow.Table", None]:
         """Create next table at the given row offset"""
@@ -880,7 +880,7 @@ 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, use_threads=True)
+        return pyarrow.concat_tables(table_chunks)
 
 
 def build_client_context(server_hostname: str, version: str, **kwargs):