Complete Volume file operations through the kernel backend

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Release History
 
 # Unreleased
+- Support Volume `PUT`, `GET` and `REMOVE` on the kernel backend, including existing Python stream uploads. Local file paths and resolved symlink targets must stay within configured allowed roots.
 - Fix the kernel backend ignoring `use_cloud_fetch=False`, so complete results can be fetched without cloud-storage downloads.
 - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629))
 - Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most.

CONNECTION_PARAMETERS.md

@@ -171,7 +171,7 @@ to change without notice.
 | `use_inline_params`           | `bool` \| `"silent"`              |   ✅   |   ✅   | `False`       | Render parameters inline (legacy) vs. native bound params (DBR 14.1+). Shared cursor code performs the rendering before dispatch, so both backends honor it identically. |
 | `ignore_transactions`         | `bool`                            |   ✅   |   ✅   | `True`        | When `True`: `commit()` is a no-op, `rollback()` raises `NotSupportedError`, and setting `autocommit` is a no-op.          |
 | `fetch_autocommit_from_server`| `bool`                            |   ✅   |   ✅   | `False`       | Query the server (`SET AUTOCOMMIT`) for autocommit state instead of returning the cached value.                           |
-| `staging_allowed_local_path`  | `str` \| `List[str]`              |   ✅   |   ❌   | `None`        | Local path(s) permitted for Unity Catalog Volume `PUT`/`GET`. **Thrift-only** — the kernel has no Volume API yet.          |
+| `staging_allowed_local_path`  | `str` \| `List[str]`              |   ✅   |   ✅   | `None`        | Local path(s) permitted for Unity Catalog Volume `PUT`/`GET`. Kernel also checks resolved symlink targets; see below. |
 
 ## Telemetry
 
@@ -208,13 +208,12 @@ are forwarded when supported by the installed kernel binding.
    `enable_query_result_lz4_compression`.
 5. Arrow-native rendering for `_use_arrow_native_decimals` /
    `_use_arrow_native_timestamps` (complex types **are** forwarded).
-6. `staging_allowed_local_path` (Volume `PUT`/`GET`).
-7. Retry fine-tuning beyond the four forwarded knobs: `_retry_delay_default`,
+6. Retry fine-tuning beyond the four forwarded knobs: `_retry_delay_default`,
    `_retry_dangerous_codes`, `_respect_server_retry_after_header`,
    `_retry_max_redirects`, `_enable_v3_retries`.
-8. `force_enable_telemetry`.
-9. `_port` and `_connection_uri`.
-10. `_pool_connections` and `_proxy_auth_method` do not configure the
+7. `force_enable_telemetry`.
+8. `_port` and `_connection_uri`.
+9. `_pool_connections` and `_proxy_auth_method` do not configure the
     kernel-owned network stack.
 
 ### Supported on Kernel, missing / ignored on Thrift
@@ -246,3 +245,17 @@ are forwarded when supported by the installed kernel binding.
 > subject to change. Generated against connector version `4.5.0`; the source of
 > truth is `Connection.__init__` (`src/databricks/sql/client.py`) and the
 > per-backend clients under `src/databricks/sql/backend/`.
+
+### Volume file paths
+
+With `use_kernel=True`, SQL `PUT`, `GET` and `REMOVE` complete their file
+operations before `execute()` or `get_async_execution_result()` returns.
+`PUT` and `GET` require paths within
+`staging_allowed_local_path`. Both the requested path and its resolved target
+must lie within a configured root. Symlinks within those roots are permitted;
+symlinks leading outside are rejected. This is stricter than the Thrift
+backend, which checks the requested path only.
+
+`execute(input_stream=stream)` retains its existing behavior: the stream is
+read only when the server requests stream input. `REMOVE` requires no local
+root. A completed transfer returns no instruction rows.

KERNEL_REV

@@ -1 +1 @@
-0b5d0edfed32c0745d5abc2468fc3070e2eb2a17
+855847bdca71e07285cfc60332fb2f33f8e47786

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

@@ -16,8 +16,6 @@
 
 - ``query_tags`` on execute is not supported (kernel exposes
   ``statement_conf`` but PyO3 doesn't surface it).
-- Volume PUT/GET (staging operations): kernel has no Volume API
-  yet. Users on Thrift-only paths.
 """
 
 from __future__ import annotations
@@ -35,6 +33,7 @@
     Optional,
     Set,
     TYPE_CHECKING,
+    Tuple,
     Union,
 )
 
@@ -84,41 +83,6 @@
 # per-request skip-and-warn.
 _KERNEL_MANAGED_HEADERS = frozenset({"authorization", "x-databricks-org-id"})
 
-# Leading verbs of SQL volume/staging statements. Detected by the
-# leading token (case-insensitive) so the kernel backend can fail loud
-# on staging ops it can't service — see ``execute_command``.
-_STAGING_VERBS = ("PUT", "GET", "REMOVE")
-
-
-def _strip_leading_sql_comments(sql: str) -> str:
-    """Strip leading whitespace and SQL comments (``-- …`` line and
-    ``/* … */`` block, possibly several) from ``sql``, returning the
-    remainder.
-
-    Needed so staging detection sees the real leading verb: a
-    comment-prefixed staging op (``-- upload\\nPUT …`` or
-    ``/* c */ PUT …``, common in ETL scripts) must still be classified
-    as staging, or it would slip past the guard into the silent-no-op
-    bug. Block comments do not nest in Databricks SQL, so a simple
-    scan-to-``*/`` is correct.
-    """
-    i = 0
-    n = len(sql)
-    while i < n:
-        if sql[i].isspace():
-            i += 1
-        elif sql.startswith("--", i):
-            # Line comment: skip to end of line (or string).
-            nl = sql.find("\n", i)
-            i = n if nl == -1 else nl + 1
-        elif sql.startswith("/*", i):
-            # Block comment: skip to closing */ (or end if unterminated).
-            close = sql.find("*/", i + 2)
-            i = n if close == -1 else close + 2
-        else:
-            break
-    return sql[i:]
-
 
 def _is_not_found(exc: BaseException) -> bool:
     """True iff ``exc`` is a kernel ``NotFound`` error (HTTP 404 /
@@ -155,21 +119,6 @@ def _catalog_or_none(value: Optional[str]) -> Optional[str]:
     return value
 
 
-def _is_staging_statement(operation: str) -> bool:
-    """True iff ``operation`` is a volume/staging statement (PUT / GET /
-    REMOVE).
-
-    Strips leading whitespace + SQL comments first (so a comment-
-    prefixed staging op is still caught), then matches the leading token
-    only — so a normal query that merely *contains* the word (e.g.
-    ``SELECT 'GET' AS x``) isn't misflagged.
-    """
-    stripped = _strip_leading_sql_comments(operation)
-    # First whitespace-delimited token, uppercased.
-    verb = stripped.split(None, 1)[0].upper() if stripped.strip() else ""
-    return verb in _STAGING_VERBS
-
-
 class _KernelSessionParameterSupport(NamedTuple):
     session_type: Any
     parameters: FrozenSet[str]
@@ -436,6 +385,8 @@ def __init__(
         # Guarded by its own lock — cancel can race execute teardown.
         self._sync_cancellers: Dict[int, Any] = {}
         self._sync_cancellers_lock = threading.RLock()
+        # Retain the mutable owner until its transfer completes; dropping it cancels the result.
+        self._volume_handles: Dict[str, Tuple[Any, Any]] = {}
 
     # ── Session lifecycle ──────────────────────────────────────────
 
@@ -565,6 +516,21 @@ def open_session(
     def close_session(self, session_id: SessionId) -> None:
         if self._kernel_session is None:
             return
+        with self._sync_cancellers_lock:
+            cancellers = list(self._sync_cancellers.values()) + [
+                entry[1]
+                for entry in self._volume_handles.values()
+                if entry[1] is not None
+            ]
+            self._sync_cancellers.clear()
+            self._volume_handles.clear()
+        for canceller in cancellers:
+            try:
+                canceller.cancel()
+            except Exception:
+                logger.debug(
+                    "Statement cancellation during session close failed", exc_info=True
+                )
         # Close any tracked async handles first so they fire their
         # server-side CloseStatement before the session goes away.
         with self._async_handles_lock:
@@ -633,24 +599,6 @@ def execute_command(
         # ``_async_statements`` and closed by ``close_command``); the sync
         # path drops it in finally. ``close_stmt`` is the post-success
         # decision flag — it stays True on sync, flips to False on async.
-        # Volume/staging (PUT/GET/REMOVE) is not supported on the kernel
-        # path: the kernel returns the staging control row as a normal
-        # result set (``KernelResultSet.is_staging_operation`` is always
-        # False), so the connector's ``_handle_staging_operation`` never
-        # fires and NO file is transferred. Rather than silently no-op
-        # (the Thrift path performs the presigned-URL upload/download),
-        # fail loud at the call site so ETL scripts don't ingest
-        # stale/missing data. Detected by the leading SQL verb — the
-        # only signal available pre-execute, since the kernel exposes no
-        # staging marker today.
-        if _is_staging_statement(operation):
-            raise NotSupportedError(
-                "Volume / staging operations (PUT / GET / REMOVE) are not "
-                "supported on the kernel backend (use_kernel=True); the file "
-                "transfer would silently not happen. Use the Thrift backend "
-                "for staging operations."
-            )
-
         close_stmt = True
         try:
             try:
@@ -692,6 +640,13 @@ def execute_command(
                     # Canceller is best-effort; never block execute on it.
                     pass
                 executed = stmt.execute()
+                if executed.is_volume_operation:
+                    command_id = CommandId.from_sea_statement_id(executed.statement_id)
+                    with self._sync_cancellers_lock:
+                        self._volume_handles[command_id.guid] = (
+                            stmt,
+                            self._sync_cancellers.get(id(cursor)),
+                        )
                 # Execute succeeded: the kernel now owns the statement
                 # lifecycle. It auto-closes the server statement when the
                 # result stream is fully drained (``ExecutedStatement::
@@ -764,16 +719,22 @@ def execute_command(
         try:
             return self._make_result_set(executed, cursor, command_id)
         except Exception as exc:
+            self.finish_volume_operation(command_id)
             raise _wrap_kernel_exception("execute_command", exc) from exc
 
+    def finish_volume_operation(self, command_id: CommandId) -> None:
+        with self._sync_cancellers_lock:
+            self._volume_handles.pop(command_id.guid, None)
+
     def cancel_command(self, command_id: CommandId) -> None:
-        with self._async_handles_lock:
-            handle = self._async_handles.get(command_id.guid)
+        with self._sync_cancellers_lock:
+            volume = self._volume_handles.get(command_id.guid)
+        handle = volume[1] if volume is not None else None
         if handle is None:
-            # Sync-execute paths fully materialise the result before
-            # ``execute_command`` returns, so by the time
-            # cancel_command can fire there's nothing in flight.
-            # Match the Thrift backend's tolerant behaviour.
+            with self._async_handles_lock:
+                handle = self._async_handles.get(command_id.guid)
+        if handle is None:
+            # Completed synchronous operations have no retained canceller.
             logger.debug("cancel_command: no in-flight async handle for %s", command_id)
             return
         try:
@@ -822,6 +783,13 @@ def cancel_running_cursor(self, cursor: "Cursor") -> bool:
         return True
 
     def close_command(self, command_id: CommandId) -> None:
+        with self._sync_cancellers_lock:
+            volume = self._volume_handles.pop(command_id.guid, None)
+        if volume is not None and volume[1] is not None:
+            try:
+                volume[1].cancel()
+            except Exception:
+                logger.debug("Volume cancellation during close failed", exc_info=True)
         with self._async_handles_lock:
             handle = self._async_handles.pop(command_id.guid, None)
             stmt = self._async_statements.pop(command_id.guid, None)

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

@@ -31,15 +31,17 @@
 from __future__ import annotations
 
 import logging
+import os
 from collections import deque
-from typing import Any, Deque, List, Optional, TYPE_CHECKING, cast
+from typing import Any, BinaryIO, Deque, List, Optional, TYPE_CHECKING, Union, cast
 
 import pyarrow
 
 from databricks.sql.backend.kernel._errors import wrap_kernel_exception
 from databricks.sql.backend.kernel.type_mapping import description_from_arrow_schema
 from databricks.sql.backend.types import CommandId, CommandState
 from databricks.sql.result_set import ResultSet
+from databricks.sql.exc import OperationalError, ProgrammingError
 from databricks.sql.types import Row
 
 if TYPE_CHECKING:
@@ -83,7 +85,7 @@ def __init__(
             has_more_rows=True,
             results_queue=None,
             description=description_from_arrow_schema(schema),
-            is_staging_operation=False,
+            is_staging_operation=kernel_handle.is_volume_operation,
             lz4_compressed=False,
             arrow_schema_bytes=None,
         )
@@ -101,6 +103,72 @@ def __init__(
         self._buffered_count: int = 0
         self._exhausted: bool = False
 
+    def complete_staging_operation(
+        self,
+        allowed_paths: Union[None, str, List[str]],
+        input_stream: Optional[BinaryIO] = None,
+    ) -> None:
+        paths = (
+            [allowed_paths] if isinstance(allowed_paths, str) else (allowed_paths or [])
+        )
+        consume_instructions = self.is_staging_operation
+        try:
+            outcome = self._complete_volume_operation(paths)
+            if outcome == "not_volume":
+                consume_instructions = False
+                return
+            consume_instructions = True
+            if outcome == "needs_input_stream":
+                try:
+                    if input_stream is None:
+                        raise ProgrammingError(
+                            "No input stream provided for streaming operation"
+                        )
+                    data = input_stream.read()
+                except BaseException:
+                    self.close()
+                    raise
+                outcome = self._complete_volume_operation(paths, data)
+            if outcome != "completed":
+                raise ProgrammingError("Volume operation did not complete")
+            self.status = CommandState.SUCCEEDED
+        finally:
+            if consume_instructions:
+                self._buffer.clear()
+                self._buffer_offset = 0
+                self._buffered_count = 0
+                self._exhausted = True
+                self.has_more_rows = False
+            cast("KernelDatabricksClient", self.backend).finish_volume_operation(
+                self.command_id
+            )
+
+    def _complete_volume_operation(
+        self, paths: List[str], input_bytes: Optional[bytes] = None
+    ) -> str:
+        try:
+            return self._kernel_handle.complete_volume_operation(
+                paths, input_bytes=input_bytes
+            )
+        except Exception as exc:
+            os_code = getattr(exc, "local_os_error", None)
+            if os_code is not None:
+                local_path = getattr(exc, "local_path", None)
+                if os.name == "nt":
+                    raise OSError(
+                        0, "Local file operation failed", local_path, os_code
+                    ) from exc
+                raise OSError(os_code, os.strerror(os_code), local_path) from exc
+            http_status = getattr(exc, "http_status", None)
+            if (
+                http_status is not None
+                and getattr(exc, "error_code", None) == "VOLUME_STORAGE_HTTP_ERROR"
+            ):
+                raise OperationalError(
+                    f"Staging operation over HTTP was unsuccessful: {http_status}"
+                ) from exc
+            raise wrap_kernel_exception("complete_volume_operation", exc) from exc
+
     # ----- internal helpers -----
 
     def _pull_one_batch(self) -> bool:
@@ -244,19 +312,20 @@ def close(self) -> None:
             self.has_been_closed_server_side = True
             self.status = CommandState.CLOSED
             return
+        backend = cast("KernelDatabricksClient", self.backend)
+        if self.is_staging_operation and self.status != CommandState.SUCCEEDED:
+            # Cancellation must precede waiting for a transfer's native lock.
+            try:
+                backend.cancel_command(self.command_id)
+            except Exception:
+                logger.debug("Volume cancellation during close failed", exc_info=True)
         try:
             self._kernel_handle.close()
         except Exception as exc:
             # close() failures are not actionable at the connector
             # level; log and swallow so the cursor's __del__ /
             # connection close path stays clean.
             logger.warning("Error closing kernel handle: %s", exc)
-        # Honor the base ``ResultSet`` contract: notify the backend.
-        # For async results, ``backend.close_command`` drops the
-        # retained owning handle and parent Statement. Sync-execute and
-        # metadata paths never registered in ``_async_handles`` to begin
-        # with, so this call is tolerant bookkeeping for them.
-        backend = cast("KernelDatabricksClient", self.backend)
         try:
             backend.close_command(self.command_id)
         except Exception as exc:

src/databricks/sql/client.py

@@ -1284,6 +1284,9 @@ def _handle_staging_operation(
         """
 
         assert self.active_result_set is not None
+        complete = getattr(self.active_result_set, "complete_staging_operation", None)
+        if complete is not None:
+            return complete(staging_allowed_local_path, input_stream)
         row = self.active_result_set.fetchone()
         assert row is not None