Reject decoded streaming PUT sources in native execution

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>

Author: aakash-saravanan-db

src/databricks/sql/backend/kernel/_errors.py

@@ -48,6 +48,7 @@
 from databricks.sql.exc import (
     DatabaseError,
     Error,
+    NotSupportedError,
     OperationalError,
     ProgrammingError,
     ServerOperationError,
@@ -135,6 +136,8 @@ def reraise_kernel_error(
     code = getattr(exc, "code", "Unknown")
     cls = _CODE_TO_EXCEPTION.get(code, DatabaseError)
     error_code = getattr(exc, "error_code", None)
+    if error_code == "FILE_TRANSFER_UNSUPPORTED_STREAM":
+        cls = NotSupportedError
     file_transfer = error_code in (
         "FILE_TRANSFER_ERROR",
         "FILE_TRANSFER_PERMISSION_ERROR",

src/databricks/sql/backend/kernel/client.py

@@ -572,13 +572,6 @@ def execute_command(
             try:
                 stripped_sql = _strip_leading_sql_comments(operation)
                 staging_verb = _staging_verb(stripped_sql)
-                if staging_verb == "PUT" and re.match(
-                    r"""(['"])__input_stream__\1""",
-                    _strip_leading_sql_comments(stripped_sql[3:]),
-                ):
-                    raise NotSupportedError(
-                        "Streaming PUT is not supported by the kernel backend."
-                    )
                 stmt.set_sql(operation)
                 allowed_paths = (
                     cursor.connection.staging_allowed_local_path
@@ -601,7 +594,7 @@ def execute_command(
                             raise TypeError("Local paths must resolve to strings")
                     except (TypeError, ValueError) as exc:
                         raise ProgrammingError(_STAGING_PATHS_TYPE_ERROR) from exc
-                    set_paths(normalized_paths)
+                    set_paths(normalized_paths, input_stream_token="__input_stream__")
                 elif staging_verb is not None:
                     raise NotSupportedError(
                         "Volume transfers require a databricks-sql-kernel wheel "

tests/e2e/test_kernel_backend.py

@@ -817,6 +817,8 @@ def execute(query):
                     execute(f"-- upload\nPUT '{source}' INTO '{remote}' OVERWRITE")
                     == []
                 )
+                with pytest.raises(NotSupportedError, match="Streaming PUT"):
+                    execute(f"PUT r'__input_stream__' INTO '{remote}' OVERWRITE")
                 assert execute(f"GET '{remote}' TO '{destination}'") == []
                 assert destination.read_bytes() == payload
                 with pytest.raises(ProgrammingError):

tests/unit/test_kernel_client.py

@@ -862,7 +862,7 @@ def test_execute_command_forwards_allowed_local_paths(
         enforce_embedded_schema_correctness=False,
     )
     stmt.set_staging_allowed_local_paths.assert_called_once_with(
-        expected if uses_paths else []
+        expected if uses_paths else [], input_stream_token="__input_stream__"
     )
     (stmt.submit if async_op else stmt.execute).assert_called_once()
 
@@ -912,7 +912,9 @@ def execute():
         stmt.set_staging_allowed_local_paths.assert_not_called()
     else:
         execute()
-        stmt.set_staging_allowed_local_paths.assert_called_once_with([])
+        stmt.set_staging_allowed_local_paths.assert_called_once_with(
+            [], input_stream_token="__input_stream__"
+        )
         stmt.execute.assert_called_once()
 
 
@@ -947,30 +949,34 @@ def test_staging_verb_matches_only_the_leading_statement(operation, expected):
     assert kernel_client._staging_verb(stripped) == expected
 
 
-@pytest.mark.parametrize(
-    "query",
-    [
-        "PUT '__input_stream__' INTO '/Volumes/c/s/v/file'",
-        "put '__input_stream__' INTO '/Volumes/c/s/v/file'",
-        "PUT '__input_stream__'INTO '/Volumes/c/s/v/file'",
-        'PUT"__input_stream__"INTO "/Volumes/c/s/v/file"',
-        "/* upload */ PUT '__input_stream__' INTO '/Volumes/c/s/v/file'",
-        "/* outer /* inner */ outer */ PUT '__input_stream__' INTO '/Volumes/c/s/v/file'",
-        "PUT'__input_stream__' INTO '/Volumes/c/s/v/file'",
-        "PUT /* upload */ '__input_stream__' INTO '/Volumes/c/s/v/file'",
-        "PUT\n-- upload\n'__input_stream__' INTO '/Volumes/c/s/v/file'",
-        "/* upload */ PUT\"__input_stream__\" INTO '/Volumes/c/s/v/file'",
-    ],
-)
-def test_kernel_rejects_streaming_put(query):
+def test_kernel_rejects_decoded_streaming_put():
     c = _make_client()
     c._kernel_session = MagicMock()
+    stmt = c._kernel_session.statement.return_value
+    error = _FakeKernelError(
+        code="InvalidArgument", message="Streaming PUT is not supported"
+    )
+    error.error_code = "FILE_TRANSFER_UNSUPPORTED_STREAM"
+    stmt.execute.side_effect = error
+    cursor = MagicMock()
+    cursor.connection.staging_allowed_local_path = None
     with pytest.raises(NotSupportedError, match="Streaming PUT"):
         c.execute_command(
-            query, MagicMock(), 1, 1, False, MagicMock(), False, [], False, False
+            "PUT r'__input_stream__' INTO '/Volumes/c/s/v/file'",
+            MagicMock(),
+            1,
+            1,
+            False,
+            cursor,
+            False,
+            [],
+            False,
+            False,
         )
-    stmt = c._kernel_session.statement.return_value
-    stmt.execute.assert_not_called()
+    stmt.set_staging_allowed_local_paths.assert_called_once_with(
+        [], input_stream_token="__input_stream__"
+    )
+    stmt.execute.assert_called_once()
     stmt.close.assert_called_once()