fix(kernel): match documented local path option types

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>

Author: aakash-saravanan-db

src/databricks/sql/backend/kernel/client.py

@@ -587,20 +587,30 @@ def execute_command(
                 stmt.set_sql(operation)
                 allowed_paths = (
                     cursor.connection.staging_allowed_local_path
-                    if _is_staging_statement(operation)
+                    if re.match(
+                        r"(?:PUT|GET)\b",
+                        _strip_leading_sql_comments(operation),
+                        re.IGNORECASE,
+                    )
                     else None
                 )
-                if isinstance(allowed_paths, (str, os.PathLike)):
-                    allowed_paths = [allowed_paths]
                 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_allowed_local_path must be a string or a list of paths"
+                        )
+                    if isinstance(allowed_paths, str):
+                        allowed_paths = [allowed_paths]
                     try:
                         normalized_paths = [
                             os.fsdecode(path) for path in allowed_paths or []
                         ]
                     except (TypeError, ValueError) as exc:
                         raise ProgrammingError(
-                            "staging_allowed_local_path must be a path or a list of paths"
+                            "staging_allowed_local_path must be a string or a list of paths"
                         ) from exc
                     set_paths(normalized_paths)
                 elif _is_staging_statement(operation):

tests/unit/test_kernel_client.py

@@ -811,7 +811,6 @@ def test_execute_command_forwards_row_limit(row_limit):
     [
         (None, []),
         ("/tmp/files", ["/tmp/files"]),
-        (Path("/tmp/files"), ["/tmp/files"]),
         ([Path("/tmp/a"), Path("/tmp/b")], ["/tmp/a", "/tmp/b"]),
         (["/tmp/a", "/tmp/b"], ["/tmp/a", "/tmp/b"]),
     ],
@@ -847,9 +846,18 @@ def test_execute_command_forwards_allowed_local_paths(allowed, expected, async_o
     (stmt.submit if async_op else stmt.execute).assert_called_once()
 
 
-@pytest.mark.parametrize("allowed", [123, [object()]])
 @pytest.mark.parametrize(
-    "operation", ["SELECT 1", "PUT 'input' INTO '/Volumes/c/s/v/file'"]
+    "allowed",
+    [123, [object()], Path("/tmp/files"), ("/tmp/files",), (), {"/tmp/files": True}],
+)
+@pytest.mark.parametrize(
+    "operation",
+    [
+        "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_malformed_local_paths_only_reject_transfers(allowed, operation):
     c = _make_client()
@@ -867,7 +875,7 @@ def execute():
             operation, MagicMock(), 1, 1, False, cursor, False, [], False, False
         )
 
-    if operation.startswith("PUT"):
+    if operation.startswith(("PUT", "GET")):
         with pytest.raises(ProgrammingError, match="staging_allowed_local_path"):
             execute()
         stmt.execute.assert_not_called()