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

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

@@ -1,5 +1,9 @@
 # Release History
 
+# 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.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: databricks-sql-connector
-Version: 4.1.1
+Version: 4.1.2
 Summary: Databricks SQL Connector for Python
 License: Apache-2.0
 Author: Databricks

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

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databricks-sql-connector"
-version = "4.1.1"
+version = "4.1.2"
 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.2}/src/databricks/sql/__init__.py

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

{databricks_sql_connector-4.1.1 → databricks_sql_connector-4.1.2}/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.2}/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):