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

@@ -114,7 +114,7 @@ to change without notice.
 
 | Option                               | Type        | Thrift | Kernel | Default Value | Note                                                                                                                                            |
 | ------------------------------------ | ----------- | :----: | :----: | ------------- | ---------------------------------------------------------------------------------------------------------------------------------------------- |
-| `_socket_timeout`                    | `float` (s) |   ✅   |   ✅   | `900` (Thrift); `120` (kernel) | Thrift: socket send/recv/connect timeout. Kernel: total HTTP request deadline from connect through response-body completion. A positive value is forwarded; unset or `0` selects the kernel's 120s default. On the kernel path, `0` is neither unlimited nor an immediate timeout. |
+| `_socket_timeout`                    | `float` (s) |   ✅   |   ✅   | `900` (Thrift); `120` (kernel) | Thrift: socket send/recv/connect timeout. Kernel: total HTTP request deadline from connect through response-body completion. A positive value is forwarded; unset or `0` selects the kernel's 120s SQL request default. Volume transfers have no total deadline: an explicit positive value caps storage connection establishment, GET/REMOVE response-header waits and all response-body reads. PUT response-header waits remain unbounded. Unset or `0` leaves storage inactivity unbounded. |
 | `_pool_connections`                  | `int`       |   ⚠️   |   ❌   | `10`          | Number of pools in the shared Python HTTP client. On Thrift this affects auxiliary auth/feature-flag/telemetry traffic, not the separate Thrift query transport. Kernel-owned network traffic does not use this client. |
 | `_pool_maxsize`                      | `int`       |   ⚠️   |   ⚠️   | `20` (Python); `100` (kernel when unset) | Max idle connections retained per host. It configures the shared Python client but not the Thrift query transport. With kernel ≥ 1.1.0, a positive value also configures the kernel's Rust HTTP pool. Unset or `0` keeps each client's default. |
 | `_proxy_auth_method`                 | `str`       |   ✅   |   ❌   | `None`        | `basic` or `negotiate` (Kerberos). Applied to both Thrift query traffic and the shared Python client. It is not forwarded to the kernel's Rust transport. See [`docs/proxy.md`](docs/proxy.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
+566fe7f0a8a2fa246db95045a7000136ee6070d5

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)