diff --git a/CHANGELOG.md b/CHANGELOG.md index 4fbb64b68..e8710805a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ # Release History # Unreleased +- Support SQL PUT, GET, and REMOVE file transfers with a native kernel built from `KERNEL_REV`. Published kernel 1.4.0 does not support these transfers; the release version and dependency floor are pending [kernel #443](https://github.com/databricks/databricks-sql-kernel/issues/443). Older wheels reject transfers and continue to support ordinary queries. 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..048834a36 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]` | ✅ | Source build | `None` | Local path(s) permitted for Unity Catalog Volume `PUT`/`GET`. Requires a native kernel built from `KERNEL_REV` with [kernel #439](https://github.com/databricks/databricks-sql-kernel/pull/439). Published 1.4.0 lacks this capability; the release version and dependency floor are pending [kernel #443](https://github.com/databricks/databricks-sql-kernel/issues/443). Older wheels reject transfers. 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..0c7999481 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -6f4b6d607a61bdde61dcda2cf82bc8e80fd32cb9 +9f7bf932aacbd76ff67c52cc1cd84082af2e5ab0 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..57f2b68ef 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 missing transfer files use the built-in ``FileNotFoundError`` class; +the original OS error code is preserved when available. The ``filename`` +attribute is unavailable. + 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,6 +42,8 @@ from __future__ import annotations import logging +import os +from typing import Union from databricks.sql.exc import ( DatabaseError, @@ -115,10 +120,13 @@ } -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 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 @@ -126,6 +134,18 @@ 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) + 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 +161,36 @@ 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) + if os_error_code is None: + file_error_args = (message,) + elif os.name == "nt": + file_error_args = (0, message, None, os_error_code) + else: + file_error_args = (os_error_code, message) + new: Union[Error, OSError] + # Local missing files have no HTTP status; remote missing files carry 404. + if ( + 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", "http_status", "retryable", "query_id", @@ -162,12 +205,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. + A local missing transfer file maps to ``FileNotFoundError`` with + its original OS error code when available, without ``filename``. + - 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 +225,10 @@ def wrap_kernel_exception(what: str, exc: BaseException) -> "Error": return reraise_kernel_error(exc) if isinstance(exc, Error): return exc + if isinstance(exc, (FileNotFoundError, PermissionError)) 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..7680dd714 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,46 @@ 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) + if staging_verb == "PUT" and re.match( + r"""(['"])__input_stream__\1""", + _strip_leading_sql_comments(stripped_sql[3:]), + ): + raise NotSupportedError( + "Streaming PUT is not supported by the kernel backend." + ) stmt.set_sql(operation) + allowed_paths = ( + cursor.connection.staging_allowed_local_path + if staging_verb in ("PUT", "GET") + else None + ) + set_paths = getattr(stmt, "set_staging_allowed_local_paths", None) + if callable(set_paths): + 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 + set_paths(normalized_paths) + elif staging_verb is not None: + raise NotSupportedError( + "Volume transfers require a databricks-sql-kernel wheel " + "with staging_allowed_local_paths support." + ) 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..5e1a1c90f 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,62 @@ 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") + + 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 '{source}' INTO '{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 '{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 ───────────── diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index b285ce6e1..b11e5de00 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,212 @@ 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( + "allowed, expected", + [ + (None, []), + ("/tmp/files", ["/tmp/files"]), + ([Path("/tmp/a"), Path("/tmp/b")], ["/tmp/a", "/tmp/b"]), + (["/tmp/a", "/tmp/b"], ["/tmp/a", "/tmp/b"]), + ], +) +@pytest.mark.parametrize("async_op", [False, True]) +@pytest.mark.parametrize( + "operation, uses_paths", + [ + ("PUT '/tmp/files/input' INTO '/Volumes/c/s/v/input'", True), + ("PUT '__INPUT_STREAM__' INTO '/Volumes/c/s/v/input'", True), + ("GET '/Volumes/c/s/v/input' TO '/tmp/files/output'", True), + ("/* download */ GET '/Volumes/c/s/v/input' TO '/tmp/files/output'", True), + ( + "/* outer /* inner */ outer */ GET '/Volumes/c/s/v/input' TO '/tmp/files/output'", + True, + ), + ("REMOVE '/Volumes/c/s/v/input'", False), + ("SELECT 1", False), + ], +) +def test_execute_command_forwards_allowed_local_paths( + allowed, expected, async_op, operation, uses_paths +): + 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 if uses_paths else [] + ) + (stmt.submit if async_op else stmt.execute).assert_called_once() +@pytest.mark.parametrize( + "allowed", + [ + 123, + [object()], + [b"/tmp/files"], + ["/tmp/files", b"/tmp/other"], + Path("/tmp/files"), + ("/tmp/files",), + (), + {"/tmp/files": True}, + ], +) @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'", + "SELECT 1", + "PUT 'input' INTO '/Volumes/c/s/v/file'", + "GET '/Volumes/c/s/v/file' TO 'output'", + "REMOVE '/Volumes/c/s/v/file'", ], ) -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_malformed_local_paths_only_reject_transfers(allowed, operation): c = _make_client() c._kernel_session = MagicMock() cursor = MagicMock() - cursor.arraysize = 100 - cursor.buffer_size_bytes = 1024 + 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(NotSupportedError, match="staging"): + def execute(): 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, + operation, MagicMock(), 1, 1, False, cursor, False, [], False, False + ) + + if operation.startswith(("PUT", "GET")): + with pytest.raises(ProgrammingError, match="staging_allowed_local_path"): + execute() + stmt.execute.assert_not_called() + stmt.set_staging_allowed_local_paths.assert_not_called() + else: + execute() + stmt.set_staging_allowed_local_paths.assert_called_once_with([]) + stmt.execute.assert_called_once() + + +@pytest.mark.parametrize( + "operation, expected", + [ + ("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_staging_verb_matches_only_the_leading_statement(operation, expected): + stripped = kernel_client._strip_leading_sql_comments(operation) + assert kernel_client._staging_verb(stripped) == expected + + +@pytest.mark.parametrize( + "query", + [ + "PUT '__input_stream__' INTO '/Volumes/c/s/v/file'", + "put '__input_stream__' INTO '/Volumes/c/s/v/file'", + "PUT '__input_stream__'INTO '/Volumes/c/s/v/file'", + 'PUT"__input_stream__"INTO "/Volumes/c/s/v/file"', + "/* upload */ PUT '__input_stream__' INTO '/Volumes/c/s/v/file'", + "/* outer /* inner */ outer */ PUT '__input_stream__' INTO '/Volumes/c/s/v/file'", + "PUT'__input_stream__' INTO '/Volumes/c/s/v/file'", + "PUT /* upload */ '__input_stream__' INTO '/Volumes/c/s/v/file'", + "PUT\n-- upload\n'__input_stream__' INTO '/Volumes/c/s/v/file'", + "/* upload */ PUT\"__input_stream__\" INTO '/Volumes/c/s/v/file'", + ], +) +def test_kernel_rejects_streaming_put(query): + c = _make_client() + c._kernel_session = MagicMock() + with pytest.raises(NotSupportedError, match="Streaming PUT"): + c.execute_command( + query, MagicMock(), 1, 1, False, MagicMock(), False, [], False, False ) + stmt = c._kernel_session.statement.return_value + stmt.execute.assert_not_called() + stmt.close.assert_called_once() @pytest.mark.parametrize( - "operation, is_staging", + "operation", [ - ("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 + "SELECT 1", + "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_is_staging_statement(operation, is_staging): - assert kernel_client._is_staging_statement(operation) is is_staging +def test_execute_command_with_older_kernel(operation): + c = _make_client() + c._kernel_session = MagicMock() + stmt = c._kernel_session.statement.return_value + del stmt.set_staging_allowed_local_paths + stmt.execute.return_value = MagicMock( + statement_id="00000000-0000-0000-0000-000000000001", + arrow_schema=MagicMock(return_value=pa.schema([("x", pa.int64())])), + ) + kwargs = dict( + operation=operation, + session_id=MagicMock(), + max_rows=1, + max_bytes=1, + lz4_compression=False, + cursor=MagicMock(), + use_cloud_fetch=False, + parameters=[], + async_op=False, + enforce_embedded_schema_correctness=False, + ) + if operation == "SELECT 1": + c.execute_command(**kwargs) + stmt.execute.assert_called_once() + else: + with pytest.raises(NotSupportedError, match="Volume transfers require"): + c.execute_command(**kwargs) + stmt.execute.assert_not_called() + stmt.close.assert_called_once() # --------------------------------------------------------------------------- @@ -2582,3 +2725,101 @@ 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 + 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 + + +@pytest.mark.parametrize( + "code,error_code,raw,expected", + [ + ("NotFound", "FILE_TRANSFER_ERROR", errno.ENOENT, FileNotFoundError), + ( + "PermissionDenied", + "FILE_TRANSFER_PERMISSION_ERROR", + errno.EACCES, + PermissionError, + ), + ( + "PermissionDenied", + "FILE_TRANSFER_PERMISSION_ERROR", + errno.EPERM, + PermissionError, + ), + ], +) +def test_local_transfer_error_preserves_os_code(code, error_code, raw, expected): + error = _FakeKernelError(code=code, message="File transfer failed") + error.error_code = error_code + if os.name == "nt": + raw = 2 if expected is FileNotFoundError else 5 + error.os_error_code = raw + result = kernel_client._reraise_kernel_error(error) + assert type(result) is expected + assert result.os_error_code == raw + assert result.filename is None + if os.name == "nt": + assert result.winerror == raw + assert result.errno == (errno.ENOENT if raw == 2 else errno.EACCES) + else: + assert result.errno == raw + + +@pytest.mark.parametrize( + "code,error_code", + [ + ("NotFound", "FILE_TRANSFER_ERROR"), + ("PermissionDenied", "FILE_TRANSFER_PERMISSION_ERROR"), + ], +) +def test_nested_wrapping_preserves_mapped_local_transfer_error(code, error_code): + 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 + mapped = wrap_kernel_exception("inner", error) + 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)