diff --git a/CHANGELOG.md b/CHANGELOG.md index 2423c20e3..043a6569e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,8 @@ # Release History +# Unreleased +- Fix: `REMOVE` staging operations no longer require `staging_allowed_local_path` to be set, since removing a remote file does not touch the local filesystem (databricks/databricks-sql-python#726) + # 4.3.0 (2026-06-12) - **New: optional Rust kernel backend (`use_kernel=True`).** Adds an alternative connection path backed by the native [`databricks-sql-kernel`](https://pypi.org/project/databricks-sql-kernel/) client (a Rust core exposed via PyO3), installable with the new `databricks-sql-connector[kernel]` extra. The kernel talks to Databricks over the **SEA (Statement Execution API) HTTP transport** — not Thrift — with CloudFetch and inline-Arrow result fetching, so `use_kernel=True` gives you a modern SEA-native client through the same DB-API surface. Supports PAT, OAuth M2M, and OAuth U2M auth. Requires Python >= 3.10 (the kernel wheel is `cp310-abi3`); on older interpreters the extra is a no-op and `use_kernel=True` raises a clear `ImportError`. The default backend remains Thrift — opt in per connection. - Kernel backend behavior is aligned with the Thrift backend so application code works the same either way: consistent cursor-state tracking (`query_id` / `get_query_state`), metadata (catalogs/schemas/tables/columns with JDBC-style filter semantics and case-insensitive `table_types`), DML `rowcount`, server-sourced async execution state, sync `cancel()`, fail-loud staging/volume operations, and structured error context (SQLSTATE, diagnostic info). Kernel logs surface through Python `logging` under the `databricks.sql.kernel` logger (databricks/databricks-sql-python#824, #825, #830, #838, #839 by @vikrantpuppala) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 7b94fc98d..a6c7b2997 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1081,6 +1081,14 @@ def _handle_staging_operation( headers=headers, ) + # REMOVE deletes a remote resource and never touches the local + # filesystem, so it does not require staging_allowed_local_path. + if row.operation == "REMOVE": + return self._handle_staging_remove( + presigned_url=row.presignedUrl, + headers=headers, + ) + # For non-streaming operations, validate staging_allowed_local_path if isinstance(staging_allowed_local_path, type(str())): _staging_allowed_local_paths = [staging_allowed_local_path] @@ -1133,14 +1141,12 @@ def _handle_staging_operation( ) # TODO: Create a retry loop here to re-attempt if the request times out or fails + # REMOVE is handled above, before staging_allowed_local_path validation, + # since it does not touch the local filesystem. if row.operation == "GET": return self._handle_staging_get(**handler_args) elif row.operation == "PUT": return self._handle_staging_put(**handler_args) - elif row.operation == "REMOVE": - # Local file isn't needed to remove a remote resource - handler_args.pop("local_file") - return self._handle_staging_remove(**handler_args) else: raise ProgrammingError( f"Operation {row.operation} is not supported. " diff --git a/tests/unit/test_staging_remove.py b/tests/unit/test_staging_remove.py new file mode 100644 index 000000000..0a5114621 --- /dev/null +++ b/tests/unit/test_staging_remove.py @@ -0,0 +1,97 @@ +from unittest.mock import patch, Mock, MagicMock + +import pytest + +import databricks.sql.client as client + + +class TestStagingRemove: + """Unit tests for staging REMOVE operations. + + REMOVE deletes a remote resource and never touches the local filesystem, + so it must not require ``staging_allowed_local_path`` to be configured + (see GitHub issue #726). + """ + + @pytest.fixture + def cursor(self): + return client.Cursor(connection=Mock(), backend=Mock()) + + def _setup_mock_remove_response(self, cursor): + """Set up a mock staging REMOVE response with no localFile.""" + mock_result_set = Mock() + mock_result_set.is_staging_operation = True + cursor.backend.execute_command.return_value = mock_result_set + + mock_row = Mock() + mock_row.operation = "REMOVE" + # The server does not include a localFile for REMOVE. + mock_row.localFile = None + mock_row.presignedUrl = "https://example.com/delete" + mock_row.headers = "{}" + mock_result_set.fetchone.return_value = mock_row + + cursor.active_result_set = mock_result_set + return mock_result_set + + def test_remove_without_staging_allowed_local_path(self, cursor): + """REMOVE must succeed when staging_allowed_local_path is None.""" + + self._setup_mock_remove_response(cursor) + + with patch.object(cursor, "_handle_staging_remove") as mock_remove: + cursor._handle_staging_operation(staging_allowed_local_path=None) + + mock_remove.assert_called_once() + # local_file must not be forwarded to the REMOVE handler. + assert "local_file" not in mock_remove.call_args.kwargs + assert ( + mock_remove.call_args.kwargs["presigned_url"] + == "https://example.com/delete" + ) + + def test_remove_with_staging_allowed_local_path(self, cursor): + """REMOVE must still work when staging_allowed_local_path is provided.""" + + self._setup_mock_remove_response(cursor) + + with patch.object(cursor, "_handle_staging_remove") as mock_remove: + cursor._handle_staging_operation(staging_allowed_local_path="/tmp") + + mock_remove.assert_called_once() + + 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.is_staging_operation = True + mock_row = Mock() + mock_row.operation = "GET" + mock_row.localFile = "/tmp/file.csv" + mock_row.presignedUrl = "https://example.com/download" + mock_row.headers = "{}" + mock_result_set.fetchone.return_value = mock_row + cursor.active_result_set = mock_result_set + + with pytest.raises(client.ProgrammingError) as excinfo: + cursor._handle_staging_operation(staging_allowed_local_path=None) + assert "You must provide at least one staging_allowed_local_path" in str( + excinfo.value + ) + + 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.is_staging_operation = True + mock_row = Mock() + mock_row.operation = "COPY" + mock_row.localFile = None + mock_row.presignedUrl = "https://example.com/copy" + mock_row.headers = "{}" + mock_result_set.fetchone.return_value = mock_row + cursor.active_result_set = mock_result_set + + with pytest.raises(client.ProgrammingError) as excinfo: + cursor._handle_staging_operation(staging_allowed_local_path="/tmp") + assert "is not supported" in str(excinfo.value)