From 3a92da2bc044242db0e7b9b2b7d288c44c6186a4 Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Fri, 13 Nov 2026 22:15:32 +0000
Subject: [PATCH 01/15] Complete Volume file operations through the native
 kernel

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 CHANGELOG.md                                  |   1 +
 CONNECTION_PARAMETERS.md                      |  24 +++-
 src/databricks/sql/backend/kernel/client.py   | 118 +++++++-----------
 .../sql/backend/kernel/result_set.py          |  84 ++++++++++---
 src/databricks/sql/client.py                  |   3 +
 tests/e2e/test_kernel_backend.py              |  25 ----
 tests/unit/test_client.py                     |   8 ++
 tests/unit/test_kernel_client.py              |  96 +++++---------
 tests/unit/test_kernel_result_set.py          |  90 +++++++++++++
 tests/unit/test_staging_remove.py             |   7 +-
 tests/unit/test_streaming_put.py              |   3 +-
 11 files changed, 268 insertions(+), 191 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2bae4c54e..64cb4ea6a 100644
--- a/CHANGELOG.md
+++ b/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.
diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md
index c15b4d9d4..046385a90 100644
--- a/CONNECTION_PARAMETERS.md
+++ b/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,16 @@ 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 execution 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.
diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py
index df52f34a5..d2909ecf2 100644
--- a/src/databricks/sql/backend/kernel/client.py
+++ b/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,19 @@ 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()
+            ]
+            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 +597,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 +638,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[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 +717,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 +781,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:
+            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)
diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py
index 15f78bf5f..67d63769d 100644
--- a/src/databricks/sql/backend/kernel/result_set.py
+++ b/src/databricks/sql/backend/kernel/result_set.py
@@ -31,8 +31,9 @@
 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
 
@@ -40,6 +41,7 @@
 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,64 @@ 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 [])
+        )
+        try:
+            outcome = self._complete_volume_operation(paths)
+            if outcome == "not_volume":
+                return
+            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._buffer.clear()
+            self._buffered_count = 0
+            self._exhausted = True
+            self.has_more_rows = False
+            self.status = CommandState.SUCCEEDED
+        finally:
+            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:
+                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,18 +304,7 @@ def close(self) -> None:
             self.has_been_closed_server_side = True
             self.status = CommandState.CLOSED
             return
-        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.
+        # Signal through the detached canceller before close waits for a transfer's native lock.
         backend = cast("KernelDatabricksClient", self.backend)
         try:
             backend.close_command(self.command_id)
@@ -263,6 +312,13 @@ def close(self) -> None:
             logger.warning(
                 "backend.close_command from result-set close failed: %s", exc
             )
+        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)
         self._buffer.clear()
         self._buffered_count = 0
         self._kernel_handle = None
diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
index 613b62386..753cb225a 100755
--- a/src/databricks/sql/client.py
+++ b/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
 
diff --git a/tests/e2e/test_kernel_backend.py b/tests/e2e/test_kernel_backend.py
index 260c49689..532187852 100644
--- a/tests/e2e/test_kernel_backend.py
+++ b/tests/e2e/test_kernel_backend.py
@@ -35,7 +35,6 @@
 from databricks.sql.backend.types import CommandState
 from databricks.sql.exc import (
     DatabaseError,
-    NotSupportedError,
     OperationalError,
     ServerOperationError,
 )
@@ -863,30 +862,6 @@ def test_scientific_notation_decimal_parameter_binds(conn):
         assert int(cur.fetchall()[0][0]) == 100
 
 
-# ── Staging / volume — fail loud, not silent no-op ────────────────
-
-
-def test_staging_put_raises_not_supported(conn):
-    """A PUT (volume/staging) statement fails loud on the kernel path
-    rather than silently no-opping (which would make ETL ingest
-    stale/missing data). (CUJ-gap audit)"""
-    with conn.cursor() as cur:
-        with pytest.raises(NotSupportedError, match="staging"):
-            cur.execute("PUT '/tmp/x.csv' INTO '/Volumes/main/default/v/x.csv'")
-
-
-def test_comment_prefixed_staging_put_raises_not_supported(conn):
-    """A comment-prefixed staging op (common in ETL scripts) must also
-    fail loud — the leading-verb detection strips SQL comments first, so
-    it can't slip past into the silent-no-op bug (PR #825 review #1)."""
-    with conn.cursor() as cur:
-        with pytest.raises(NotSupportedError, match="staging"):
-            cur.execute(
-                "-- upload the daily extract\n"
-                "PUT '/tmp/x.csv' INTO '/Volumes/main/default/v/x.csv'"
-            )
-
-
 # ── Error fidelity — diagnostic-info reaches .context ─────────────
 
 
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
index 047252488..4cb81a329 100644
--- a/tests/unit/test_client.py
+++ b/tests/unit/test_client.py
@@ -409,6 +409,14 @@ def test_get_columns_parameters_passed_to_thrift_backend(self, mock_thrift_backe
                 for k, v in req_args.items():
                     self.assertEqual(v, call_args[k])
 
+    def test_staging_uses_native_completion_without_fetching_instructions(self):
+        cursor = client.Cursor(Mock(), Mock())
+        result = Mock(spec=["complete_staging_operation"])
+        cursor.active_result_set = result
+        stream = Mock()
+        cursor._handle_staging_operation("/allowed", stream)
+        result.complete_staging_operation.assert_called_once_with("/allowed", stream)
+
     def test_cancel_command_calls_the_backend(self):
         mock_thrift_backend = Mock()
         cursor = client.Cursor(Mock(), mock_thrift_backend)
diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py
index f94aef0c5..2e63d9798 100644
--- a/tests/unit/test_kernel_client.py
+++ b/tests/unit/test_kernel_client.py
@@ -1006,72 +1006,6 @@ def test_execute_command_forwards_row_limit(row_limit):
     stmt.set_row_limit.assert_called_once_with(row_limit)
 
 
-# ---------------------------------------------------------------------------
-# Staging / volume operations — fail loud (not silently no-op)
-# ---------------------------------------------------------------------------
-
-
-@pytest.mark.parametrize(
-    "operation",
-    [
-        "PUT '/local/f.csv' INTO '/Volumes/c/s/v/f.csv'",
-        "  put '/local/f' into '/Volumes/...'",  # leading ws + lowercase
-        "GET '/Volumes/c/s/v/f' TO '/local/f'",
-        "REMOVE '/Volumes/c/s/v/f'",
-    ],
-)
-def test_staging_operation_raises_not_supported(operation):
-    """Volume/staging PUT/GET/REMOVE must FAIL LOUD on the kernel path
-    (the kernel can't perform the presigned-URL transfer; silently
-    no-opping would make ETL ingest stale/missing data)."""
-    c = _make_client()
-    c._kernel_session = MagicMock()
-    cursor = MagicMock()
-    cursor.arraysize = 100
-    cursor.buffer_size_bytes = 1024
-
-    with pytest.raises(NotSupportedError, match="staging"):
-        c.execute_command(
-            operation=operation,
-            session_id=MagicMock(),
-            max_rows=1,
-            max_bytes=1,
-            lz4_compression=False,
-            cursor=cursor,
-            use_cloud_fetch=False,
-            parameters=[],
-            async_op=False,
-            enforce_embedded_schema_correctness=False,
-        )
-
-
-@pytest.mark.parametrize(
-    "operation, is_staging",
-    [
-        ("PUT '/f' INTO '/v'", True),
-        ("get '/v' to '/f'", True),
-        ("REMOVE '/v'", True),
-        # Comment-prefixed staging ops MUST still be caught — otherwise
-        # they slip into the silent-no-op bug this guard exists to close
-        # (regression: review #1 on PR #825). ETL scripts commonly
-        # prefix statements with comments.
-        ("-- upload the file\nPUT '/f' INTO '/v'", True),
-        ("/* staging */ PUT '/f' INTO '/v'", True),
-        ("/* c1 */\n  -- c2\n  get '/v' to '/f'", True),  # mixed, multiple
-        ("   \n\t PUT '/f' INTO '/v'", True),  # leading whitespace only
-        ("SELECT 'GET' AS x", False),  # word appears but not leading verb
-        ("SELECT * FROM puts", False),
-        ("-- PUT in a comment\nSELECT 1", False),  # verb only in comment
-        ("/* PUT */ SELECT 1", False),
-        ("INSERT INTO t VALUES (1)", False),
-        ("", False),
-        ("-- just a comment", False),  # comment only, no statement
-    ],
-)
-def test_is_staging_statement(operation, is_staging):
-    assert kernel_client._is_staging_statement(operation) is is_staging
-
-
 # ---------------------------------------------------------------------------
 # Sync cancel wiring (cursor.cancel() during a blocking execute())
 # ---------------------------------------------------------------------------
@@ -1140,6 +1074,7 @@ def fake_execute():
             )
         return MagicMock(
             statement_id="stmt-id",
+            is_volume_operation=False,
             arrow_schema=MagicMock(return_value=pa.schema([("x", pa.int64())])),
         )
 
@@ -2781,3 +2716,32 @@ def test_only_reserved_headers_omits_kwarg(self, monkeypatch):
             monkeypatch, [("Authorization", "Bearer x"), ("x-databricks-org-id", "1")]
         )
         assert "http_headers" not in captured
+
+
+@pytest.mark.parametrize("is_volume", [False, True])
+def test_sync_volume_canceller_survives_until_completion(is_volume):
+    c = _make_client()
+    c._kernel_session = MagicMock()
+    stmt = c._kernel_session.statement.return_value
+    stmt.execute.return_value = MagicMock(
+        statement_id="volume-id",
+        is_volume_operation=is_volume,
+        arrow_schema=MagicMock(return_value=pa.schema([("x", pa.int64())])),
+    )
+    result = c.execute_command(
+        operation="/* command */ PUT '/local/f' INTO '/Volumes/c/s/v/f'",
+        session_id=MagicMock(),
+        max_rows=1,
+        max_bytes=1,
+        lz4_compression=False,
+        cursor=MagicMock(arraysize=100, buffer_size_bytes=1024),
+        use_cloud_fetch=False,
+        parameters=[],
+        async_op=False,
+        enforce_embedded_schema_correctness=False,
+    )
+    assert result.is_staging_operation is is_volume
+    c.cancel_command(result.command_id)
+    assert stmt.canceller.return_value.cancel.called is is_volume
+    c.finish_volume_operation(result.command_id)
+    assert not c._volume_handles
diff --git a/tests/unit/test_kernel_result_set.py b/tests/unit/test_kernel_result_set.py
index cfa2d3f2e..8c7a3dbb5 100644
--- a/tests/unit/test_kernel_result_set.py
+++ b/tests/unit/test_kernel_result_set.py
@@ -5,6 +5,8 @@
 from __future__ import annotations
 
 from collections import deque
+import errno
+import io
 from typing import Deque
 from unittest.mock import MagicMock
 
@@ -28,6 +30,7 @@ def __init__(self, schema: pa.Schema, batches):
         self._schema = schema
         self._batches: Deque[pa.RecordBatch] = deque(batches)
         self.closed = False
+        self.is_volume_operation = False
 
     def arrow_schema(self) -> pa.Schema:
         return self._schema
@@ -243,3 +246,90 @@ def test_close_skips_kernel_call_when_connection_already_closed(int_schema):
     # …but local state is still terminal so __del__ is idempotent:
     assert rs.status == CommandState.CLOSED
     assert rs.has_been_closed_server_side is True
+
+
+@pytest.mark.parametrize("data", [b"payload", b""])
+def test_volume_reads_input_only_when_native_requests_it(int_schema, data):
+    handle = _FakeKernelHandle(int_schema, [_batch(int_schema, [1])])
+    handle.is_volume_operation = True
+    handle.complete_volume_operation = MagicMock(
+        side_effect=["needs_input_stream", "completed"]
+    )
+    rs = _make_rs(handle)
+    stream = io.BytesIO(data)
+    rs.complete_staging_operation("/allowed", stream)
+    assert stream.tell() == len(data)
+    assert handle.complete_volume_operation.call_args_list == [
+        ((["/allowed"],), {"input_bytes": None}),
+        ((["/allowed"],), {"input_bytes": data}),
+    ]
+    assert rs.fetchall() == []
+    rs.backend.finish_volume_operation.assert_called_once_with(rs.command_id)
+
+
+@pytest.mark.parametrize("outcome", ["completed", "not_volume"])
+def test_volume_does_not_read_irrelevant_stream(int_schema, outcome):
+    handle = _FakeKernelHandle(int_schema, [_batch(int_schema, [1])])
+    handle.complete_volume_operation = MagicMock(return_value=outcome)
+    rs = _make_rs(handle)
+    stream = MagicMock()
+    rs.complete_staging_operation(None, stream)
+    stream.read.assert_not_called()
+    assert len(rs.fetchall()) == (1 if outcome == "not_volume" else 0)
+
+
+@pytest.mark.parametrize("missing", [False, True])
+def test_volume_missing_or_failed_input_closes_without_masking_error(
+    int_schema, missing
+):
+    from databricks.sql.exc import ProgrammingError
+
+    handle = _FakeKernelHandle(int_schema, [])
+    handle.complete_volume_operation = MagicMock(return_value="needs_input_stream")
+    rs = _make_rs(handle)
+    error = RuntimeError("stream read failed")
+    stream = MagicMock()
+    stream.read.side_effect = error
+    with pytest.raises(ProgrammingError if missing else RuntimeError) as raised:
+        rs.complete_staging_operation([], None if missing else stream)
+    if not missing:
+        assert raised.value is error
+    assert handle.closed
+    handle.complete_volume_operation.assert_called_once()
+
+
+def test_volume_preserves_local_os_error_and_filename(int_schema):
+    handle = _FakeKernelHandle(int_schema, [])
+    error = RuntimeError("local file failure")
+    error.local_os_error = errno.ENOENT
+    error.local_path = "/allowed/missing"
+    handle.complete_volume_operation = MagicMock(side_effect=error)
+    rs = _make_rs(handle)
+    with pytest.raises(FileNotFoundError) as raised:
+        rs.complete_staging_operation(["/allowed"])
+    assert raised.value.errno == errno.ENOENT
+    assert raised.value.filename == "/allowed/missing"
+
+
+def test_volume_close_signals_backend_before_native_lock(int_schema):
+    handle = _FakeKernelHandle(int_schema, [])
+    rs = _make_rs(handle)
+    handle.close = MagicMock(
+        side_effect=lambda: rs.backend.close_command.assert_called_once_with(
+            rs.command_id
+        )
+    )
+    rs.close()
+    handle.close.assert_called_once()
+
+
+def test_volume_http_404_is_an_operational_error(int_schema):
+    from databricks.sql.exc import OperationalError
+
+    handle = _FakeKernelHandle(int_schema, [])
+    error = RuntimeError("storage response")
+    error.http_status = 404
+    handle.complete_volume_operation = MagicMock(side_effect=error)
+    rs = _make_rs(handle)
+    with pytest.raises(OperationalError, match="404"):
+        rs.complete_staging_operation([])
diff --git a/tests/unit/test_staging_remove.py b/tests/unit/test_staging_remove.py
index 0a5114621..87c1df7c1 100644
--- a/tests/unit/test_staging_remove.py
+++ b/tests/unit/test_staging_remove.py
@@ -3,6 +3,7 @@
 import pytest
 
 import databricks.sql.client as client
+from databricks.sql.result_set import ResultSet
 
 
 class TestStagingRemove:
@@ -19,7 +20,7 @@ def cursor(self):
 
     def _setup_mock_remove_response(self, cursor):
         """Set up a mock staging REMOVE response with no localFile."""
-        mock_result_set = Mock()
+        mock_result_set = Mock(spec=ResultSet)
         mock_result_set.is_staging_operation = True
         cursor.backend.execute_command.return_value = mock_result_set
 
@@ -63,7 +64,7 @@ def test_remove_with_staging_allowed_local_path(self, cursor):
     def test_get_still_requires_staging_allowed_local_path(self, cursor):
         """GET must still fail when staging_allowed_local_path is None."""
 
-        mock_result_set = Mock()
+        mock_result_set = Mock(spec=ResultSet)
         mock_result_set.is_staging_operation = True
         mock_row = Mock()
         mock_row.operation = "GET"
@@ -82,7 +83,7 @@ def test_get_still_requires_staging_allowed_local_path(self, cursor):
     def test_unsupported_operation_still_errors(self, cursor):
         """An unknown operation must still raise, even with a path configured."""
 
-        mock_result_set = Mock()
+        mock_result_set = Mock(spec=ResultSet)
         mock_result_set.is_staging_operation = True
         mock_row = Mock()
         mock_row.operation = "COPY"
diff --git a/tests/unit/test_streaming_put.py b/tests/unit/test_streaming_put.py
index 2b9a9e6d6..e0d4001ef 100644
--- a/tests/unit/test_streaming_put.py
+++ b/tests/unit/test_streaming_put.py
@@ -4,6 +4,7 @@
 import pytest
 
 import databricks.sql.client as client
+from databricks.sql.result_set import ResultSet
 
 
 class TestStreamingPut:
@@ -15,7 +16,7 @@ def cursor(self):
 
     def _setup_mock_staging_put_stream_response(self, mock_backend):
         """Helper method to set up mock staging PUT stream response."""
-        mock_result_set = Mock()
+        mock_result_set = Mock(spec=ResultSet)
         mock_result_set.is_staging_operation = True
         mock_backend.execute_command.return_value = mock_result_set
 

From c54d0912fc9261ddd8c79b0d5cc1cc625a9ea2eb Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Fri, 13 Nov 2026 23:00:24 +0000
Subject: [PATCH 02/15] Consume Volume instructions after failed completion

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 src/databricks/sql/backend/kernel/result_set.py | 13 +++++++++----
 tests/unit/test_kernel_result_set.py            |  7 +++++--
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py
index 67d63769d..1c1559b4e 100644
--- a/src/databricks/sql/backend/kernel/result_set.py
+++ b/src/databricks/sql/backend/kernel/result_set.py
@@ -111,10 +111,13 @@ def complete_staging_operation(
         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:
@@ -128,12 +131,14 @@ def complete_staging_operation(
                 outcome = self._complete_volume_operation(paths, data)
             if outcome != "completed":
                 raise ProgrammingError("Volume operation did not complete")
-            self._buffer.clear()
-            self._buffered_count = 0
-            self._exhausted = True
-            self.has_more_rows = False
             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
             )
diff --git a/tests/unit/test_kernel_result_set.py b/tests/unit/test_kernel_result_set.py
index 8c7a3dbb5..ab5e99851 100644
--- a/tests/unit/test_kernel_result_set.py
+++ b/tests/unit/test_kernel_result_set.py
@@ -298,8 +298,9 @@ def test_volume_missing_or_failed_input_closes_without_masking_error(
     handle.complete_volume_operation.assert_called_once()
 
 
-def test_volume_preserves_local_os_error_and_filename(int_schema):
-    handle = _FakeKernelHandle(int_schema, [])
+def test_failed_volume_consumes_instructions_and_preserves_os_error(int_schema):
+    handle = _FakeKernelHandle(int_schema, [_batch(int_schema, [1])])
+    handle.is_volume_operation = True
     error = RuntimeError("local file failure")
     error.local_os_error = errno.ENOENT
     error.local_path = "/allowed/missing"
@@ -309,6 +310,8 @@ def test_volume_preserves_local_os_error_and_filename(int_schema):
         rs.complete_staging_operation(["/allowed"])
     assert raised.value.errno == errno.ENOENT
     assert raised.value.filename == "/allowed/missing"
+    assert rs.fetchall() == []
+    rs.backend.finish_volume_operation.assert_called_once_with(rs.command_id)
 
 
 def test_volume_close_signals_backend_before_native_lock(int_schema):

From edd760888edd6e42a020aa09bfd2d0fba4704826 Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Fri, 13 Nov 2026 23:15:12 +0000
Subject: [PATCH 03/15] Fetch completed async results without a redundant
 status poll

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 CHANGELOG.md                                |  1 +
 src/databricks/sql/backend/kernel/client.py | 13 +++---
 tests/unit/test_kernel_client.py            | 48 ++++++++++-----------
 3 files changed, 32 insertions(+), 30 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2bae4c54e..d56207546 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,7 @@
 # Release History
 
 # Unreleased
+- Avoid an extra status poll when retrieving a completed asynchronous query result with the kernel backend.
 - 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.
diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py
index df52f34a5..3d4599db4 100644
--- a/src/databricks/sql/backend/kernel/client.py
+++ b/src/databricks/sql/backend/kernel/client.py
@@ -966,7 +966,7 @@ def get_execution_result(
         # derive a reader for re-fetch. This preserves the Thrift-parity
         # behavior where results remain re-callable until explicit close.
         #
-        # Concurrency: the owning handle is shared, and ``await_result()``
+        # Concurrency: the owning handle is shared, and ``fetch_result()``
         # below runs OUTSIDE the lock, so it must not run on the same
         # handle a concurrent ``get_query_state`` poll is already using
         # for ``status()``. Mirror that method's guard here — if a status
@@ -1004,7 +1004,8 @@ def get_execution_result(
                     if owner is not None
                     else self._kernel_session.attach_async_statement(command_id.guid)
                 )
-            stream = handle.await_result()
+            # Cursor.get_async_execution_result has already confirmed success.
+            stream = handle.fetch_result()
         except Exception as exc:
             if uses_owning_handle:
                 with self._async_handles_lock:
@@ -1013,16 +1014,16 @@ def get_execution_result(
         # ``KernelResultSet.__init__`` calls ``arrow_schema()`` which
         # can raise — map that to PEP 249 too.
         #
-        # Unlike the ``await_result()`` failure above, we deliberately do
+        # Unlike the ``fetch_result()`` failure above, we deliberately do
         # NOT discard the ``_async_result_stream_started`` marker here.
-        # By this point ``await_result()`` has already succeeded, so the
+        # By this point ``fetch_result()`` has already succeeded, so the
         # owning handle's result stream has been started (and may be
-        # partially consumed); re-awaiting that same handle on a retry is
+        # partially consumed); re-fetching that same handle on a retry is
         # not safe. Leaving the marker set routes any retry through the
         # reader fallback, which re-materialises a fresh stream.
         # The trade-off is that such a retry loses the async-statement
         # telemetry — an accepted, narrow gap limited to the case where
-        # result-set construction fails after a successful await.
+        # result-set construction fails after a successful fetch.
         try:
             return self._make_result_set(stream, cursor, command_id)
         except Exception as exc:
diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py
index f94aef0c5..b7771be81 100644
--- a/tests/unit/test_kernel_client.py
+++ b/tests/unit/test_kernel_client.py
@@ -1285,11 +1285,11 @@ def test_close_command_tolerant_when_handle_missing():
     c.close_command(fake_command_id)  # must not raise
 
 
-def _attach_returns(c, *, status=None, status_error=None, await_result=None):
+def _attach_returns(c, *, status=None, status_error=None, fetch_result=None):
     """Wire ``_kernel_session.attach_async_statement`` to return a fake
     handle. ``status`` is a ``(state, failure)`` tuple for ``handle.status()``;
     ``status_error`` makes ``attach`` itself raise (e.g. NotFound);
-    ``await_result`` sets ``handle.await_result()`` return value."""
+    ``fetch_result`` sets ``handle.fetch_result()`` return value."""
     c._kernel_session = MagicMock()
     if status_error is not None:
         c._kernel_session.attach_async_statement.side_effect = status_error
@@ -1297,8 +1297,8 @@ def _attach_returns(c, *, status=None, status_error=None, await_result=None):
     handle = MagicMock()
     if status is not None:
         handle.status.return_value = status
-    if await_result is not None:
-        handle.await_result.return_value = await_result
+    if fetch_result is not None:
+        handle.fetch_result.return_value = fetch_result
     c._kernel_session.attach_async_statement.return_value = handle
     return handle
 
@@ -1319,7 +1319,7 @@ def test_get_query_state_returns_succeeded_when_server_404s():
 def test_independent_async_read_requires_reader_api(read_result):
     c = _make_client()
     c._kernel_session = MagicMock()
-    owner = MagicMock(spec=["status", "await_result"])
+    owner = MagicMock(spec=["status", "fetch_result"])
     cid = CommandId.from_sea_statement_id("tracked-async")
     c._async_handles[cid.guid] = owner
     c._async_result_stream_started.add(cid.guid)
@@ -1440,7 +1440,7 @@ def test_get_execution_result_uses_retained_owning_handle_first():
     fake_stream = MagicMock()
     fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
     handle = MagicMock()
-    handle.await_result.return_value = fake_stream
+    handle.fetch_result.return_value = fake_stream
     cursor = MagicMock()
     cursor.arraysize = 100
     cursor.buffer_size_bytes = 1024
@@ -1452,7 +1452,7 @@ def test_get_execution_result_uses_retained_owning_handle_first():
 
     assert rs is not None
     c._kernel_session.attach_async_statement.assert_not_called()
-    handle.await_result.assert_called_once_with()
+    handle.fetch_result.assert_called_once_with()
     assert cid.guid in c._async_result_stream_started
 
 
@@ -1462,7 +1462,7 @@ def test_get_execution_result_attaches_by_id_when_no_retained_handle():
     c = _make_client()
     fake_stream = MagicMock()
     fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
-    handle = _attach_returns(c, await_result=fake_stream)
+    handle = _attach_returns(c, fetch_result=fake_stream)
     cursor = MagicMock()
     cursor.arraysize = 100
     cursor.buffer_size_bytes = 1024
@@ -1472,7 +1472,7 @@ def test_get_execution_result_attaches_by_id_when_no_retained_handle():
 
     assert rs is not None
     c._kernel_session.attach_async_statement.assert_called_once_with("async-1")
-    handle.await_result.assert_called_once_with()
+    handle.fetch_result.assert_called_once_with()
 
 
 def test_get_execution_result_owning_handle_failure_can_retry_owning_handle():
@@ -1482,7 +1482,7 @@ def test_get_execution_result_owning_handle_failure_can_retry_owning_handle():
     c = _make_client()
     c._kernel_session = MagicMock()
     handle = MagicMock()
-    handle.await_result.side_effect = _FakeKernelError(code="Unavailable")
+    handle.fetch_result.side_effect = _FakeKernelError(code="Unavailable")
     cid = CommandId.from_sea_statement_id("async-retry-owning")
     c._async_handles[cid.guid] = handle
 
@@ -1494,11 +1494,11 @@ def test_get_execution_result_owning_handle_failure_can_retry_owning_handle():
 
 
 def test_get_execution_result_construction_failure_retains_marker_and_uses_reader():
-    """If the owning handle's ``await_result()`` succeeds but result-set
+    """If the owning handle's ``fetch_result()`` succeeds but result-set
     construction then raises, the ``_async_result_stream_started`` marker
     is deliberately left set (the owning stream may be partially
-    consumed, so re-awaiting it is unsafe). A subsequent call must route
-    through the derived reader fallback rather than re-awaiting the owning
+    consumed, so re-fetching it is unsafe). A subsequent call must route
+    through the derived reader fallback rather than re-fetching the owning
     handle."""
     c = _make_client()
     c._kernel_session = MagicMock()
@@ -1507,7 +1507,7 @@ def test_get_execution_result_construction_failure_retains_marker_and_uses_reade
     # raise so ``_make_result_set`` fails after a successful await.
     owning_stream.arrow_schema.side_effect = _FakeKernelError(code="Internal")
     owning_handle = MagicMock()
-    owning_handle.await_result.return_value = owning_stream
+    owning_handle.fetch_result.return_value = owning_stream
     cid = CommandId.from_sea_statement_id("async-construct-fail")
     c._async_handles[cid.guid] = owning_handle
 
@@ -1516,15 +1516,15 @@ def test_get_execution_result_construction_failure_retains_marker_and_uses_reade
 
     # Marker stays set even though construction failed.
     assert cid.guid in c._async_result_stream_started
-    owning_handle.await_result.assert_called_once_with()
+    owning_handle.fetch_result.assert_called_once_with()
     c._kernel_session.attach_async_statement.assert_not_called()
 
-    # A retry derives a reader (fresh stream) instead of re-awaiting
+    # A retry derives a reader (fresh stream) instead of re-fetching
     # the partially-consumed owning handle.
     retry_stream = MagicMock()
     retry_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
     reader = MagicMock()
-    reader.await_result.return_value = retry_stream
+    reader.fetch_result.return_value = retry_stream
     owning_handle.reader.return_value = reader
     cursor = MagicMock()
     cursor.arraysize = 100
@@ -1535,9 +1535,9 @@ def test_get_execution_result_construction_failure_retains_marker_and_uses_reade
     assert rs is not None
     owning_handle.reader.assert_called_once_with()
     c._kernel_session.attach_async_statement.assert_not_called()
-    reader.await_result.assert_called_once_with()
-    # The owning handle was not re-awaited on the retry.
-    owning_handle.await_result.assert_called_once_with()
+    reader.fetch_result.assert_called_once_with()
+    # The owning handle was not re-fetched on the retry.
+    owning_handle.fetch_result.assert_called_once_with()
 
 
 def test_get_execution_result_maps_not_found_to_programming_error():
@@ -1779,9 +1779,9 @@ def test_get_execution_result_is_re_callable():
     fake_stream = MagicMock()
     fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
     owning_handle = MagicMock()
-    owning_handle.await_result.return_value = fake_stream
+    owning_handle.fetch_result.return_value = fake_stream
     reader = MagicMock()
-    reader.await_result.return_value = fake_stream
+    reader.fetch_result.return_value = fake_stream
     owning_handle.reader.return_value = reader
     cid = CommandId.from_sea_statement_id("async-recall-twice")
     c._async_handles[cid.guid] = owning_handle
@@ -1793,10 +1793,10 @@ def test_get_execution_result_is_re_callable():
     rs2 = c.get_execution_result(cid, cursor=cursor)
 
     assert rs1 is not None and rs2 is not None
-    owning_handle.await_result.assert_called_once_with()
+    owning_handle.fetch_result.assert_called_once_with()
     owning_handle.reader.assert_called_once_with()
     c._kernel_session.attach_async_statement.assert_not_called()
-    reader.await_result.assert_called_once_with()
+    reader.fetch_result.assert_called_once_with()
 
 
 # ---------------------------------------------------------------------------

From 07dbcc4684b7458cd99e3492452063cce85b0072 Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Fri, 13 Nov 2026 23:38:00 +0000
Subject: [PATCH 04/15] Clarify asynchronous Volume completion

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 CONNECTION_PARAMETERS.md | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md
index 046385a90..d8fd7ae83 100644
--- a/CONNECTION_PARAMETERS.md
+++ b/CONNECTION_PARAMETERS.md
@@ -249,7 +249,8 @@ are forwarded when supported by the installed kernel binding.
 ### Volume file paths
 
 With `use_kernel=True`, SQL `PUT`, `GET` and `REMOVE` complete their file
-operations before execution returns. `PUT` and `GET` require paths within
+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

From dd8bc9ca53e6f7ea97d5b95dcbeca154c8326441 Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Sat, 14 Nov 2026 00:07:09 +0000
Subject: [PATCH 05/15] Pin the native source for Volume completion

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 KERNEL_REV | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/KERNEL_REV b/KERNEL_REV
index de1b52ca7..f85bb4685 100644
--- a/KERNEL_REV
+++ b/KERNEL_REV
@@ -1 +1 @@
-0b5d0edfed32c0745d5abc2468fc3070e2eb2a17
+bdc4d35693ab94ee1991799d12a40503fdf66670

From 56ff1a0fd32965ca39bd03b673e92d5b6d4e54f3 Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Sat, 14 Nov 2026 00:13:02 +0000
Subject: [PATCH 06/15] Document completed async result retrieval

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 src/databricks/sql/backend/kernel/client.py     | 1 +
 src/databricks/sql/backend/kernel/result_set.py | 2 +-
 tests/unit/test_kernel_client.py                | 2 +-
 3 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py
index d4ffabbba..2e026719e 100644
--- a/src/databricks/sql/backend/kernel/client.py
+++ b/src/databricks/sql/backend/kernel/client.py
@@ -923,6 +923,7 @@ def get_execution_result(
         command_id: CommandId,
         cursor: "Cursor",
     ) -> "ResultSet":
+        """Fetch results after the cursor has confirmed successful completion."""
         # Prefer the original owning async handle for the first
         # in-process result stream. The kernel attaches the real
         # ExecuteStatementAsync telemetry to that handle; attached
diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py
index 1c1559b4e..5da84a859 100644
--- a/src/databricks/sql/backend/kernel/result_set.py
+++ b/src/databricks/sql/backend/kernel/result_set.py
@@ -6,7 +6,7 @@
 - ``ExecutedStatement`` — returned by ``Statement.execute()``. Has a
   ``statement_id`` and a ``cancel()`` method.
 - ``ResultStream`` — returned by ``Session.metadata().list_*`` and by
-  ``ExecutedAsyncStatement.await_result()``. No statement id; no
+  ``ExecutedAsyncStatement.fetch_result()`` after completion. No statement id; no
   cancel.
 
 Both implement the same three methods this class actually calls:
diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py
index 39940871f..8d2f4ea31 100644
--- a/tests/unit/test_kernel_client.py
+++ b/tests/unit/test_kernel_client.py
@@ -1439,7 +1439,7 @@ def test_get_execution_result_construction_failure_retains_marker_and_uses_reade
     c._kernel_session = MagicMock()
     owning_stream = MagicMock()
     # ``KernelResultSet.__init__`` calls ``arrow_schema()``; make that
-    # raise so ``_make_result_set`` fails after a successful await.
+    # raise so ``_make_result_set`` fails after a successful fetch.
     owning_stream.arrow_schema.side_effect = _FakeKernelError(code="Internal")
     owning_handle = MagicMock()
     owning_handle.fetch_result.return_value = owning_stream

From ea8e8a96111414e723d369b3b55a2ff97a126f2f Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Sat, 14 Nov 2026 00:15:37 +0000
Subject: [PATCH 07/15] Select the Windows-compatible Volume kernel

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 KERNEL_REV | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/KERNEL_REV b/KERNEL_REV
index f85bb4685..a34a99f05 100644
--- a/KERNEL_REV
+++ b/KERNEL_REV
@@ -1 +1 @@
-bdc4d35693ab94ee1991799d12a40503fdf66670
+c22bd606abc48599ab780a65a1389d141c3414bf

From 1e271bbb6b586c3d849fd38a8f87c960b500550f Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Sat, 14 Nov 2026 02:42:17 +0000
Subject: [PATCH 08/15] Keep async result reuse in the kernel and tolerate
 missing cancellers

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 CHANGELOG.md                                  |  1 -
 src/databricks/sql/backend/kernel/client.py   | 20 +++----
 .../sql/backend/kernel/result_set.py          |  2 +-
 tests/unit/test_kernel_client.py              | 57 ++++++++++---------
 4 files changed, 40 insertions(+), 40 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fa652a2ec..64cb4ea6a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,7 +2,6 @@
 
 # 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.
-- Avoid an extra status poll when retrieving a completed asynchronous query result with the kernel backend.
 - 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.
diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py
index 2e026719e..f22b9abc3 100644
--- a/src/databricks/sql/backend/kernel/client.py
+++ b/src/databricks/sql/backend/kernel/client.py
@@ -518,7 +518,7 @@ def close_session(self, session_id: SessionId) -> None:
             return
         with self._sync_cancellers_lock:
             cancellers = list(self._sync_cancellers.values()) + [
-                entry[1] for entry in self._volume_handles.values()
+                entry[1] for entry in self._volume_handles.values() if entry[1] is not None
             ]
             self._sync_cancellers.clear()
             self._volume_handles.clear()
@@ -643,7 +643,7 @@ def execute_command(
                     with self._sync_cancellers_lock:
                         self._volume_handles[command_id.guid] = (
                             stmt,
-                            self._sync_cancellers[id(cursor)],
+                            self._sync_cancellers.get(id(cursor)),
                         )
                 # Execute succeeded: the kernel now owns the statement
                 # lifecycle. It auto-closes the server statement when the
@@ -783,7 +783,7 @@ def cancel_running_cursor(self, cursor: "Cursor") -> bool:
     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:
+        if volume is not None and volume[1] is not None:
             try:
                 volume[1].cancel()
             except Exception:
@@ -923,7 +923,6 @@ def get_execution_result(
         command_id: CommandId,
         cursor: "Cursor",
     ) -> "ResultSet":
-        """Fetch results after the cursor has confirmed successful completion."""
         # Prefer the original owning async handle for the first
         # in-process result stream. The kernel attaches the real
         # ExecuteStatementAsync telemetry to that handle; attached
@@ -933,7 +932,7 @@ def get_execution_result(
         # derive a reader for re-fetch. This preserves the Thrift-parity
         # behavior where results remain re-callable until explicit close.
         #
-        # Concurrency: the owning handle is shared, and ``fetch_result()``
+        # Concurrency: the owning handle is shared, and ``await_result()``
         # below runs OUTSIDE the lock, so it must not run on the same
         # handle a concurrent ``get_query_state`` poll is already using
         # for ``status()``. Mirror that method's guard here — if a status
@@ -971,8 +970,7 @@ def get_execution_result(
                     if owner is not None
                     else self._kernel_session.attach_async_statement(command_id.guid)
                 )
-            # Cursor.get_async_execution_result has already confirmed success.
-            stream = handle.fetch_result()
+            stream = handle.await_result()
         except Exception as exc:
             if uses_owning_handle:
                 with self._async_handles_lock:
@@ -981,16 +979,16 @@ def get_execution_result(
         # ``KernelResultSet.__init__`` calls ``arrow_schema()`` which
         # can raise — map that to PEP 249 too.
         #
-        # Unlike the ``fetch_result()`` failure above, we deliberately do
+        # Unlike the ``await_result()`` failure above, we deliberately do
         # NOT discard the ``_async_result_stream_started`` marker here.
-        # By this point ``fetch_result()`` has already succeeded, so the
+        # By this point ``await_result()`` has already succeeded, so the
         # owning handle's result stream has been started (and may be
-        # partially consumed); re-fetching that same handle on a retry is
+        # partially consumed); re-awaiting that same handle on a retry is
         # not safe. Leaving the marker set routes any retry through the
         # reader fallback, which re-materialises a fresh stream.
         # The trade-off is that such a retry loses the async-statement
         # telemetry — an accepted, narrow gap limited to the case where
-        # result-set construction fails after a successful fetch.
+        # result-set construction fails after a successful await.
         try:
             return self._make_result_set(stream, cursor, command_id)
         except Exception as exc:
diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py
index 5da84a859..1c1559b4e 100644
--- a/src/databricks/sql/backend/kernel/result_set.py
+++ b/src/databricks/sql/backend/kernel/result_set.py
@@ -6,7 +6,7 @@
 - ``ExecutedStatement`` — returned by ``Statement.execute()``. Has a
   ``statement_id`` and a ``cancel()`` method.
 - ``ResultStream`` — returned by ``Session.metadata().list_*`` and by
-  ``ExecutedAsyncStatement.fetch_result()`` after completion. No statement id; no
+  ``ExecutedAsyncStatement.await_result()``. No statement id; no
   cancel.
 
 Both implement the same three methods this class actually calls:
diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py
index 8d2f4ea31..c13b6324a 100644
--- a/tests/unit/test_kernel_client.py
+++ b/tests/unit/test_kernel_client.py
@@ -1220,11 +1220,11 @@ def test_close_command_tolerant_when_handle_missing():
     c.close_command(fake_command_id)  # must not raise
 
 
-def _attach_returns(c, *, status=None, status_error=None, fetch_result=None):
+def _attach_returns(c, *, status=None, status_error=None, await_result=None):
     """Wire ``_kernel_session.attach_async_statement`` to return a fake
     handle. ``status`` is a ``(state, failure)`` tuple for ``handle.status()``;
     ``status_error`` makes ``attach`` itself raise (e.g. NotFound);
-    ``fetch_result`` sets ``handle.fetch_result()`` return value."""
+    ``await_result`` sets ``handle.await_result()`` return value."""
     c._kernel_session = MagicMock()
     if status_error is not None:
         c._kernel_session.attach_async_statement.side_effect = status_error
@@ -1232,8 +1232,8 @@ def _attach_returns(c, *, status=None, status_error=None, fetch_result=None):
     handle = MagicMock()
     if status is not None:
         handle.status.return_value = status
-    if fetch_result is not None:
-        handle.fetch_result.return_value = fetch_result
+    if await_result is not None:
+        handle.await_result.return_value = await_result
     c._kernel_session.attach_async_statement.return_value = handle
     return handle
 
@@ -1254,7 +1254,7 @@ def test_get_query_state_returns_succeeded_when_server_404s():
 def test_independent_async_read_requires_reader_api(read_result):
     c = _make_client()
     c._kernel_session = MagicMock()
-    owner = MagicMock(spec=["status", "fetch_result"])
+    owner = MagicMock(spec=["status", "await_result"])
     cid = CommandId.from_sea_statement_id("tracked-async")
     c._async_handles[cid.guid] = owner
     c._async_result_stream_started.add(cid.guid)
@@ -1375,7 +1375,7 @@ def test_get_execution_result_uses_retained_owning_handle_first():
     fake_stream = MagicMock()
     fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
     handle = MagicMock()
-    handle.fetch_result.return_value = fake_stream
+    handle.await_result.return_value = fake_stream
     cursor = MagicMock()
     cursor.arraysize = 100
     cursor.buffer_size_bytes = 1024
@@ -1387,7 +1387,7 @@ def test_get_execution_result_uses_retained_owning_handle_first():
 
     assert rs is not None
     c._kernel_session.attach_async_statement.assert_not_called()
-    handle.fetch_result.assert_called_once_with()
+    handle.await_result.assert_called_once_with()
     assert cid.guid in c._async_result_stream_started
 
 
@@ -1397,7 +1397,7 @@ def test_get_execution_result_attaches_by_id_when_no_retained_handle():
     c = _make_client()
     fake_stream = MagicMock()
     fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
-    handle = _attach_returns(c, fetch_result=fake_stream)
+    handle = _attach_returns(c, await_result=fake_stream)
     cursor = MagicMock()
     cursor.arraysize = 100
     cursor.buffer_size_bytes = 1024
@@ -1407,7 +1407,7 @@ def test_get_execution_result_attaches_by_id_when_no_retained_handle():
 
     assert rs is not None
     c._kernel_session.attach_async_statement.assert_called_once_with("async-1")
-    handle.fetch_result.assert_called_once_with()
+    handle.await_result.assert_called_once_with()
 
 
 def test_get_execution_result_owning_handle_failure_can_retry_owning_handle():
@@ -1417,7 +1417,7 @@ def test_get_execution_result_owning_handle_failure_can_retry_owning_handle():
     c = _make_client()
     c._kernel_session = MagicMock()
     handle = MagicMock()
-    handle.fetch_result.side_effect = _FakeKernelError(code="Unavailable")
+    handle.await_result.side_effect = _FakeKernelError(code="Unavailable")
     cid = CommandId.from_sea_statement_id("async-retry-owning")
     c._async_handles[cid.guid] = handle
 
@@ -1429,20 +1429,20 @@ def test_get_execution_result_owning_handle_failure_can_retry_owning_handle():
 
 
 def test_get_execution_result_construction_failure_retains_marker_and_uses_reader():
-    """If the owning handle's ``fetch_result()`` succeeds but result-set
+    """If the owning handle's ``await_result()`` succeeds but result-set
     construction then raises, the ``_async_result_stream_started`` marker
     is deliberately left set (the owning stream may be partially
-    consumed, so re-fetching it is unsafe). A subsequent call must route
-    through the derived reader fallback rather than re-fetching the owning
+    consumed, so re-awaiting it is unsafe). A subsequent call must route
+    through the derived reader fallback rather than re-awaiting the owning
     handle."""
     c = _make_client()
     c._kernel_session = MagicMock()
     owning_stream = MagicMock()
     # ``KernelResultSet.__init__`` calls ``arrow_schema()``; make that
-    # raise so ``_make_result_set`` fails after a successful fetch.
+    # raise so ``_make_result_set`` fails after a successful await.
     owning_stream.arrow_schema.side_effect = _FakeKernelError(code="Internal")
     owning_handle = MagicMock()
-    owning_handle.fetch_result.return_value = owning_stream
+    owning_handle.await_result.return_value = owning_stream
     cid = CommandId.from_sea_statement_id("async-construct-fail")
     c._async_handles[cid.guid] = owning_handle
 
@@ -1451,15 +1451,15 @@ def test_get_execution_result_construction_failure_retains_marker_and_uses_reade
 
     # Marker stays set even though construction failed.
     assert cid.guid in c._async_result_stream_started
-    owning_handle.fetch_result.assert_called_once_with()
+    owning_handle.await_result.assert_called_once_with()
     c._kernel_session.attach_async_statement.assert_not_called()
 
-    # A retry derives a reader (fresh stream) instead of re-fetching
+    # A retry derives a reader (fresh stream) instead of re-awaiting
     # the partially-consumed owning handle.
     retry_stream = MagicMock()
     retry_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
     reader = MagicMock()
-    reader.fetch_result.return_value = retry_stream
+    reader.await_result.return_value = retry_stream
     owning_handle.reader.return_value = reader
     cursor = MagicMock()
     cursor.arraysize = 100
@@ -1470,9 +1470,9 @@ def test_get_execution_result_construction_failure_retains_marker_and_uses_reade
     assert rs is not None
     owning_handle.reader.assert_called_once_with()
     c._kernel_session.attach_async_statement.assert_not_called()
-    reader.fetch_result.assert_called_once_with()
-    # The owning handle was not re-fetched on the retry.
-    owning_handle.fetch_result.assert_called_once_with()
+    reader.await_result.assert_called_once_with()
+    # The owning handle was not re-awaited on the retry.
+    owning_handle.await_result.assert_called_once_with()
 
 
 def test_get_execution_result_maps_not_found_to_programming_error():
@@ -1714,9 +1714,9 @@ def test_get_execution_result_is_re_callable():
     fake_stream = MagicMock()
     fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
     owning_handle = MagicMock()
-    owning_handle.fetch_result.return_value = fake_stream
+    owning_handle.await_result.return_value = fake_stream
     reader = MagicMock()
-    reader.fetch_result.return_value = fake_stream
+    reader.await_result.return_value = fake_stream
     owning_handle.reader.return_value = reader
     cid = CommandId.from_sea_statement_id("async-recall-twice")
     c._async_handles[cid.guid] = owning_handle
@@ -1728,10 +1728,10 @@ def test_get_execution_result_is_re_callable():
     rs2 = c.get_execution_result(cid, cursor=cursor)
 
     assert rs1 is not None and rs2 is not None
-    owning_handle.fetch_result.assert_called_once_with()
+    owning_handle.await_result.assert_called_once_with()
     owning_handle.reader.assert_called_once_with()
     c._kernel_session.attach_async_statement.assert_not_called()
-    reader.fetch_result.assert_called_once_with()
+    reader.await_result.assert_called_once_with()
 
 
 # ---------------------------------------------------------------------------
@@ -2719,10 +2719,13 @@ def test_only_reserved_headers_omits_kwarg(self, monkeypatch):
 
 
 @pytest.mark.parametrize("is_volume", [False, True])
-def test_sync_volume_canceller_survives_until_completion(is_volume):
+@pytest.mark.parametrize("has_canceller", [False, True])
+def test_sync_volume_canceller_survives_until_completion(is_volume, has_canceller):
     c = _make_client()
     c._kernel_session = MagicMock()
     stmt = c._kernel_session.statement.return_value
+    if not has_canceller:
+        stmt.canceller.side_effect = RuntimeError("canceller unavailable")
     stmt.execute.return_value = MagicMock(
         statement_id="volume-id",
         is_volume_operation=is_volume,
@@ -2742,6 +2745,6 @@ def test_sync_volume_canceller_survives_until_completion(is_volume):
     )
     assert result.is_staging_operation is is_volume
     c.cancel_command(result.command_id)
-    assert stmt.canceller.return_value.cancel.called is is_volume
+    assert stmt.canceller.return_value.cancel.called is (is_volume and has_canceller)
     c.finish_volume_operation(result.command_id)
     assert not c._volume_handles

From ab1efce18f6654362c2f7155a959ca756b5d87cc Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Sat, 14 Nov 2026 02:52:59 +0000
Subject: [PATCH 09/15] Pin native Volume storage retry and terminal response
 fixes

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 KERNEL_REV | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/KERNEL_REV b/KERNEL_REV
index a34a99f05..8a2f8e599 100644
--- a/KERNEL_REV
+++ b/KERNEL_REV
@@ -1 +1 @@
-c22bd606abc48599ab780a65a1389d141c3414bf
+22f96facc276836bcac7230c8a2c0987f734f2c7

From ccc5b46e7a30e1f912a685134f009dad62c3e11a Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Sat, 14 Nov 2026 03:32:02 +0000
Subject: [PATCH 10/15] Preserve Volume error provenance and stream teardown

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 KERNEL_REV                                    |  2 +-
 .../sql/backend/kernel/result_set.py          | 24 +++++++----
 tests/unit/test_kernel_client.py              | 20 ++++++++++
 tests/unit/test_kernel_result_set.py          | 40 ++++++++++++++++++-
 4 files changed, 76 insertions(+), 10 deletions(-)

diff --git a/KERNEL_REV b/KERNEL_REV
index 8a2f8e599..47d2f1c8a 100644
--- a/KERNEL_REV
+++ b/KERNEL_REV
@@ -1 +1 @@
-22f96facc276836bcac7230c8a2c0987f734f2c7
+08a921377dd79e5a301e007da7290c1984140952
diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py
index 1c1559b4e..b7278e10d 100644
--- a/src/databricks/sql/backend/kernel/result_set.py
+++ b/src/databricks/sql/backend/kernel/result_set.py
@@ -160,7 +160,10 @@ def _complete_volume_operation(
                     ) 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:
+            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
@@ -309,14 +312,13 @@ def close(self) -> None:
             self.has_been_closed_server_side = True
             self.status = CommandState.CLOSED
             return
-        # Signal through the detached canceller before close waits for a transfer's native lock.
         backend = cast("KernelDatabricksClient", self.backend)
-        try:
-            backend.close_command(self.command_id)
-        except Exception as exc:
-            logger.warning(
-                "backend.close_command from result-set close failed: %s", exc
-            )
+        if self.is_staging_operation:
+            # 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:
@@ -324,6 +326,12 @@ def close(self) -> None:
             # level; log and swallow so the cursor's __del__ /
             # connection close path stays clean.
             logger.warning("Error closing kernel handle: %s", exc)
+        try:
+            backend.close_command(self.command_id)
+        except Exception as exc:
+            logger.warning(
+                "backend.close_command from result-set close failed: %s", exc
+            )
         self._buffer.clear()
         self._buffered_count = 0
         self._kernel_handle = None
diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py
index c13b6324a..db46bbdb9 100644
--- a/tests/unit/test_kernel_client.py
+++ b/tests/unit/test_kernel_client.py
@@ -1574,6 +1574,9 @@ def test_close_session_clears_async_handles_even_if_close_fails():
     good = MagicMock()
     bad = MagicMock()
     bad.close.side_effect = _FakeKernelError(code="Unavailable")
+    bad.cancel.side_effect = _FakeKernelError(code="Unavailable")
+    c._sync_cancellers[1] = bad
+    c._volume_handles["volume"] = (MagicMock(), good)
     c._async_handles["a"] = good
     c._async_handles["b"] = bad
     c._kernel_session = MagicMock()
@@ -1581,6 +1584,10 @@ def test_close_session_clears_async_handles_even_if_close_fails():
     assert c._async_handles == {}
     assert good.close.called
     assert bad.close.called
+    bad.cancel.assert_called_once_with()
+    good.cancel.assert_called_once_with()
+    assert c._sync_cancellers == {}
+    assert c._volume_handles == {}
 
 
 def test_close_session_closes_and_drops_swept_handles():
@@ -2748,3 +2755,16 @@ def test_sync_volume_canceller_survives_until_completion(is_volume, has_cancelle
     assert stmt.canceller.return_value.cancel.called is (is_volume and has_canceller)
     c.finish_volume_operation(result.command_id)
     assert not c._volume_handles
+
+
+@pytest.mark.parametrize("cancel_fails", [False, True])
+def test_close_volume_command_cancels_and_removes_handle(cancel_fails):
+    c = _make_client()
+    command_id = CommandId.from_sea_statement_id("volume-id")
+    canceller = MagicMock()
+    if cancel_fails:
+        canceller.cancel.side_effect = RuntimeError("cancel failed")
+    c._volume_handles[command_id.guid] = (MagicMock(), canceller)
+    c.close_command(command_id)
+    canceller.cancel.assert_called_once_with()
+    assert not c._volume_handles
diff --git a/tests/unit/test_kernel_result_set.py b/tests/unit/test_kernel_result_set.py
index ab5e99851..9e9067489 100644
--- a/tests/unit/test_kernel_result_set.py
+++ b/tests/unit/test_kernel_result_set.py
@@ -316,9 +316,10 @@ def test_failed_volume_consumes_instructions_and_preserves_os_error(int_schema):
 
 def test_volume_close_signals_backend_before_native_lock(int_schema):
     handle = _FakeKernelHandle(int_schema, [])
+    handle.is_volume_operation = True
     rs = _make_rs(handle)
     handle.close = MagicMock(
-        side_effect=lambda: rs.backend.close_command.assert_called_once_with(
+        side_effect=lambda: rs.backend.cancel_command.assert_called_once_with(
             rs.command_id
         )
     )
@@ -326,12 +327,49 @@ def test_volume_close_signals_backend_before_native_lock(int_schema):
     handle.close.assert_called_once()
 
 
+def test_ordinary_close_releases_stream_before_backend(int_schema):
+    handle = _FakeKernelHandle(int_schema, [])
+    rs = _make_rs(handle)
+
+    def close_backend(_):
+        assert handle.closed
+
+    rs.backend.close_command.side_effect = close_backend
+    rs.close()
+    rs.backend.cancel_command.assert_not_called()
+    rs.backend.close_command.assert_called_once_with(rs.command_id)
+
+
+@pytest.mark.parametrize("code,status", [("NotFound", 404), ("InvalidArgument", 400)])
+def test_volume_instruction_http_error_preserves_mapping(int_schema, code, status):
+    from databricks.sql.backend.kernel import _errors
+    from databricks.sql.exc import ProgrammingError
+
+    error = _errors._kernel.KernelError("instruction read failed")
+    error.code = code
+    error.message = "instruction read failed"
+    error.http_status = status
+    error.sql_state = "22023"
+    error.query_id = "query-id"
+    handle = _FakeKernelHandle(int_schema, [])
+    handle.complete_volume_operation = MagicMock(side_effect=error)
+    rs = _make_rs(handle)
+    with pytest.raises(ProgrammingError, match="instruction read failed") as raised:
+        rs.complete_staging_operation([])
+    assert raised.value.http_status == status
+    assert raised.value.code == code
+    assert raised.value.sql_state == "22023"
+    assert raised.value.query_id == "query-id"
+    assert raised.value.__cause__ is error
+
+
 def test_volume_http_404_is_an_operational_error(int_schema):
     from databricks.sql.exc import OperationalError
 
     handle = _FakeKernelHandle(int_schema, [])
     error = RuntimeError("storage response")
     error.http_status = 404
+    error.error_code = "VOLUME_STORAGE_HTTP_ERROR"
     handle.complete_volume_operation = MagicMock(side_effect=error)
     rs = _make_rs(handle)
     with pytest.raises(OperationalError, match="404"):

From 64c9838d34cf81da31f8f80facca8d64d469964b Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Sat, 14 Nov 2026 03:50:07 +0000
Subject: [PATCH 11/15] Avoid cancelling completed Volume transfers during
 close

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 .../sql/backend/kernel/result_set.py          |  2 +-
 tests/unit/test_kernel_result_set.py          | 19 +++++++++++++------
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py
index b7278e10d..76accc06b 100644
--- a/src/databricks/sql/backend/kernel/result_set.py
+++ b/src/databricks/sql/backend/kernel/result_set.py
@@ -313,7 +313,7 @@ def close(self) -> None:
             self.status = CommandState.CLOSED
             return
         backend = cast("KernelDatabricksClient", self.backend)
-        if self.is_staging_operation:
+        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)
diff --git a/tests/unit/test_kernel_result_set.py b/tests/unit/test_kernel_result_set.py
index 9e9067489..31d8e68bf 100644
--- a/tests/unit/test_kernel_result_set.py
+++ b/tests/unit/test_kernel_result_set.py
@@ -314,15 +314,22 @@ def test_failed_volume_consumes_instructions_and_preserves_os_error(int_schema):
     rs.backend.finish_volume_operation.assert_called_once_with(rs.command_id)
 
 
-def test_volume_close_signals_backend_before_native_lock(int_schema):
+@pytest.mark.parametrize("completed", [False, True])
+def test_volume_close_cancels_only_pending_transfer(int_schema, completed):
     handle = _FakeKernelHandle(int_schema, [])
     handle.is_volume_operation = True
     rs = _make_rs(handle)
-    handle.close = MagicMock(
-        side_effect=lambda: rs.backend.cancel_command.assert_called_once_with(
-            rs.command_id
-        )
-    )
+    if completed:
+        handle.complete_volume_operation = MagicMock(return_value="completed")
+        rs.complete_staging_operation([])
+
+    def close_native():
+        if completed:
+            rs.backend.cancel_command.assert_not_called()
+        else:
+            rs.backend.cancel_command.assert_called_once_with(rs.command_id)
+
+    handle.close = MagicMock(side_effect=close_native)
     rs.close()
     handle.close.assert_called_once()
 

From f4292d36246af3ae665af72ee0b4050b497b50a6 Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Sat, 14 Nov 2026 04:01:50 +0000
Subject: [PATCH 12/15] Format Volume cancellation cleanup comprehension

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 src/databricks/sql/backend/kernel/client.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py
index f22b9abc3..01d98c477 100644
--- a/src/databricks/sql/backend/kernel/client.py
+++ b/src/databricks/sql/backend/kernel/client.py
@@ -518,7 +518,9 @@ def close_session(self, session_id: SessionId) -> 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
+                entry[1]
+                for entry in self._volume_handles.values()
+                if entry[1] is not None
             ]
             self._sync_cancellers.clear()
             self._volume_handles.clear()

From 07538ed6a4243081466d8ca021bba0e40ef82e89 Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Sat, 14 Nov 2026 05:19:16 +0000
Subject: [PATCH 13/15] Select native Volume inactivity timeout correction

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 KERNEL_REV | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/KERNEL_REV b/KERNEL_REV
index 47d2f1c8a..bc57504ba 100644
--- a/KERNEL_REV
+++ b/KERNEL_REV
@@ -1 +1 @@
-08a921377dd79e5a301e007da7290c1984140952
+855847bdca71e07285cfc60332fb2f33f8e47786

From ff77a54484482f5874d0ff6c9f4412dafa2442eb Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Sat, 14 Nov 2026 23:54:43 +0000
Subject: [PATCH 14/15] build(kernel): select corrected Volume completion

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 CONNECTION_PARAMETERS.md | 2 +-
 KERNEL_REV               | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md
index d8fd7ae83..884018669 100644
--- a/CONNECTION_PARAMETERS.md
+++ b/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 bounds response inactivity after upload and between response chunks; 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). |
diff --git a/KERNEL_REV b/KERNEL_REV
index bc57504ba..53687ee12 100644
--- a/KERNEL_REV
+++ b/KERNEL_REV
@@ -1 +1 @@
-855847bdca71e07285cfc60332fb2f33f8e47786
+a543bbc572b627f059461fce351713da8cb59df9

From cd42b4e04af0557815b26db578033e7739fee9ba Mon Sep 17 00:00:00 2001
From: Aakash Saravanan <aakash.saravanan@databricks.com>
Date: Sun, 15 Nov 2026 01:48:43 +0000
Subject: [PATCH 15/15] build: select native Volume cancellation and error
 handling

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
---
 CONNECTION_PARAMETERS.md | 2 +-
 KERNEL_REV               | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md
index 884018669..8a7c5e7f8 100644
--- a/CONNECTION_PARAMETERS.md
+++ b/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 SQL request default. Volume transfers have no total deadline: an explicit positive value bounds response inactivity after upload and between response chunks; unset or `0` leaves storage inactivity unbounded. |
+| `_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). |
diff --git a/KERNEL_REV b/KERNEL_REV
index 53687ee12..3232e69ae 100644
--- a/KERNEL_REV
+++ b/KERNEL_REV
@@ -1 +1 @@
-a543bbc572b627f059461fce351713da8cb59df9
+566fe7f0a8a2fa246db95045a7000136ee6070d5