fix(volume): preserve filesystem errors and unify transfer detection

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

Author: aakash-saravanan-db

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

@@ -120,20 +120,23 @@
 
 def reraise_kernel_error(
     exc: "_kernel.KernelError",
-) -> "Union[Error, FileNotFoundError]":
+) -> "Union[Error, OSError]":
     """Map a kernel error and forward its structured attributes.
 
-    Local ``FILE_TRANSFER_ERROR`` errors with ``NotFound`` and no HTTP
-    status use the built-in ``FileNotFoundError``; other errors use
-    PEP 249 exception classes.
+    Local missing files use ``FileNotFoundError``; filesystem permission
+    failures use ``PermissionError``. Other errors use PEP 249 classes.
 
     The returned exception is raised by callers with ``raise ... from
     exc``; the ``from`` clause is what sets ``__cause__``, so we don't
     touch it here.
     """
     code = getattr(exc, "code", "Unknown")
     cls = _CODE_TO_EXCEPTION.get(code, DatabaseError)
-    file_transfer = getattr(exc, "error_code", None) == "FILE_TRANSFER_ERROR"
+    error_code = getattr(exc, "error_code", None)
+    file_transfer = error_code in (
+        "FILE_TRANSFER_ERROR",
+        "FILE_TRANSFER_PERMISSION_ERROR",
+    )
     if file_transfer:
         cls = (
             ProgrammingError
@@ -157,9 +160,11 @@ def reraise_kernel_error(
             "diagnostic-info": getattr(exc, "diagnostic_info", None),
         }
     message = getattr(exc, "message", str(exc))
-    new: Union[Error, FileNotFoundError]
+    new: Union[Error, OSError]
     # Local missing files have no HTTP status; remote missing files carry 404.
-    if (
+    if error_code == "FILE_TRANSFER_PERMISSION_ERROR":
+        new = PermissionError(message)
+    elif (
         file_transfer
         and code == "NotFound"
         and getattr(exc, "http_status", None) is None
@@ -187,9 +192,7 @@ def reraise_kernel_error(
     return new
 
 
-def wrap_kernel_exception(
-    what: str, exc: BaseException
-) -> "Union[Error, FileNotFoundError]":
+def wrap_kernel_exception(what: str, exc: BaseException) -> "Union[Error, OSError]":
     """Map a PyO3 exception to the connector's public error contract.
 
     - ``KernelError`` → mapped class with structured attrs forwarded.

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

@@ -74,10 +74,6 @@
 # per-request skip-and-warn.
 _KERNEL_MANAGED_HEADERS = frozenset({"authorization", "x-databricks-org-id"})
 
-# Transfer detection controls local-path forwarding and the older-wheel guard.
-_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
@@ -143,19 +139,10 @@ 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 _staging_verb(stripped_sql: str) -> Optional[str]:
+    """Match a leading verb so SELECT 'GET' cannot enable local file access."""
+    match = re.match(r"(PUT|GET|REMOVE)\b", stripped_sql, re.IGNORECASE)
+    return match.group(1).upper() if match else None
 
 
 def _kernel_session_accepts_kwarg(name: str) -> bool:
@@ -576,9 +563,11 @@ def execute_command(
         close_stmt = True
         try:
             try:
+                stripped_sql = _strip_leading_sql_comments(operation)
+                staging_verb = _staging_verb(stripped_sql)
                 if re.match(
                     r"""PUT\s+(['"])__input_stream__\1(?:\s|$)""",
-                    _strip_leading_sql_comments(operation),
+                    stripped_sql,
                     re.IGNORECASE,
                 ):
                     raise NotSupportedError(
@@ -587,11 +576,7 @@ def execute_command(
                 stmt.set_sql(operation)
                 allowed_paths = (
                     cursor.connection.staging_allowed_local_path
-                    if re.match(
-                        r"(?:PUT|GET)\b",
-                        _strip_leading_sql_comments(operation),
-                        re.IGNORECASE,
-                    )
+                    if staging_verb in ("PUT", "GET")
                     else None
                 )
                 set_paths = getattr(stmt, "set_staging_allowed_local_paths", None)
@@ -613,7 +598,7 @@ def execute_command(
                             "staging_allowed_local_path must be a string or a list of paths"
                         ) from exc
                     set_paths(normalized_paths)
-                elif _is_staging_statement(operation):
+                elif staging_verb is not None:
                     raise NotSupportedError(
                         "Volume transfers require a databricks-sql-kernel wheel "
                         "with staging_allowed_local_paths support."

tests/unit/test_kernel_client.py

@@ -926,6 +926,8 @@ def test_kernel_rejects_streaming_put(query):
         "PUT 'a' INTO 'b'",
         "-- comment\nGET 'a' TO 'b'",
         "/* comment */ REMOVE 'a'",
+        "GET'/Volumes/c/s/v/file' TO '/tmp/file'",
+        "/* comment */ REMOVE'/Volumes/c/s/v/file'",
     ],
 )
 def test_execute_command_with_older_kernel(operation):
@@ -2679,6 +2681,7 @@ def test_only_reserved_headers_omits_kwarg(self, monkeypatch):
         ("PermissionDenied", 403, "FILE_TRANSFER_ERROR", OperationalError),
         ("NotFound", None, "FILE_TRANSFER_ERROR", FileNotFoundError),
         ("NotFound", 404, "FILE_TRANSFER_ERROR", OperationalError),
+        ("PermissionDenied", None, "FILE_TRANSFER_PERMISSION_ERROR", PermissionError),
         ("NetworkError", None, "FILE_TRANSFER_ERROR", OperationalError),
         ("Internal", None, "FILE_TRANSFER_ERROR", OperationalError),
         ("Internal", None, None, DatabaseError),