Fix: REMOVE staging operation no longer requires staging_allowed_local_path (#726)

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)

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. "

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)