addressing review comments

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

Author: sreekanth-db

examples/streaming_put.py

@@ -9,28 +9,26 @@
 import os
 from databricks import sql
 
-def main():
-    """Simple streaming PUT example."""
-    
-    # Connect to Databricks
-    connection = sql.connect(
-        server_hostname=os.getenv("DATABRICKS_SERVER_HOSTNAME"),
-        http_path=os.getenv("DATABRICKS_HTTP_PATH"),
-        access_token=os.getenv("DATABRICKS_TOKEN"),
-    )
-    
+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(
-            "PUT '__input_stream__' INTO '/Volumes/my_catalog/my_schema/my_volume/hello.txt'",
+            f"PUT '__input_stream__' INTO '/Volumes/{catalog}/{schema}/{volume}/hello.txt' OVERWRITE",
             input_stream=stream
         )
 
-        print("File uploaded successfully!")
-
-if __name__ == "__main__":
-    main() 
+        print("File uploaded successfully!") 

src/databricks/sql/client.py

@@ -67,6 +67,7 @@
 )
 from databricks.sql.telemetry.latency_logger import log_latency
 from databricks.sql.telemetry.models.enums import StatementType
+from databricks.sql.common.http import DatabricksHttpClient, HttpMethod
 
 logger = logging.getLogger(__name__)
 
@@ -455,7 +456,6 @@ def __init__(
         self.active_command_id = None
         self.escaper = ParamEscaper()
         self.lastrowid = None
-        self._input_stream_data: Optional[BinaryIO] = None
 
         self.ASYNC_DEFAULT_POLLING_INTERVAL = 2
 
@@ -616,8 +616,29 @@ def _check_not_closed(self):
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
+    def _validate_staging_http_response(self, response: requests.Response, operation_name: str = "staging operation") -> None:
+
+        # Check response codes
+        OK = requests.codes.ok  # 200
+        CREATED = requests.codes.created  # 201
+        ACCEPTED = requests.codes.accepted  # 202
+        NO_CONTENT = requests.codes.no_content  # 204
+
+        if response.status_code not in [OK, CREATED, NO_CONTENT, ACCEPTED]:
+            raise OperationalError(
+                f"{operation_name} over HTTP was unsuccessful: {response.status_code}-{response.text}",
+                session_id_hex=self.connection.get_session_id_hex(),
+            )
+
+        if response.status_code == ACCEPTED:
+            logger.debug(
+                "Response code %s from server indicates %s was accepted "
+                "but not yet applied on the server. It's possible this command may fail later.",
+                ACCEPTED, operation_name
+            )
+
     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.
@@ -641,14 +662,9 @@ def _handle_staging_operation(
             row.operation == "PUT"
             and getattr(row, "localFile", None) == "__input_stream__"
         ):
-            if not self._input_stream_data:
-                raise ProgrammingError(
-                    "No input stream provided for streaming operation",
-                    session_id_hex=self.connection.get_session_id_hex(),
-                )
             return self._handle_staging_put_stream(
                 presigned_url=row.presignedUrl,
-                stream=self._input_stream_data,
+                stream=input_stream,
                 headers=headers,
             )
 
@@ -696,7 +712,8 @@ def _handle_staging_operation(
         }
 
         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
@@ -720,54 +737,37 @@ def _handle_staging_put_stream(
         self,
         presigned_url: str,
         stream: BinaryIO,
-        headers: Optional[dict] = None,
+        headers: dict = {},
     ) -> None:
         """Handle PUT operation with streaming data.
 
         Args:
             presigned_url: The presigned URL for upload
             stream: Binary stream to upload
-            headers: Optional HTTP headers
+            headers: HTTP headers
 
         Raises:
+            ProgrammingError: If no input stream is provided
             OperationalError: If the upload fails
         """
 
-        # Prepare headers
-        http_headers = dict(headers) if headers else {}
-
-        try:
-            # Stream directly to presigned URL
-            response = requests.put(
-                url=presigned_url,
-                data=stream,
-                headers=http_headers,
-                timeout=300,  # 5 minute timeout
+        if not stream:
+            raise ProgrammingError(
+                "No input stream provided for streaming operation",
+                session_id_hex=self.connection.get_session_id_hex(),
             )
 
-            # Check response codes
-            OK = requests.codes.ok  # 200
-            CREATED = requests.codes.created  # 201
-            ACCEPTED = requests.codes.accepted  # 202
-            NO_CONTENT = requests.codes.no_content  # 204
+        http_client = DatabricksHttpClient.get_instance()
 
-            if response.status_code not in [OK, CREATED, NO_CONTENT, ACCEPTED]:
-                raise OperationalError(
-                    f"Staging operation over HTTP was unsuccessful: {response.status_code}-{response.text}",
-                    session_id_hex=self.connection.get_session_id_hex(),
-                )
-
-            if response.status_code == ACCEPTED:
-                logger.debug(
-                    f"Response code {ACCEPTED} from server indicates upload was accepted "
-                    "but not yet applied on the server. It's possible this command may fail later."
-                )
-
-        except requests.exceptions.RequestException as e:
-            raise OperationalError(
-                f"HTTP request failed during stream upload: {str(e)}",
-                session_id_hex=self.connection.get_session_id_hex(),
-            ) from e
+        # Stream directly to presigned URL
+        with http_client.execute(
+            method=HttpMethod.PUT,
+            url=presigned_url,
+            data=stream,
+            headers=headers,
+            timeout=300,  # 5 minute timeout
+        ) as response:
+            self._validate_staging_http_response(response, "stream upload")
 
     @log_latency(StatementType.SQL)
     def _handle_staging_put(
@@ -787,27 +787,7 @@ def _handle_staging_put(
         with open(local_file, "rb") as fh:
             r = requests.put(url=presigned_url, data=fh, headers=headers)
 
-        # fmt: off
-        # Design borrowed from: https://stackoverflow.com/a/2342589/5093960
-
-        OK = requests.codes.ok                  # 200
-        CREATED = requests.codes.created        # 201
-        ACCEPTED = requests.codes.accepted      # 202
-        NO_CONTENT = requests.codes.no_content  # 204
-
-        # fmt: on
-
-        if r.status_code not in [OK, CREATED, NO_CONTENT, ACCEPTED]:
-            raise OperationalError(
-                f"Staging operation over HTTP was unsuccessful: {r.status_code}-{r.text}",
-                session_id_hex=self.connection.get_session_id_hex(),
-            )
-
-        if r.status_code == ACCEPTED:
-            logger.debug(
-                f"Response code {ACCEPTED} from server indicates ingestion command was accepted "
-                + "but not yet applied on the server. It's possible this command may fail later."
-            )
+        self._validate_staging_http_response(r, "file upload")
 
     @log_latency(StatementType.SQL)
     def _handle_staging_get(
@@ -856,8 +836,8 @@ def execute(
         self,
         operation: str,
         parameters: Optional[TParameterCollection] = None,
-        input_stream: Optional[BinaryIO] = None,
         enforce_embedded_schema_correctness=False,
+        input_stream: Optional[BinaryIO] = None,
     ) -> "Cursor":
         """
         Execute a query and wait for execution to complete.
@@ -894,62 +874,49 @@ def execute(
         logger.debug(
             "Cursor.execute(operation=%s, parameters=%s)", operation, parameters
         )
-        try:
-            # Store stream data if provided
-            self._input_stream_data = None
-            if input_stream is not None:
-                # Validate stream has required methods
-                if not hasattr(input_stream, "read"):
-                    raise TypeError(
-                        "input_stream must be a binary stream with read() method"
-                    )
-                self._input_stream_data = input_stream
+        param_approach = self._determine_parameter_approach(parameters)
+        if param_approach == ParameterApproach.NONE:
+            prepared_params = NO_NATIVE_PARAMS
+            prepared_operation = operation
 
-            param_approach = self._determine_parameter_approach(parameters)
-            if param_approach == ParameterApproach.NONE:
-                prepared_params = NO_NATIVE_PARAMS
-                prepared_operation = operation
+        elif param_approach == ParameterApproach.INLINE:
+            prepared_operation, prepared_params = self._prepare_inline_parameters(
+                operation, parameters
+            )
+        elif param_approach == ParameterApproach.NATIVE:
+            normalized_parameters = self._normalize_tparametercollection(parameters)
+            param_structure = self._determine_parameter_structure(
+                normalized_parameters
+            )
+            transformed_operation = transform_paramstyle(
+                operation, normalized_parameters, param_structure
+            )
+            prepared_operation, prepared_params = self._prepare_native_parameters(
+                transformed_operation, normalized_parameters, param_structure
+            )
 
-            elif param_approach == ParameterApproach.INLINE:
-                prepared_operation, prepared_params = self._prepare_inline_parameters(
-                    operation, parameters
-                )
-            elif param_approach == ParameterApproach.NATIVE:
-                normalized_parameters = self._normalize_tparametercollection(parameters)
-                param_structure = self._determine_parameter_structure(
-                    normalized_parameters
-                )
-                transformed_operation = transform_paramstyle(
-                    operation, normalized_parameters, param_structure
-                )
-                prepared_operation, prepared_params = self._prepare_native_parameters(
-                    transformed_operation, normalized_parameters, param_structure
-                )
+        self._check_not_closed()
+        self._close_and_clear_active_result_set()
+        self.active_result_set = self.backend.execute_command(
+            operation=prepared_operation,
+            session_id=self.connection.session.session_id,
+            max_rows=self.arraysize,
+            max_bytes=self.buffer_size_bytes,
+            lz4_compression=self.connection.lz4_compression,
+            cursor=self,
+            use_cloud_fetch=self.connection.use_cloud_fetch,
+            parameters=prepared_params,
+            async_op=False,
+            enforce_embedded_schema_correctness=enforce_embedded_schema_correctness,
+        )
 
-            self._check_not_closed()
-            self._close_and_clear_active_result_set()
-            self.active_result_set = self.backend.execute_command(
-                operation=prepared_operation,
-                session_id=self.connection.session.session_id,
-                max_rows=self.arraysize,
-                max_bytes=self.buffer_size_bytes,
-                lz4_compression=self.connection.lz4_compression,
-                cursor=self,
-                use_cloud_fetch=self.connection.use_cloud_fetch,
-                parameters=prepared_params,
-                async_op=False,
-                enforce_embedded_schema_correctness=enforce_embedded_schema_correctness,
+        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,
+                input_stream=input_stream
             )
 
-            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
-                )
-
-            return self
-        finally:
-            # Clean up stream data
-            self._input_stream_data = None
+        return self
 
     @log_latency(StatementType.QUERY)
     def execute_async(