streaming put supoprt

Signed-off-by: Sreekanth Vadigi <[email protected]>

Author: sreekanth-db

examples/streaming_put.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python3
+"""
+Simple example of streaming PUT operations.
+
+This demonstrates the basic usage of streaming PUT with the __input_stream__ token.
+"""
+
+import io
+import os
+from databricks import sql
+
+with sql.connect(
+    server_hostname=os.getenv("DATABRICKS_SERVER_HOSTNAME"),
+    http_path=os.getenv("DATABRICKS_HTTP_PATH"),
+    access_token=os.getenv("DATABRICKS_TOKEN"),
+) as connection:
+
+    with connection.cursor() as cursor:
+        # Create a simple data stream
+        data = b"Hello, streaming world!"
+        stream = io.BytesIO(data)
+        
+        # Get catalog, schema, and volume from environment variables
+        catalog = os.getenv("DATABRICKS_CATALOG")
+        schema = os.getenv("DATABRICKS_SCHEMA")
+        volume = os.getenv("DATABRICKS_VOLUME")
+        
+        # Upload to Unity Catalog volume
+        cursor.execute(
+            f"PUT '__input_stream__' INTO '/Volumes/{catalog}/{schema}/{volume}/hello.txt' OVERWRITE",
+            input_stream=stream
+        )
+        
+        print("File uploaded successfully!") 

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 @@ def _check_not_closed(self):
             )
 
     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 @@ def _handle_staging_operation(
         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 @@ def _handle_staging_operation(
             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 @@ def _handle_staging_operation(
                     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 @@ def _handle_staging_put(
                 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,38 @@ def _handle_staging_put(
                 + "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 +893,7 @@ def execute(
         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 +968,8 @@ def execute(
 
         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

tests/e2e/common/streaming_put_tests.py

@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+"""
+E2E tests for streaming PUT operations.
+"""
+
+import io
+import logging
+import pytest
+
+logger = logging.getLogger(__name__)
+
+
+class PySQLStreamingPutTestSuiteMixin:
+    """Test suite for streaming PUT operations."""
+
+    def test_streaming_put_basic(self, catalog, schema):
+        """Test basic streaming PUT functionality."""
+        
+        # Create test data
+        test_data = b"Hello, streaming world! This is test data."
+        filename = "streaming_put_test.txt"
+        file_path = f"/Volumes/{catalog}/{schema}/e2etests/{filename}"
+        
+        try:
+            with self.connection() as conn:
+                with conn.cursor() as cursor:
+                    self._cleanup_test_file(file_path)
+                    
+                    with io.BytesIO(test_data) as stream:
+                        cursor.execute(
+                            f"PUT '__input_stream__' INTO '{file_path}'",
+                            input_stream=stream
+                        )
+                    
+                    # Verify file exists
+                    cursor.execute(f"LIST '/Volumes/{catalog}/{schema}/e2etests/'")
+                    files = cursor.fetchall()
+                    
+                    # Check if our file is in the list
+                    file_paths = [row[0] for row in files]
+                    assert file_path in file_paths, f"File {file_path} not found in {file_paths}"
+        finally:
+            self._cleanup_test_file(file_path)
+    
+    def test_streaming_put_missing_stream(self, catalog, schema):
+        """Test that missing stream raises appropriate error."""
+        
+        with self.connection() as conn:
+            with conn.cursor() as cursor:
+                # Test without providing stream
+                with pytest.raises(Exception):  # Should fail
+                    cursor.execute(
+                        f"PUT '__input_stream__' INTO '/Volumes/{catalog}/{schema}/e2etests/test.txt'"
+                        # Note: No input_stream parameter
+                    ) 
+
+    def _cleanup_test_file(self, file_path):
+        """Clean up a test file if it exists."""
+        try:
+            with self.connection(extra_params={"staging_allowed_local_path": "/"}) as conn:
+                with conn.cursor() as cursor:
+                    cursor.execute(f"REMOVE '{file_path}'")
+                    logger.info("Successfully cleaned up test file: %s", file_path)
+        except Exception as e:
+            logger.error("Cleanup failed for %s: %s", file_path, e)

tests/e2e/test_driver.py

@@ -50,6 +50,7 @@
 from tests.e2e.common.retry_test_mixins import PySQLRetryTestsMixin
 
 from tests.e2e.common.uc_volume_tests import PySQLUCVolumeTestSuiteMixin
+from tests.e2e.common.streaming_put_tests import PySQLStreamingPutTestSuiteMixin
 
 from databricks.sql.exc import SessionAlreadyClosedError
 
@@ -290,6 +291,7 @@ class TestPySQLCoreSuite(
     PySQLStagingIngestionTestSuiteMixin,
     PySQLRetryTestsMixin,
     PySQLUCVolumeTestSuiteMixin,
+    PySQLStreamingPutTestSuiteMixin,
 ):
     validate_row_value_type = True
     validate_result = True