feat(volume): enable native SQL file transfers

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

Author: aakash-saravanan-db

KERNEL_REV

@@ -1 +1 @@
-6f4b6d607a61bdde61dcda2cf82bc8e80fd32cb9
+d41cc3657748966d8b69a802bc40e8d019022ceb

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

@@ -39,6 +39,7 @@
 from __future__ import annotations
 
 import logging
+from typing import Union
 
 from databricks.sql.exc import (
     DatabaseError,
@@ -115,7 +116,9 @@
 }
 
 
-def reraise_kernel_error(exc: "_kernel.KernelError") -> "Error":
+def reraise_kernel_error(
+    exc: "_kernel.KernelError",
+) -> "Union[Error, FileNotFoundError]":
     """Convert a ``databricks_sql_kernel.KernelError`` to a PEP 249
     exception with the kernel's structured attributes forwarded onto
     the new instance.
@@ -126,6 +129,9 @@ def reraise_kernel_error(exc: "_kernel.KernelError") -> "Error":
     """
     code = getattr(exc, "code", "Unknown")
     cls = _CODE_TO_EXCEPTION.get(code, DatabaseError)
+    file_transfer = getattr(exc, "error_code", None) == "FILE_TRANSFER_ERROR"
+    if file_transfer:
+        cls = ProgrammingError if code == "PermissionDenied" else OperationalError
 
     # For ServerOperationError, reproduce the Thrift backend's
     # ``context`` dict so callers that read
@@ -141,7 +147,15 @@ def reraise_kernel_error(exc: "_kernel.KernelError") -> "Error":
             "operation-id": getattr(exc, "query_id", None),
             "diagnostic-info": getattr(exc, "diagnostic_info", None),
         }
-    new = cls(getattr(exc, "message", str(exc)), context)
+    message = getattr(exc, "message", str(exc))
+    if (
+        file_transfer
+        and code == "NotFound"
+        and getattr(exc, "http_status", None) is None
+    ):
+        new = FileNotFoundError(message)
+    else:
+        new = cls(message, context)
 
     for attr in (
         "code",
@@ -162,7 +176,9 @@ def reraise_kernel_error(exc: "_kernel.KernelError") -> "Error":
     return new
 
 
-def wrap_kernel_exception(what: str, exc: BaseException) -> "Error":
+def wrap_kernel_exception(
+    what: str, exc: BaseException
+) -> "Union[Error, FileNotFoundError]":
     """Map any exception from a PyO3 call site to a PEP 249 exception.
 
     - ``KernelError`` → mapped class with structured attrs forwarded.

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
@@ -74,41 +72,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 /
@@ -145,21 +108,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
-
-
 def _kernel_session_accepts_kwarg(name: str) -> bool:
     """True iff the installed ``databricks_sql_kernel.Session`` constructor
     declares keyword ``name``.
@@ -575,28 +523,14 @@ 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:
                 stmt.set_sql(operation)
+                allowed_paths = cursor.connection.staging_allowed_local_path
+                if isinstance(allowed_paths, str):
+                    allowed_paths = [allowed_paths]
+                stmt.set_staging_allowed_local_paths(allowed_paths or [])
                 stmt.set_row_limit(row_limit)
                 if query_tags:
                     # Per-statement query tags. The kernel serialises the

tests/e2e/test_kernel_backend.py

@@ -32,6 +32,7 @@
     DatabaseError,
     NotSupportedError,
     OperationalError,
+    ProgrammingError,
     ServerOperationError,
 )
 
@@ -366,9 +367,7 @@ def test_schemas_with_empty_string_filter_matches_nothing(conn):
         assert cur.fetchall() == []
 
 
-@pytest.mark.parametrize(
-    "empty_filter", ["schema_name", "table_name", "column_name"]
-)
+@pytest.mark.parametrize("empty_filter", ["schema_name", "table_name", "column_name"])
 def test_columns_with_empty_string_filter_matches_nothing(conn, empty_filter):
     filters = {
         "catalog_name": "system",
@@ -777,28 +776,55 @@ 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'"
-            )
+@pytest.mark.parametrize("async_op", [False, True])
+def test_volume_file_lifecycle(kernel_conn_params, tmp_path, catalog, async_op):
+    allowed = tmp_path / "allowed"
+    allowed.mkdir()
+    source = allowed / "input.bin"
+    destination = allowed / "output.bin"
+    payload = bytes(range(256)) * 4096
+    source.write_bytes(payload)
+    denied = tmp_path / "denied.bin"
+    denied.write_bytes(payload)
+    catalog = catalog or "main"
+    namespace = f"volume_{uuid4().hex}"
+    qualified = f"`{catalog.replace('`', '``')}`.`{namespace}`"
+    remote = f"/Volumes/{catalog}/{namespace}/files/input.bin"
+    with sql.connect(
+        **kernel_conn_params, staging_allowed_local_path=str(allowed)
+    ) as connection:
+        with connection.cursor() as cur:
+            cur.execute(f"CREATE SCHEMA {qualified}")
+            try:
+                cur.execute(f"CREATE VOLUME {qualified}.files")
+
+                def execute(query):
+                    if async_op:
+                        cur.execute_async(query)
+                        cur.get_async_execution_result()
+                    else:
+                        cur.execute(query)
+                    return cur.fetchall()
+
+                assert (
+                    execute(f"-- upload\nPUT '{source}' INTO '{remote}' OVERWRITE")
+                    == []
+                )
+                assert execute(f"GET '{remote}' TO '{destination}'") == []
+                assert destination.read_bytes() == payload
+                with pytest.raises(ProgrammingError):
+                    execute(f"PUT '{denied}' INTO '{remote}' OVERWRITE")
+                with pytest.raises(FileNotFoundError):
+                    execute(
+                        f"PUT '{allowed / 'missing.bin'}' INTO '{remote}' OVERWRITE"
+                    )
+                assert execute(f"REMOVE '{remote}'") == []
+                assert execute(f"REMOVE '{remote}'") == []
+                with pytest.raises(OperationalError, match="404"):
+                    execute(f"GET '{remote}' TO '{allowed / 'deleted.bin'}'")
+                assert execute("SELECT 1")[0][0] == 1
+            finally:
+                cur.execute(f"DROP SCHEMA {qualified} CASCADE")
 
 
 # ── Error fidelity — diagnostic-info reaches .context ─────────────

tests/unit/test_kernel_client.py

@@ -393,9 +393,7 @@ def fake_session(**kw):
 
 
 @pytest.mark.parametrize("enabled", [True, False])
-def test_open_session_passes_geospatial_representation_to_kernel(
-    monkeypatch, enabled
-):
+def test_open_session_passes_geospatial_representation_to_kernel(monkeypatch, enabled):
     captured = {}
 
     def fake_session(*, enable_geospatial_support=True, **kw):
@@ -807,70 +805,43 @@ 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",
+    "allowed, expected",
     [
-        "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'",
+        (None, []),
+        ("/tmp/files", ["/tmp/files"]),
+        (["/tmp/a", "/tmp/b"], ["/tmp/a", "/tmp/b"]),
     ],
 )
-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)."""
+@pytest.mark.parametrize("async_op", [False, True])
+def test_execute_command_forwards_allowed_local_paths(allowed, expected, async_op):
     c = _make_client()
     c._kernel_session = MagicMock()
     cursor = MagicMock()
+    cursor.connection.staging_allowed_local_path = allowed
     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
+    stmt = c._kernel_session.statement.return_value
+    executed = MagicMock(
+        statement_id="00000000-0000-0000-0000-000000000001",
+        arrow_schema=MagicMock(return_value=pa.schema([("x", pa.int64())])),
+    )
+    stmt.execute.return_value = executed
+    stmt.submit.return_value = executed
+    c.execute_command(
+        operation="PUT '/tmp/files/input' INTO '/Volumes/c/s/v/input'",
+        session_id=MagicMock(),
+        max_rows=1,
+        max_bytes=1,
+        lz4_compression=False,
+        cursor=cursor,
+        use_cloud_fetch=False,
+        parameters=[],
+        async_op=async_op,
+        enforce_embedded_schema_correctness=False,
+    )
+    stmt.set_staging_allowed_local_paths.assert_called_once_with(expected)
+    (stmt.submit if async_op else stmt.execute).assert_called_once()
 
 
 # ---------------------------------------------------------------------------
@@ -2582,3 +2553,21 @@ 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(
+    "code, http_status, expected",
+    [
+        ("PermissionDenied", None, ProgrammingError),
+        ("NotFound", None, FileNotFoundError),
+        ("NotFound", 404, OperationalError),
+        ("NetworkError", None, OperationalError),
+    ],
+)
+def test_file_transfer_error_matches_thrift_exception(code, http_status, expected):
+    error = _FakeKernelError(code=code, message="File transfer failed")
+    error.error_code = "FILE_TRANSFER_ERROR"
+    error.http_status = http_status
+    result = kernel_client._reraise_kernel_error(error)
+    assert type(result) is expected
+    assert result.error_code == "FILE_TRANSFER_ERROR"