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..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 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 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). | @@ -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,17 @@ 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 `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 +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/KERNEL_REV b/KERNEL_REV index de1b52ca7..3232e69ae 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -0b5d0edfed32c0745d5abc2468fc3070e2eb2a17 +566fe7f0a8a2fa246db95045a7000136ee6070d5 diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index df52f34a5..01d98c477 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,21 @@ 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() + if entry[1] is not None + ] + 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 +599,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 +640,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.get(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 +719,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 +783,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 and volume[1] 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..76accc06b 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,72 @@ 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 []) + ) + 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: + 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.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 + ) + + 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 + and getattr(exc, "error_code", None) == "VOLUME_STORAGE_HTTP_ERROR" + ): + 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,6 +312,13 @@ def close(self) -> None: self.has_been_closed_server_side = True self.status = CommandState.CLOSED return + backend = cast("KernelDatabricksClient", self.backend) + 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) + except Exception: + logger.debug("Volume cancellation during close failed", exc_info=True) try: self._kernel_handle.close() except Exception as exc: @@ -251,12 +326,6 @@ 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) - # 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. - backend = cast("KernelDatabricksClient", self.backend) try: backend.close_command(self.command_id) except Exception as exc: 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..db46bbdb9 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())])), ) @@ -1639,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() @@ -1646,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(): @@ -2781,3 +2723,48 @@ 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]) +@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, + 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 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 cfa2d3f2e..31d8e68bf 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,138 @@ 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_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" + 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" + assert rs.fetchall() == [] + rs.backend.finish_volume_operation.assert_called_once_with(rs.command_id) + + +@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) + 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() + + +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"): + 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