diff --git a/CHANGELOG.md b/CHANGELOG.md index 4fbb64b68..acad0505c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ # Release History # Unreleased +- Support SQL PUT, GET, and REMOVE file transfers on the kernel backend with local-path permission checks and native filesystem exceptions. Streaming PUT remains unsupported in kernel mode. - 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)) # 4.9.0 (2026-11-04) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 37736ed1a..92de73b33 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`. `REMOVE` requires no local path. Kernel mode does not support streaming PUT. | ## Telemetry @@ -206,13 +206,12 @@ installed kernel binding. `force_enable_telemetry` remains Python-driver-only. `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 diff --git a/KERNEL_REV b/KERNEL_REV index d025be44a..b1078d374 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -6f4b6d607a61bdde61dcda2cf82bc8e80fd32cb9 +37c7a3ec1b54f4690d9325d65f2220e0bb64b83e diff --git a/examples/staging_ingestion.py b/examples/staging_ingestion.py index a55be4778..4a0dce397 100644 --- a/examples/staging_ingestion.py +++ b/examples/staging_ingestion.py @@ -30,7 +30,8 @@ 2. Set the FILEPATH constant to the path of a file that will be uploaded (this example assumes its a CSV file) 3. Run this file -Note: staging_allowed_local_path can be either a Pathlike object or a list of Pathlike objects. +Note: staging_allowed_local_path accepts a path string or a list of strings or +os.PathLike objects that resolve to strings. """ INGESTION_USER = "some.user@example.com" diff --git a/src/databricks/sql/backend/kernel/_errors.py b/src/databricks/sql/backend/kernel/_errors.py index bcf3f56b1..2584f7b87 100644 --- a/src/databricks/sql/backend/kernel/_errors.py +++ b/src/databricks/sql/backend/kernel/_errors.py @@ -13,10 +13,13 @@ - Anything else — ``TypeError`` / ``OverflowError`` / ``ValueError`` from PyO3 argument conversion, or arbitrary extension-internal Python errors. These would otherwise propagate - raw to connector callers, breaking the DB-API contract that says - "only PEP 249 exception types cross the boundary". Wrapped in + raw to connector callers. Wrapped in ``OperationalError`` here. +Local transfer failures use built-in ``OSError`` classes, including +``FileNotFoundError``, ``PermissionError`` and other platform subclasses. +The original OS error code and requested filename are preserved when available. + These primitives live in their own module so both ``client.py`` (which orchestrates PyO3 calls) and ``result_set.py`` (which calls ``fetch_next_batch`` on the same kernel handles) can share them @@ -39,10 +42,13 @@ from __future__ import annotations import logging +import os +from typing import Union from databricks.sql.exc import ( DatabaseError, Error, + NotSupportedError, OperationalError, ProgrammingError, ServerOperationError, @@ -115,10 +121,14 @@ } -def reraise_kernel_error(exc: "_kernel.KernelError") -> "Error": - """Convert a ``databricks_sql_kernel.KernelError`` to a PEP 249 - exception with the kernel's structured attributes forwarded onto - the new instance. +def reraise_kernel_error( + exc: "_kernel.KernelError", +) -> "Union[Error, OSError]": + """Map a kernel error and forward its structured attributes. + + Local transfer failures use built-in ``OSError`` classes, including + missing-file, permission, directory and disk-full errors. Non-local + 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 @@ -126,6 +136,20 @@ def reraise_kernel_error(exc: "_kernel.KernelError") -> "Error": """ code = getattr(exc, "code", "Unknown") cls = _CODE_TO_EXCEPTION.get(code, DatabaseError) + error_code = getattr(exc, "error_code", None) + if error_code == "FILE_TRANSFER_UNSUPPORTED_STREAM": + cls = NotSupportedError + file_transfer = error_code in ( + "FILE_TRANSFER_ERROR", + "FILE_TRANSFER_PERMISSION_ERROR", + ) + if file_transfer: + cls = ( + ProgrammingError + if code in ("PermissionDenied", "InvalidArgument") + and getattr(exc, "http_status", None) is None + else OperationalError + ) # For ServerOperationError, reproduce the Thrift backend's # ``context`` dict so callers that read @@ -141,13 +165,49 @@ 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)) + os_error_code = getattr(exc, "os_error_code", None) + os_error_filename = getattr(exc, "os_error_filename", None) + file_error_args: tuple + if os_error_code is None: + file_error_args = ( + (None, message, os_error_filename) + if os_error_filename is not None + else (message,) + ) + elif os.name == "nt": + file_error_args = (0, message, os_error_filename, os_error_code) + else: + file_error_args = (os_error_code, message, os_error_filename) + new: Union[Error, OSError] + # Local missing files have no HTTP status; remote missing files carry 404. + if ( + file_transfer + and os_error_code is not None + and getattr(exc, "http_status", None) is None + ): + new = OSError(*file_error_args) + elif ( + error_code == "FILE_TRANSFER_PERMISSION_ERROR" + and getattr(exc, "http_status", None) is None + ): + new = PermissionError(*file_error_args) + elif ( + file_transfer + and code == "NotFound" + and getattr(exc, "http_status", None) is None + ): + new = FileNotFoundError(*file_error_args) + else: + new = cls(message, context) for attr in ( "code", "sql_state", "error_code", "vendor_code", + "os_error_code", + "os_error_filename", "http_status", "retryable", "query_id", @@ -162,12 +222,14 @@ def reraise_kernel_error(exc: "_kernel.KernelError") -> "Error": return new -def wrap_kernel_exception(what: str, exc: BaseException) -> "Error": - """Map any exception from a PyO3 call site to a PEP 249 exception. +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. - - Already-PEP-249 ``Error`` (e.g. raised by an inner caller that - already mapped) → passed through unchanged. + Local transfer failures map to built-in ``OSError`` classes with + their original OS error code and requested filename when available. + - Already-mapped PEP 249 or local transfer errors → passed through + unchanged so nested callers preserve their class and attributes. - Anything else (``TypeError`` / ``ValueError`` / etc. from PyO3 argument conversion, extension-internal errors) → wrapped in ``OperationalError``. @@ -180,6 +242,11 @@ def wrap_kernel_exception(what: str, exc: BaseException) -> "Error": return reraise_kernel_error(exc) if isinstance(exc, Error): return exc + if isinstance(exc, OSError) and getattr(exc, "error_code", None) in ( + "FILE_TRANSFER_ERROR", + "FILE_TRANSFER_PERMISSION_ERROR", + ): + return exc return OperationalError( f"Unexpected error from databricks_sql_kernel during {what}: {exc!r}" ) diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 0ca7ba6a8..a64f9e086 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -16,14 +16,14 @@ - ``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 import inspect import logging +import os +import re import threading import uuid from typing import Any, Dict, List, Optional, Set, TYPE_CHECKING, Union @@ -73,11 +73,10 @@ # http_path). Forwarding either is redundant and trips the kernel's # 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") +_STAGING_PATHS_TYPE_ERROR = ( + "staging_allowed_local_path must be a string or a list of " + "strings or os.PathLike objects" +) def _strip_leading_sql_comments(sql: str) -> str: @@ -85,12 +84,8 @@ def _strip_leading_sql_comments(sql: str) -> str: ``/* … */`` 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. + Comment-prefixed PUT/GET statements need the permitted local paths + just as statements that begin with the verb do. """ i = 0 n = len(sql) @@ -102,9 +97,15 @@ def _strip_leading_sql_comments(sql: str) -> str: 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 + depth = 1 + i += 2 + for delimiter in re.finditer(r"/\*|\*/", sql[i:]): + depth += 1 if delimiter.group() == "/*" else -1 + if depth == 0: + i += delimiter.end() + break + else: + i = n else: break return sql[i:] @@ -145,19 +146,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: @@ -575,28 +567,32 @@ 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: + stripped_sql = _strip_leading_sql_comments(operation) + staging_verb = _staging_verb(stripped_sql) stmt.set_sql(operation) + allowed_paths = ( + cursor.connection.staging_allowed_local_path + if staging_verb in ("PUT", "GET") + else None + ) + if allowed_paths is not None and not isinstance( + allowed_paths, (str, list) + ): + raise ProgrammingError(_STAGING_PATHS_TYPE_ERROR) + if isinstance(allowed_paths, str): + allowed_paths = [allowed_paths] + try: + normalized_paths = [os.fspath(path) for path in allowed_paths or []] + if any(not isinstance(path, str) for path in normalized_paths): + raise TypeError("Local paths must resolve to strings") + except (TypeError, ValueError) as exc: + raise ProgrammingError(_STAGING_PATHS_TYPE_ERROR) from exc + stmt.set_staging_allowed_local_paths( + normalized_paths, input_stream_token="__input_stream__" + ) stmt.set_row_limit(row_limit) if query_tags: # Per-statement query tags. The kernel serialises the diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 613b62386..4e7eedb3a 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1394,7 +1394,7 @@ def _handle_staging_put( session_id_hex=self.connection.get_session_id_hex(), ) - # Kernel backends reject staging operations before reaching these handlers. + # Kernel backends transfer files natively and never dispatch to these handlers. assert self.connection.http_client is not None with open(local_file, "rb") as fh: r = self.connection.http_client.request( diff --git a/tests/e2e/test_kernel_backend.py b/tests/e2e/test_kernel_backend.py index ae0301fb6..eced0810c 100644 --- a/tests/e2e/test_kernel_backend.py +++ b/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,72 @@ 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's.bin" + destination = allowed / "output's.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" + + def literal(path): + return "'" + str(path).replace("\\", "\\\\").replace("'", "\\'") + "'" + + 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") + + with sql.connect(**kernel_conn_params) as no_paths: + with no_paths.cursor() as denied_cursor: + with pytest.raises(ProgrammingError): + denied_cursor.execute( + f"PUT {literal(source)} INTO {literal(remote)} OVERWRITE" + ) + + 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 {literal(source)} INTO {literal(remote)} OVERWRITE" + ) + == [] + ) + with pytest.raises(NotSupportedError, match="Streaming PUT"): + execute(f"PUT r'__input_stream__' INTO {literal(remote)} OVERWRITE") + assert execute(f"GET {literal(remote)} TO {literal(destination)}") == [] + assert destination.read_bytes() == payload + with pytest.raises(ProgrammingError): + execute(f"PUT {literal(denied)} INTO {literal(remote)} OVERWRITE") + with pytest.raises(FileNotFoundError): + execute( + f"PUT {literal(allowed / 'missing.bin')} INTO {literal(remote)} OVERWRITE" + ) + assert execute(f"REMOVE {literal(remote)}") == [] + assert execute(f"REMOVE {literal(remote)}") == [] + with pytest.raises(OperationalError, match="404"): + execute( + f"GET {literal(remote)} TO {literal(allowed / 'deleted.bin')}" + ) + assert execute("SELECT 1")[0][0] == 1 + finally: + cur.execute(f"DROP SCHEMA {qualified} CASCADE") # ── Error fidelity — diagnostic-info reaches .context ───────────── diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index b285ce6e1..3eac2a5cf 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -13,8 +13,11 @@ from __future__ import annotations +import errno +import os import sys import types +from pathlib import Path from typing import Optional from unittest.mock import MagicMock @@ -393,9 +396,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 +808,162 @@ 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", + "operation, allowed, expected, async_op", [ - "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'", + ("PUT 'input' INTO '/Volumes/c/s/v/file'", None, [], False), + ("PUT 'input' INTO '/Volumes/c/s/v/file'", "/tmp/files", ["/tmp/files"], False), + ("GET '/Volumes/c/s/v/file' TO 'output'", [], [], False), + ( + "GET '/Volumes/c/s/v/file' TO 'output'", + [Path("/tmp/a"), Path("/tmp/b")], + ["/tmp/a", "/tmp/b"], + False, + ), + ( + "/* outer /* inner */ outer */ GET '/Volumes/c/s/v/file' TO 'output'", + ["/tmp/a", "/tmp/b"], + ["/tmp/a", "/tmp/b"], + True, + ), + ("PUT 'input' INTO '/Volumes/c/s/v/file'", "/tmp/files", ["/tmp/files"], True), + ("REMOVE '/Volumes/c/s/v/file'", object(), [], False), + ("SELECT 1", object(), [], False), ], ) -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).""" +def test_execute_command_forwards_allowed_local_paths( + operation, 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 + 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=operation, + 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, input_stream_token="__input_stream__" + ) + (stmt.submit if async_op else stmt.execute).assert_called_once() - with pytest.raises(NotSupportedError, match="staging"): + +@pytest.mark.parametrize( + "allowed", + [ + 123, + [object()], + ["/tmp/files", b"/tmp/other"], + Path("/tmp/files"), + ("/tmp/files",), + ], +) +def test_execute_command_rejects_malformed_local_paths(allowed): + c = _make_client() + c._kernel_session = MagicMock() + cursor = MagicMock() + cursor.connection.staging_allowed_local_path = allowed + stmt = c._kernel_session.statement.return_value + stmt.execute.return_value = MagicMock( + statement_id="00000000-0000-0000-0000-000000000001", + arrow_schema=MagicMock(return_value=pa.schema([("x", pa.int64())])), + ) + + with pytest.raises(ProgrammingError, match="staging_allowed_local_path"): 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, + "PUT 'input' INTO '/Volumes/c/s/v/file'", + MagicMock(), + 1, + 1, + False, + cursor, + False, + [], + False, + False, ) + stmt.execute.assert_not_called() + stmt.set_staging_allowed_local_paths.assert_not_called() @pytest.mark.parametrize( - "operation, is_staging", + "operation, expected", [ - ("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 + ("PUT '/f' INTO '/v'", "PUT"), + ("get '/v' to '/f'", "GET"), + ("REMOVE '/v'", "REMOVE"), + ("-- upload the file\nPUT '/f' INTO '/v'", "PUT"), + ("/* staging */ PUT '/f' INTO '/v'", "PUT"), + ("/* c1 */\n -- c2\n get '/v' to '/f'", "GET"), + ("/* outer /* inner */ outer */ PUT '/f' INTO '/v'", "PUT"), + ("/* a /* b /* c */ b */ a */ GET '/v' TO '/f'", "GET"), + ("/* outer /* closed */ PUT '/f' INTO '/v'", None), + (" \n\t PUT '/f' INTO '/v'", "PUT"), + ("SELECT 'GET' AS x", None), + ("SELECT * FROM puts", None), + ("-- PUT in a comment\nSELECT 1", None), + ("/* PUT */ SELECT 1", None), + ("INSERT INTO t VALUES (1)", None), + ("", None), + ("-- just a comment", None), + ("/* just a comment */", None), + ("PUTTER", None), + ("GETTING", None), + ("REMOVED", None), ], ) -def test_is_staging_statement(operation, is_staging): - assert kernel_client._is_staging_statement(operation) is is_staging +def test_staging_verb_matches_only_the_leading_statement(operation, expected): + stripped = kernel_client._strip_leading_sql_comments(operation) + assert kernel_client._staging_verb(stripped) == expected + + +def test_kernel_rejects_decoded_streaming_put(): + c = _make_client() + c._kernel_session = MagicMock() + stmt = c._kernel_session.statement.return_value + error = _FakeKernelError( + code="InvalidArgument", message="Streaming PUT is not supported" + ) + error.error_code = "FILE_TRANSFER_UNSUPPORTED_STREAM" + stmt.execute.side_effect = error + cursor = MagicMock() + cursor.connection.staging_allowed_local_path = None + with pytest.raises(NotSupportedError, match="Streaming PUT"): + c.execute_command( + "PUT r'__input_stream__' INTO '/Volumes/c/s/v/file'", + MagicMock(), + 1, + 1, + False, + cursor, + False, + [], + False, + False, + ) + stmt.set_staging_allowed_local_paths.assert_called_once_with( + [], input_stream_token="__input_stream__" + ) + stmt.execute.assert_called_once() + stmt.close.assert_called_once() # --------------------------------------------------------------------------- @@ -2582,3 +2675,117 @@ 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, error_code, expected", + [ + ("PermissionDenied", None, "FILE_TRANSFER_ERROR", ProgrammingError), + ("InvalidArgument", None, "FILE_TRANSFER_ERROR", ProgrammingError), + ("InvalidArgument", 400, "FILE_TRANSFER_ERROR", OperationalError), + ("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), + ("PermissionDenied", 403, "FILE_TRANSFER_PERMISSION_ERROR", OperationalError), + ("NetworkError", None, "FILE_TRANSFER_ERROR", OperationalError), + ("Internal", None, "FILE_TRANSFER_ERROR", OperationalError), + ("Internal", None, None, DatabaseError), + ], +) +def test_file_transfer_error_matches_thrift_exception( + code, http_status, error_code, expected +): + error = _FakeKernelError(code=code, message="File transfer failed") + error.error_code = error_code + error.http_status = http_status + error.os_error_filename = "local-file.bin" + result = kernel_client._reraise_kernel_error(error) + assert type(result) is expected + assert result.error_code == error_code + assert result.code == code + assert result.http_status == http_status + if issubclass(expected, OSError): + assert result.errno is None + assert result.filename == "local-file.bin" + + +@pytest.mark.parametrize( + "code,error_code,raw,windows_raw,expected", + [ + ("NotFound", "FILE_TRANSFER_ERROR", errno.ENOENT, 2, FileNotFoundError), + ( + "PermissionDenied", + "FILE_TRANSFER_PERMISSION_ERROR", + errno.EACCES, + 5, + PermissionError, + ), + ( + "PermissionDenied", + "FILE_TRANSFER_PERMISSION_ERROR", + errno.EPERM, + 5, + PermissionError, + ), + ("Internal", "FILE_TRANSFER_ERROR", errno.ENOTDIR, 267, NotADirectoryError), + ("Internal", "FILE_TRANSFER_ERROR", errno.EISDIR, 5, IsADirectoryError), + ("ResourceExhausted", "FILE_TRANSFER_ERROR", errno.ENOSPC, 112, OSError), + ], +) +def test_local_transfer_error_preserves_os_code( + code, error_code, raw, windows_raw, expected +): + error = _FakeKernelError(code=code, message="File transfer failed") + error.error_code = error_code + if os.name == "nt": + raw = windows_raw + if expected is IsADirectoryError: + expected = PermissionError + error.os_error_code = raw + error.os_error_filename = "local-file.bin" + result = kernel_client._reraise_kernel_error(error) + assert type(result) is expected + assert result.os_error_code == raw + assert result.filename == "local-file.bin" + if os.name == "nt": + assert result.winerror == raw + assert result.errno == OSError(0, "", None, raw).errno + else: + assert result.errno == raw + + +@pytest.mark.parametrize( + "code,error_code,raw", + [ + ("NotFound", "FILE_TRANSFER_ERROR", None), + ("PermissionDenied", "FILE_TRANSFER_PERMISSION_ERROR", None), + ("Internal", "FILE_TRANSFER_ERROR", errno.ENOTDIR), + ("ResourceExhausted", "FILE_TRANSFER_ERROR", errno.ENOSPC), + ], +) +def test_nested_wrapping_preserves_mapped_local_transfer_error(code, error_code, raw): + from databricks.sql.backend.kernel._errors import wrap_kernel_exception + + error = _FakeKernelError(code=code, sql_state="HY000", query_id="transfer-id") + error.error_code = error_code + error.os_error_code = raw + mapped = wrap_kernel_exception("inner", error) + assert isinstance(mapped, OSError) + wrapped = wrap_kernel_exception("outer", mapped) + assert wrapped is mapped + assert wrapped.code == code + assert wrapped.error_code == error_code + assert wrapped.sql_state == "HY000" + assert wrapped.query_id == "transfer-id" + + +@pytest.mark.parametrize( + "error", [FileNotFoundError("missing"), PermissionError("denied")] +) +def test_raw_local_errors_keep_unexpected_error_wrapping(error): + from databricks.sql.backend.kernel._errors import wrap_kernel_exception + + wrapped = wrap_kernel_exception("execute_command", error) + assert type(wrapped) is OperationalError + assert "execute_command" in str(wrapped)