fsspec-databricks 0.2.1__tar.gz → 0.3.0__tar.gz

{fsspec_databricks-0.2.1 → fsspec_databricks-0.3.0}/DEVELOPMENT.md

@@ -147,6 +147,9 @@ Prefer federation: it needs no stored secret. GitHub's OIDC token identifies eac
 
 - **`github-oidc`**: a federation policy on the service principal, in the Databricks account console. Its audience is
   `DATABRICKS_TOKEN_AUDIENCE`, e.g. the account ID.
+  The SDK's `github-oidc` exchanges a new token for each request, which slows the tests down.
+  So `.github/scripts/with-id-token-file.sh` keeps GitHub's token in a file and runs the tests with `file-oidc`, which
+  caches the exchanged token.
 - **`github-oidc-azure`**: a federated credential on the Entra ID app, with the audience `api://AzureADTokenExchange`.
   The app is the workspace's service principal. This needs no access to the account console.
 
@@ -216,10 +219,14 @@ exception.
 
 ### Picklable
 
-File systems are pickled to run in other processes, e.g. by Dask. Clients, event loops, threads and sessions are left
-out, and made again on first use.
+File systems are pickled to run in other processes, e.g. by Dask. fsspec pickles a file system as its constructor
+arguments, and unpickling calls the constructor again:
 
-A `WorkspaceClient` given as `client` is kept, and is not picklable. Other ways of authenticating are.
+- clients, event loops, threads and sessions are not pickled, and are made again on first use;
+- attributes set after construction are lost;
+- every constructor argument must be picklable.
+
+A `WorkspaceClient` given as `client` is not picklable. Other ways of authenticating are.
 
 ### Known limits
 
@@ -271,7 +278,6 @@ Within one backend, both go to the backend.
 
 - the Databricks SDK `Config` and `WorkspaceClient`, made on first use from `client`, `config` or `config_params`;
 - `close()`, also called at exit and on garbage collection;
-- pickling without the client and config;
 - the `dbfs` protocol: `_strip_protocol` and `unstrip_protocol` use `path.py`.
 
 Its `mkdir`, `makedirs`, `rmdir` and `_rm` only check preconditions, such as a missing parent or a file where a folder
@@ -309,17 +315,20 @@ when no glob pattern is involved. Files are opened with the SDK's `dbfs.open`.
 
 ### Volumes
 
-`VolumeFileSystem`, in `volume.py`, subclasses `DBFS`, as the DBFS API also lists Volume folders.
+`VolumeFileSystem`, in `volume.py`, subclasses `DBFS`. It keeps the SDK's `dbfs` for `exists` and for creating
+folders. Listings, copies, moves and deletes inside a volume use the Files API and the backend's own file objects.
 
 The APIs called depend on the depth of the path:
 
-| Path                                   | `info`                                                    | `ls`            |
-|----------------------------------------|-----------------------------------------------------------|-----------------|
-| `/Volumes`                             | None: a fixed folder                                      | `catalogs.list` |
-| `/Volumes/<catalog>`                   | `catalogs.get`                                            | `schemas.list`  |
-| `/Volumes/<catalog>/<schema>`          | `schemas.get`                                             | `volumes.list`  |
-| `/Volumes/<catalog>/<schema>/<volume>` | `volumes.read`                                            | `dbfs.list`     |
-| Deeper                                 | `files.get_metadata`, else `files.get_directory_metadata` | `dbfs.list`     |
+| Path                                   | `info`                                                    | `ls`                            |
+|----------------------------------------|-----------------------------------------------------------|---------------------------------|
+| `/Volumes`                             | None: a fixed folder                                      | `catalogs.list`                 |
+| `/Volumes/<catalog>`                   | `catalogs.get`                                            | `schemas.list`                  |
+| `/Volumes/<catalog>/<schema>`          | `schemas.get`                                             | `volumes.list`                  |
+| `/Volumes/<catalog>/<schema>/<volume>` | `volumes.read`                                            | `files.list_directory_contents` |
+| Deeper                                 | `files.get_metadata`, else `files.get_directory_metadata` | `files.list_directory_contents` |
+
+`ls` of a file returns the file itself, from `files.get_metadata` after the listing fails.
 
 `include_browse` is passed to the Unity Catalog APIs.
 
@@ -332,9 +341,22 @@ Each metadata call is a round trip, so the checks inside a volume make as few as
 | `isdir`, so also `mkdir` and `rmdir`          | 1: `files.get_directory_metadata`               |
 | `isfile`                                      | 1: `files.get_metadata`                         |
 | Opening for writing                           | 1: `isdir`, to refuse directories               |
+| `ls` of a directory                           | 1: `files.list_directory_contents`              |
 
 `DatabricksFileSystem` passes `isdir` and `isfile` to the backends. `test_volume_fs_metadata_calls` counts the calls.
 
+Copies, moves and deletes inside a volume:
+
+- `cp_file` streams the file from a `VolumeReadableFile` into a `VolumeWritableFile`, as the Files API has no copy. The
+  SDK's `dbfs.copy` would hold the whole file in memory.
+- `mv` is fsspec's default: a copy, then a delete.
+- `rm` of a file checks it with `info`, then calls `files.delete`.
+- A recursive `rm` lists the tree one depth at a time, deletes the files, then the folders from the deepest. Each step
+  runs up to `_max_delete_concurrency` requests in a thread pool.
+- Paths of volumes, schemas and catalogs go to `DBFS`.
+
+`test_volume_fs_rm_calls` and `test_volume_fs_cp_file_streams` check these on a fake `WorkspaceClient.files`.
+
 #### The IO event loop
 
 Each `VolumeFileSystem` runs an asyncio event loop in a daemon thread, `dbfs-io-thread`, started on first use. It
@@ -370,6 +392,9 @@ loop:
 `_compute_block_size` gives the block size: the smallest power-of-2 multiple of the minimum that covers the target with
 `max_concurrency` blocks, capped at the maximum.
 
+A range request that gets a `429`, `500`, `502`, `503` or `504` is retried up to 3 times, with exponential backoff and
+jitter. Writes are not retried: a repeated `PUT` can fail with "already exists" when the first one went through.
+
 #### Writing
 
 `VolumeWritableFile` subclasses `AbstractAsyncWritableFile`. Writes are buffered:

{fsspec_databricks-0.2.1 → fsspec_databricks-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: fsspec-databricks
-Version: 0.2.1
+Version: 0.3.0
 Summary: `fsspec`-based file system interface for Databricks file system"s"
 Project-URL: Homepage, https://github.com/sugibuchi/fsspec-databricks
 Author-email: Tsuyoshi SUGIBUCHI <sugibuchi@gmail.com>

{fsspec_databricks-0.2.1 → fsspec_databricks-0.3.0}/USAGE.md

@@ -166,6 +166,14 @@ fs.ls("dbfs:/Volumes/my_catalog/my_schema/my_volume/data")
 Reads and writes run concurrent requests. Large writes use multipart upload, or resumable upload where the workspace
 requires it.
 
+Reads retry after transient server errors. Writes do not.
+
+Copies, moves and deletes:
+
+- A copy reads the file and writes it again, as the Files API cannot copy. It uses the options for reads and writes.
+- A move is a copy, then a delete. It replaces an existing file at the destination.
+- A recursive delete sends concurrent requests.
+
 ### Modes
 
 | Mode  | Supported | Notes                            |

{fsspec_databricks-0.2.1 → fsspec_databricks-0.3.0}/src/fsspec_databricks/base.py

@@ -90,17 +90,6 @@ class AbstractDatabricksFileSystem(AbstractFileSystem):
         _register_cleanup(self)
         self.log.debug("Opened: %s", self)
 
-    def __getstate__(self):
-        state = self.__dict__.copy()
-        del state["_AbstractDatabricksFileSystem__client"]
-        del state["_AbstractDatabricksFileSystem__config"]
-        return state
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-        self.__client = None
-        self.__config = None
-
     def close(self):
         """Close the file system and release any resources."""
         if not self.closed:
@@ -197,6 +186,16 @@ class AbstractDatabricksFileSystem(AbstractFileSystem):
     def cp_file(self, path1, path2, **kwargs):
         self.log.debug("Copying file: path1=%s, path2=%s", path1, path2)
 
+    def _cp_directory(self, path2):
+        """Copy a directory: create the destination, as directories have no content."""
+        try:
+            info = self.info(path2)
+        except FileNotFoundError:
+            self.mkdirs(path2)
+            return
+        if info["type"] != "directory":
+            raise file_exists_error(path2)
+
     def _rm(self, path):
         # Implement only the precondition check here
         # Actual operations must be implemented by subclasses

{fsspec_databricks-0.2.1 → fsspec_databricks-0.3.0}/src/fsspec_databricks/dbfs.py

@@ -8,7 +8,7 @@ from databricks.sdk.config import Config
 from databricks.sdk.service.files import FileInfo
 
 from .base import AbstractDatabricksFileSystem
-from .error import error_mapping, file_exists_error, is_a_directory_error
+from .error import error_mapping, is_a_directory_error
 from .path import parse_dbfs_path
 from .utils import to_datetime
 
@@ -80,15 +80,8 @@ class DBFS(AbstractDatabricksFileSystem):
 
         source_info = self.info(path1)
         if source_info["type"] == "directory":
-            try:
-                dest_info = self.info(path2)
-                if dest_info["type"] == "directory":
-                    return
-                else:
-                    raise file_exists_error(path2)
-            except FileNotFoundError:
-                self.mkdirs(path2)
-                return
+            self._cp_directory(path2)
+            return
 
         with error_mapping(path1):
             self.client.dbfs.copy(

{fsspec_databricks-0.2.1 → fsspec_databricks-0.3.0}/src/fsspec_databricks/spec.py

@@ -10,7 +10,7 @@ from fsspec.spec import AbstractFileSystem
 
 from .base import AbstractDatabricksFileSystem, root_dir_info
 from .dbfs import DBFS
-from .error import file_exists_error, file_not_found_error
+from .error import file_not_found_error
 from .local import DatabricksLocalFileSystem
 from .path import (
     FileSystemType,
@@ -293,23 +293,6 @@ class DatabricksFileSystem(AbstractDatabricksFileSystem):
             self.__local_fs = None
             super().close()
 
-    def __getstate__(self):
-        state = super().__getstate__()
-        del state["_DatabricksFileSystem__local_fs"]
-        del state["_DatabricksFileSystem__workspace_fs"]
-        del state["_DatabricksFileSystem__volume_fs"]
-        del state["_DatabricksFileSystem__legacy_fs"]
-        del state["_running_in_workspace"]
-        return state
-
-    def __setstate__(self, state):
-        super().__setstate__(state)
-        self.__workspace_fs = None
-        self.__volume_fs = None
-        self.__legacy_fs = None
-        self.__local_fs = None
-        self._running_in_workspace = None
-
     @property
     def _volume_fs(self) -> VolumeFileSystem:
         """Get or create the UnityCatalogVolumeFileSystem instance."""
@@ -495,15 +478,8 @@ class DatabricksFileSystem(AbstractDatabricksFileSystem):
         else:
             source_info = self.info(path1)
             if source_info["type"] == "directory":
-                try:
-                    dest_info = self.info(path2)
-                    if dest_info["type"] == "directory":
-                        return
-                    else:
-                        raise file_exists_error(path2)
-                except FileNotFoundError:
-                    self.mkdirs(path2)
-                    return
+                self._cp_directory(path2)
+                return
 
             src_fs, src_path = self._resolve(path1)
             dst_fs, dst_path = self._resolve(path2)

{fsspec_databricks-0.2.1 → fsspec_databricks-0.3.0}/src/fsspec_databricks/volume.py

@@ -1,13 +1,20 @@
 import asyncio
+import glob
 import logging
 import re
+import shutil
 import urllib.parse
+from collections.abc import Callable, Iterable
+from concurrent.futures import ThreadPoolExecutor
 from datetime import UTC, datetime, timedelta
 from io import UnsupportedOperation
+from os import PathLike
+from random import random
 from threading import Event, RLock, Thread
 from typing import Any, Literal
 
 from aiohttp import (
+    ClientConnectionError,
     ClientHandlerType,
     ClientMiddlewareType,
     ClientRequest,
@@ -56,61 +63,49 @@ def _map_info(
     detail: bool,
     path: str | None = None,
 ) -> dict[str, Any] | str:
-    """Map Unity Catalog CatalogInfo, SchemaInfo, VolumeInfo or GetMetadataResponse to fsspec file info dict."""
-    if isinstance(info, CatalogInfo):
-        if detail:
-            return {
-                "name": f"/Volumes/{info.name}",
-                "size": None,
-                "type": "directory",
-                "created": to_datetime(info.created_at),
-                "modified": to_datetime(info.updated_at),
-                "islink": False,
-            }
-        else:
-            return f"/Volumes/{info.name}"
-    elif isinstance(info, SchemaInfo):
-        if detail:
-            return {
-                "name": f"/Volumes/{info.catalog_name}/{info.name}",
-                "size": None,
-                "type": "directory",
-                "created": to_datetime(info.created_at),
-                "modified": to_datetime(info.updated_at),
-                "islink": False,
-            }
+    """Map Unity Catalog CatalogInfo, SchemaInfo, VolumeInfo, DirectoryEntry or GetMetadataResponse
+    to fsspec file info dict, or to the name only if not detail."""
+    if isinstance(info, CatalogInfo | SchemaInfo | VolumeInfo):
+        if isinstance(info, CatalogInfo):
+            name = f"/Volumes/{info.name}"
+        elif isinstance(info, SchemaInfo):
+            name = f"/Volumes/{info.catalog_name}/{info.name}"
         else:
-            return f"/Volumes/{info.catalog_name}/{info.name}"
-    elif isinstance(info, VolumeInfo):
-        if detail:
-            return {
-                "name": f"/Volumes/{info.catalog_name}/{info.schema_name}/{info.name}",
-                "size": None,
-                "type": "directory",
-                "created": to_datetime(info.created_at),
-                "modified": to_datetime(info.updated_at),
-                "islink": False,
-            }
-        else:
-            return f"/Volumes/{info.catalog_name}/{info.schema_name}/{info.name}"
+            name = f"/Volumes/{info.catalog_name}/{info.schema_name}/{info.name}"
+        mapped = {
+            "name": name,
+            "size": None,
+            "type": "directory",
+            "created": to_datetime(info.created_at),
+            "modified": to_datetime(info.updated_at),
+            "islink": False,
+        }
     elif isinstance(info, GetMetadataResponse):
         if path is None:
             raise ValueError("Path must be provided for GetMetadataResponse mapping.")
-        if detail:
-            return {
-                "name": path.rstrip("/"),
-                "size": info.content_length,
-                "type": "file",
-                "created": None,
-                "modified": to_datetime(info.last_modified),
-                "islink": False,
-                "content-type": info.content_type,
-            }
-        else:
-            return path.rstrip("/")
+        mapped = {
+            "name": path.rstrip("/"),
+            "size": info.content_length,
+            "type": "file",
+            "created": None,
+            "modified": to_datetime(info.last_modified),
+            "islink": False,
+            "content-type": info.content_type,
+        }
+    elif isinstance(info, DirectoryEntry):
+        mapped = {
+            "name": info.path.rstrip("/"),
+            "size": info.file_size,
+            "type": "directory" if info.is_directory else "file",
+            "created": None,
+            "modified": to_datetime(info.last_modified),
+            "islink": False,
+        }
     else:
         raise ValueError(f"Unsupported info type: {info} {type(info)}")
 
+    return mapped if detail else mapped["name"]
+
 
 class VolumeFileSystem(DBFS):
     """Unity Catalog Volume file system.
@@ -149,6 +144,9 @@ class VolumeFileSystem(DBFS):
     _storage_proxy_probe_timeout: float = 3.0
     """The timeout in seconds for the storage proxy probe request."""
 
+    _max_delete_concurrency: int = 16
+    """The maximum number of concurrent delete requests of a recursive `rm`."""
+
     log = logging.getLogger(__name__)
 
     def __init__(
@@ -202,21 +200,6 @@ class VolumeFileSystem(DBFS):
         self.__lock = RLock()
         self.__session: ClientSession | None = None
 
-    def __getstate__(self):
-        state = super().__getstate__()
-        del state["_VolumeFileSystem__loop"]
-        del state["_VolumeFileSystem__io_thread"]
-        del state["_VolumeFileSystem__lock"]
-        del state["_VolumeFileSystem__session"]
-        return state
-
-    def __setstate__(self, state):
-        super().__setstate__(state)
-        self.__loop = None
-        self.__io_thread = None
-        self.__lock = RLock()
-        self.__session = None
-
     @property
     def _loop(self):
         if self.closed:
@@ -340,12 +323,18 @@ class VolumeFileSystem(DBFS):
                         f"/DatabricksInternal/Probe/fullstack/wis?ew={workspace_id}"
                     )
                     auth = _workspace_authenticator(self.config)
-                    async with session.head(
-                        probe_url,
-                        middlewares=(auth,),
-                        timeout=ClientTimeout(total=self._storage_proxy_probe_timeout),
-                    ) as response:
-                        return response.status == 200
+                    try:
+                        async with session.head(
+                            probe_url,
+                            middlewares=(auth,),
+                            timeout=ClientTimeout(
+                                total=self._storage_proxy_probe_timeout
+                            ),
+                        ) as response:
+                            return response.status == 200
+                    except (TimeoutError, ClientConnectionError):
+                        # The proxy can be unresolvable or unreachable from this network.
+                        return False
 
                 available = asyncio.run_coroutine_threadsafe(
                     probe(self._session, self.workspace_id), self._loop
@@ -472,7 +461,7 @@ class VolumeFileSystem(DBFS):
             return [volumes_root_info] if detail else [volumes_root]
 
         try:
-            catalog, schema, volume, _, _ = parse_volume_path(path)
+            catalog, schema, volume, _, full_path = parse_volume_path(path)
         except ValueError:
             raise file_not_found_error(path) from None
 
@@ -497,23 +486,146 @@ class VolumeFileSystem(DBFS):
                     )
                 ]
             else:
-                return super().ls(path, detail=detail)
+                try:
+                    entries = list(self.client.files.list_directory_contents(full_path))
+                except (NotFound, BadRequest):
+                    # A file lists as itself
+                    info = self._file_info(full_path)
+                    return [info if detail else info["name"]]
+                return [_map_info(e, detail=detail) for e in entries]
+
+    ### File/directory operations
+    def cp_file(self, path1, path2, **kwargs):
+        # Only logs: the base class has no precondition for copying
+        AbstractDatabricksFileSystem.cp_file(self, path1=path1, path2=path2, **kwargs)
 
-    def rmdir(self, path):
+        try:
+            src = self.open(path1, "rb")
+        except IsADirectoryError:
+            self._cp_directory(path2)
+            return
+
+        # The Files API has no copy: stream the content with concurrent reads and uploads
+        with src, self.open(path2, "wb") as dst:
+            shutil.copyfileobj(src, dst, length=self.max_write_block_size)
+
+    def mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs):
+        # fsspec's copy and delete, with the methods of this class
+        AbstractDatabricksFileSystem.mv(
+            self, path1, path2, recursive=recursive, maxdepth=maxdepth, **kwargs
+        )
+
+    def _rm(self, path):
         # The precondition check is implemented in the base class.
-        AbstractDatabricksFileSystem.rmdir(self, path=path)
+        AbstractDatabricksFileSystem._rm(self, path=path)
 
         _, _, _, _, full_path = parse_volume_path(path)
+        self._delete_file(full_path)
 
-        with error_mapping(path):
+    def rm(self, path, recursive=False, maxdepth=None):
+        self.log.debug(
+            "Deleting path: path=%s, recursive=%s, maxdepth=%s",
+            path,
+            recursive,
+            maxdepth,
+        )
+
+        paths = [path] if isinstance(path, (str, PathLike)) else path
+        for p in paths:
+            if glob.has_magic(self._strip_protocol(p)):
+                for matched in self.expand_path(p):
+                    self._rm_path(matched, recursive, maxdepth)
+            else:
+                self._rm_path(p, recursive, maxdepth)
+
+    def _rm_path(self, path, recursive: bool, maxdepth: int | None) -> None:
+        """Delete a file, or a directory with its content if recursive."""
+        _, _, _, path_in_volume, full_path = parse_volume_path(path)
+        if path_in_volume is None or path_in_volume == "/":
+            # Volumes, schemas and catalogs. A list, as DBFS.rm calls self.rm for a string.
+            DBFS.rm(self, [path], recursive=recursive, maxdepth=maxdepth)
+            return
+
+        if not recursive:
+            self._rm(path)
+            return
+
+        if self.info(path)["type"] == "file":
+            self._delete_file(full_path)
+            return
+
+        files, directories = self._walk_tree(full_path, maxdepth)
+        self._run_concurrently(self._delete_file, files)
+        # Children before parents. Directories of the same depth are independent.
+        directories.append(full_path)
+        for depth in sorted({d.count("/") for d in directories}, reverse=True):
+            self._run_concurrently(
+                self._delete_directory,
+                [d for d in directories if d.count("/") == depth],
+            )
+
+    def _walk_tree(
+        self, full_path: str, maxdepth: int | None
+    ) -> tuple[list[str], list[str]]:
+        """List the files and directories under a directory, listing each depth concurrently."""
+        files: list[str] = []
+        directories: list[str] = []
+
+        def list_directory(directory: str) -> list[DirectoryEntry]:
+            with error_mapping(directory):
+                return list(self.client.files.list_directory_contents(directory))
+
+        level = [full_path]
+        depth = 0
+        while level and (maxdepth is None or depth < maxdepth):
+            depth += 1
+            next_level = []
+            with ThreadPoolExecutor(
+                max_workers=min(len(level), self._max_delete_concurrency)
+            ) as executor:
+                for entries in executor.map(list_directory, level):
+                    for e in entries:
+                        if e.is_directory:
+                            next_level.append(e.path.rstrip("/"))
+                        else:
+                            files.append(e.path)
+            directories.extend(next_level)
+            level = next_level
+
+        return files, directories
+
+    def _run_concurrently(self, func: Callable[[str], None], paths: Iterable[str]):
+        """Call the function on each path in a thread pool, and raise the first error."""
+        paths = list(paths)
+        if not paths:
+            return
+        with ThreadPoolExecutor(
+            max_workers=min(len(paths), self._max_delete_concurrency)
+        ) as executor:
+            for _ in executor.map(func, paths):
+                pass
+
+    def _delete_file(self, full_path: str) -> None:
+        with error_mapping(full_path):
+            self.client.files.delete(full_path)
+
+    def _delete_directory(self, full_path: str) -> None:
+        with error_mapping(full_path):
             try:
                 self.client.files.delete_directory(full_path)
             except BadRequest as e:
                 if "is not empty" in e.args[0]:
-                    raise not_empty_error(path) from e
+                    raise not_empty_error(full_path) from e
                 else:
                     raise
 
+    def rmdir(self, path):
+        # The precondition check is implemented in the base class.
+        AbstractDatabricksFileSystem.rmdir(self, path=path)
+
+        _, _, _, _, full_path = parse_volume_path(path)
+        self._delete_directory(full_path)
+
     def _check_open_target(self, path: str, mode: str) -> dict[str, Any] | None:
         """Check that the path can be opened in the mode, and return its info if needed."""
         if "r" not in mode and "x" not in mode:
@@ -669,6 +781,9 @@ def _workspace_authenticator(config: Config) -> ClientMiddlewareType:
     return auth
 
 
+_retryable_statuses = frozenset({429, 500, 502, 503, 504})
+
+
 class VolumeReadableFile(AbstractAsyncReadableFile):
     log = VolumeFileSystem.log
 
@@ -704,6 +819,12 @@ class VolumeReadableFile(AbstractAsyncReadableFile):
     presigned_url_refresh_after: int = 30  # min
     """The duration in minutes after which the presigned URL should be refreshed."""
 
+    fetch_max_retries: int = 3
+    """The maximal number of retries of a range request after a transient error."""
+
+    fetch_retry_delay: float = 1.0
+    """The base delay in seconds before the first retry of a range request. It doubles with each retry."""
+
     def __init__(
         self,
         path: str,
@@ -803,6 +924,26 @@ class VolumeReadableFile(AbstractAsyncReadableFile):
         return self._presigned_url, self._presigned_url_headers
 
     async def _fetch_range(self, start: int, end: int) -> bytes:
+        delay = self.fetch_retry_delay
+        for _ in range(self.fetch_max_retries):
+            try:
+                return await self._request_range(start, end)
+            except ClientResponseError as e:
+                if e.status not in _retryable_statuses:
+                    raise
+                self.log.warning(
+                    "Retrying range request after status %d: path=%s, range=%d-%d",
+                    e.status,
+                    self.path,
+                    start,
+                    end,
+                )
+            # The jitter keeps concurrent range requests from retrying together.
+            await asyncio.sleep(delay * (1 + random()))
+            delay *= 2
+        return await self._request_range(start, end)
+
+    async def _request_range(self, start: int, end: int) -> bytes:
         if self.use_presigned_url:
             path, headers = await self._get_presigned_url()
             middlewares = ()

{fsspec_databricks-0.2.1 → fsspec_databricks-0.3.0}/tests/conftest.py

@@ -109,6 +109,7 @@ def dummy_api_context(dummy_api):
         context.files.clear()
         context.multipart_upload_sessions.clear()
         context.resumable_upload_sessions.clear()
+        context.get_file_errors.clear()
 
         context.check_signed_url = False
         context.upload_mode = "multipart"

{fsspec_databricks-0.2.1 → fsspec_databricks-0.3.0}/tests/dummy_api.py

@@ -33,6 +33,9 @@ class DummyAPIContext:
     check_signed_url: bool = False
     abort_supported: bool = True
 
+    get_file_errors: list[int] = field(default_factory=list)
+    """Statuses that the next requests to read a file get, one each."""
+
     files: dict[str, bytes] = field(default_factory=dict)
 
     multipart_upload_sessions: dict[str, dict[tuple[int, str], bytes | None]] = field(
@@ -118,6 +121,13 @@ async def get_file(
     range: Annotated[str | None, Header()] = None,
     signed: Annotated[str | None, Query()] = None,
 ):
+    if context.get_file_errors:
+        raise http_exception(
+            "INTERNAL_ERROR",
+            "Injected error",
+            context.get_file_errors.pop(0),
+        )
+
     if context.check_signed_url and signed != "true":
         raise http_exception(
             "BAD_REQUEST",

{fsspec_databricks-0.2.1 → fsspec_databricks-0.3.0}/tests/test_volume.py

@@ -1,6 +1,7 @@
+import socket
 from io import BytesIO
 from random import randbytes
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pandas as pd
 import pyarrow as pa
@@ -12,7 +13,11 @@ from aiohttp import (
     ClientResponse,
 )
 
-from fsspec_databricks.volume import VolumeReadableFile, VolumeWritableFile
+from fsspec_databricks.volume import (
+    VolumeFileSystem,
+    VolumeReadableFile,
+    VolumeWritableFile,
+)
 
 from .utils import bytes_sig
 
@@ -23,6 +28,39 @@ async def _dummy_auth(
     return await handler(request)
 
 
+def test_volume_fs_storage_proxy_probe_timeout():
+    client = MagicMock()
+    client.config.authenticate.return_value = {}
+    client.get_workspace_id.return_value = 1234
+
+    # Connections are accepted by the backlog, but no response ever comes.
+    with socket.create_server(("127.0.0.1", 0)) as server:
+        port = server.getsockname()[1]
+
+        with VolumeFileSystem(client=client) as fs:
+            fs._storage_proxy_hostname = f"http://127.0.0.1:{port}"
+            fs._storage_proxy_probe_timeout = 0.1
+
+            assert fs.storage_proxy_available is False
+
+
+def test_volume_fs_storage_proxy_probe_connection_refused():
+    client = MagicMock()
+    client.config.authenticate.return_value = {}
+    client.get_workspace_id.return_value = 1234
+
+    # Nothing listens on the port once the server is closed.
+    with socket.create_server(("127.0.0.1", 0)) as server:
+        port = server.getsockname()[1]
+
+    with VolumeFileSystem(client=client) as fs:
+        fs._storage_proxy_hostname = f"http://127.0.0.1:{port}"
+        # Long enough for a refused connection on Windows, which retries for seconds.
+        fs._storage_proxy_probe_timeout = 30.0
+
+        assert fs.storage_proxy_available is False
+
+
 def test_volume_file_read(
     dummy_api, dummy_api_context, client_event_loop, aiohttp_session
 ):
@@ -58,6 +96,83 @@ def test_volume_file_read(
             _ = f.read()
 
 
+@pytest.mark.parametrize("status", [429, 500, 502, 503, 504])
+def test_volume_file_read_retries(
+    dummy_api_context, dummy_api, client_event_loop, aiohttp_session, status
+):
+    data = randbytes(1000)
+    path = "/Volumes/catalog_a/schema/b/volume_c/path/to/file"
+
+    dummy_api_context.files[path] = data
+    dummy_api_context.get_file_errors.extend([status, status])
+
+    with patch.object(VolumeReadableFile, "fetch_retry_delay", 0):
+        with VolumeReadableFile(
+            path=path,
+            session=aiohttp_session,
+            base_url=dummy_api,
+            auth=_dummy_auth,
+            loop=client_event_loop,
+            size=len(data),
+        ) as f:
+            assert f.read() == data
+
+    assert not dummy_api_context.get_file_errors
+
+
+def test_volume_file_read_retries_exhausted(
+    dummy_api_context, dummy_api, client_event_loop, aiohttp_session
+):
+    data = randbytes(1000)
+    path = "/Volumes/catalog_a/schema/b/volume_c/path/to/file"
+
+    dummy_api_context.files[path] = data
+    dummy_api_context.get_file_errors.extend([500] * 10)
+
+    with patch.object(VolumeReadableFile, "fetch_retry_delay", 0):
+        with VolumeReadableFile(
+            path=path,
+            session=aiohttp_session,
+            base_url=dummy_api,
+            auth=_dummy_auth,
+            loop=client_event_loop,
+            size=len(data),
+            max_concurrency=1,
+        ) as f:
+            with pytest.raises(OSError):
+                f.read()
+
+    # The first request and its retries.
+    attempts = 10 - len(dummy_api_context.get_file_errors)
+    assert attempts == VolumeReadableFile.fetch_max_retries + 1
+
+
+@pytest.mark.parametrize("status", [400, 403, 404])
+def test_volume_file_read_no_retry(
+    dummy_api_context, dummy_api, client_event_loop, aiohttp_session, status
+):
+    data = randbytes(1000)
+    path = "/Volumes/catalog_a/schema/b/volume_c/path/to/file"
+
+    dummy_api_context.files[path] = data
+    dummy_api_context.get_file_errors.extend([status] * 2)
+
+    with patch.object(VolumeReadableFile, "fetch_retry_delay", 0):
+        with VolumeReadableFile(
+            path=path,
+            session=aiohttp_session,
+            base_url=dummy_api,
+            auth=_dummy_auth,
+            loop=client_event_loop,
+            size=len(data),
+            max_concurrency=1,
+        ) as f:
+            with pytest.raises(OSError):
+                f.read()
+
+    assert len(dummy_api_context.get_file_errors) == 1
+
+
 def test_volume_file_read_parquet(
     client_event_loop,
     dummy_api,

{fsspec_databricks-0.2.1 → fsspec_databricks-0.3.0}/tests/test_volume_fs.py

@@ -11,8 +11,8 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 from databricks.sdk import WorkspaceClient
-from databricks.sdk.errors import NotFound, ResourceDoesNotExist
-from databricks.sdk.service.files import GetMetadataResponse
+from databricks.sdk.errors import BadRequest, NotFound, ResourceDoesNotExist
+from databricks.sdk.service.files import DirectoryEntry, GetMetadataResponse
 
 from fsspec_databricks import DatabricksFileSystem
 from fsspec_databricks.path import parse_volume_path
@@ -152,6 +152,24 @@ def test_volume_fs_serialization():
     assert fs._loop is not fs2._loop
 
 
+def test_volume_fs_pickled_as_arguments():
+    fs = VolumeFileSystem(
+        config_params={"host": "https://example.cloud.databricks.com", "token": "x"},
+        max_read_concurrency=3,
+    )
+    loop = fs._loop
+
+    fs2 = pickle.loads(pickle.dumps(fs))
+    try:
+        # fsspec rebuilds the file system from its constructor arguments
+        assert fs2.storage_options == fs.storage_options
+        assert fs2.max_read_concurrency == 3
+        assert fs2._loop is not loop
+    finally:
+        fs.close()
+        fs2.close()
+
+
 def test_volume_fs_session():
     with VolumeFileSystem() as fs:
         # Created from a thread without a running event loop
@@ -184,9 +202,9 @@ def test_volume_fs_session_in_event_loop():
 
 
 class _CountingFilesAPI:
-    """A fake of `WorkspaceClient.files` for metadata, recording the calls."""
+    """A fake of `WorkspaceClient.files` for metadata, listing and deletion, recording the calls."""
 
-    def __init__(self, files: list[str], directories: list[str]):
+    def __init__(self, files: dict[str, bytes], directories: list[str]):
         self.files = files
         self.directories = directories
         self.calls: list[tuple[str, str]] = []
@@ -196,7 +214,7 @@ class _CountingFilesAPI:
         if path not in self.files:
             raise NotFound(f"Not found: {path}")
         return GetMetadataResponse(
-            content_length=123,
+            content_length=len(self.files[path]),
             content_type="application/octet-stream",
             last_modified="Wed, 21 Oct 2026 07:28:00 GMT",
         )
@@ -206,6 +224,37 @@ class _CountingFilesAPI:
         if path not in self.directories:
             raise NotFound(f"Not found: {path}")
 
+    def _children(self, path):
+        return [
+            p for p in [*self.files, *self.directories] if p.rsplit("/", 1)[0] == path
+        ]
+
+    def list_directory_contents(self, path):
+        self.calls.append(("list", path))
+        if path not in self.directories:
+            raise NotFound(f"Not found: {path}")
+        for p in self._children(path):
+            yield DirectoryEntry(
+                path=p,
+                name=p.rsplit("/", 1)[1],
+                is_directory=p in self.directories,
+                file_size=len(self.files[p]) if p in self.files else None,
+            )
+
+    def delete(self, path):
+        self.calls.append(("delete", path))
+        if path not in self.files:
+            raise NotFound(f"Not found: {path}")
+        del self.files[path]
+
+    def delete_directory(self, path):
+        self.calls.append(("delete_directory", path))
+        if path not in self.directories:
+            raise NotFound(f"Not found: {path}")
+        if self._children(path):
+            raise BadRequest(f"Directory {path} is not empty")
+        self.directories.remove(path)
+
 
 def test_volume_fs_metadata_calls():
     """Each check of a file or a directory makes as few API calls as possible."""
@@ -213,7 +262,7 @@ def test_volume_fs_metadata_calls():
     directory = "/Volumes/c/s/v/dir"
     missing = "/Volumes/c/s/v/dir/missing"
 
-    files_api = _CountingFilesAPI(files=[file], directories=[directory])
+    files_api = _CountingFilesAPI(files={file: b"data"}, directories=[directory])
     client = MagicMock()
     client.files = files_api
 
@@ -254,6 +303,134 @@ def test_volume_fs_metadata_calls():
         assert calls_of(lambda: fs.isfile("/Volumes/c/s/v")) == []
         assert not fs.isfile("/Volumes/c/s/v")
 
+        # Listing a directory: one call. A file lists as itself.
+        assert calls_of(lambda: fs.ls(directory)) == [("list", directory)]
+        assert fs.ls(directory, detail=False) == [file]
+        assert fs.ls(directory)[0]["size"] == len(b"data")
+        assert calls_of(lambda: fs.ls(file)) == [("list", file), ("file", file)]
+        assert fs.ls(file, detail=False) == [file]
+        with pytest.raises(FileNotFoundError):
+            fs.ls(missing)
+
+
+def _volume_fs_with_fake_files(
+    files: dict[str, bytes], directories: list[str]
+) -> tuple[VolumeFileSystem, _CountingFilesAPI, MagicMock]:
+    """A Volume file system on `_CountingFilesAPI`, whose files are opened as `BytesIO`."""
+    files_api = _CountingFilesAPI(files=files, directories=directories)
+    client = MagicMock()
+    client.files = files_api
+
+    class _Written(BytesIO):
+        def __init__(self, path):
+            super().__init__()
+            self.path = path
+
+        def close(self):
+            files_api.files[self.path] = self.getvalue()
+            super().close()
+
+    def fake_open(fs, path, mode="rb", **kwargs):
+        fs._check_open_target(path, mode)
+        return BytesIO(files_api.files[path]) if "r" in mode else _Written(path)
+
+    patcher = patch.object(VolumeFileSystem, "_open", fake_open)
+    patcher.start()
+    fs = VolumeFileSystem(client=client)
+    fs._fake_patcher = patcher
+    return fs, files_api, client
+
+
+@pytest.fixture
+def fake_volume_fs():
+    created = []
+
+    def create(files: dict[str, bytes], directories: list[str]):
+        fs, files_api, client = _volume_fs_with_fake_files(files, directories)
+        created.append(fs)
+        return fs, files_api, client
+
+    yield create
+
+    for fs in created:
+        fs._fake_patcher.stop()
+        fs.close()
+
+
+def test_volume_fs_cp_file_streams(fake_volume_fs):
+    """Copying a file streams it with the file objects, not the SDK's `dbfs.copy`."""
+    src = "/Volumes/c/s/v/dir/src.bin"
+    dst = "/Volumes/c/s/v/dir/dst.bin"
+    data = randbytes(12345)
+
+    fs, files_api, client = fake_volume_fs({src: data}, ["/Volumes/c/s/v/dir"])
+    fs.cp_file(src, dst)
+
+    assert files_api.files[dst] == data
+    # One metadata call each: the source is a file, the destination is no directory
+    assert files_api.calls == [("file", src), ("directory", dst)]
+    assert client.dbfs.mock_calls == []
+
+    with pytest.raises(FileNotFoundError):
+        fs.cp_file("/Volumes/c/s/v/dir/missing", dst)
+
+
+def test_volume_fs_mv_file(fake_volume_fs):
+    src = "/Volumes/c/s/v/dir/src.bin"
+    dst = "/Volumes/c/s/v/dir/dst.bin"
+    data = randbytes(12345)
+
+    fs, files_api, client = fake_volume_fs({src: data}, ["/Volumes/c/s/v/dir"])
+    fs.mv(src, dst)
+
+    assert files_api.files == {dst: data}
+    assert client.dbfs.mock_calls == []
+
+
+def test_volume_fs_rm_calls(fake_volume_fs):
+    """Deleting calls the Files API directly, children before their parents."""
+    root = "/Volumes/c/s/v/dir"
+    files = {
+        f"{root}/a.bin": b"a",
+        f"{root}/sub/b.bin": b"b",
+        f"{root}/sub/c.bin": b"c",
+    }
+    directories = [root, f"{root}/sub", f"{root}/sub/empty", f"{root}/other"]
+
+    fs, files_api, client = fake_volume_fs(dict(files), list(directories))
+
+    # A file: one metadata call and the delete
+    fs.rm(f"{root}/a.bin")
+    assert files_api.calls == [("file", f"{root}/a.bin"), ("delete", f"{root}/a.bin")]
+
+    with pytest.raises(IsADirectoryError):
+        fs.rm(f"{root}/sub")
+
+    # A directory: one listing per directory, then the deletes
+    files_api.calls.clear()
+    fs.rm(root, recursive=True)
+
+    assert files_api.files == {}
+    assert files_api.directories == []
+    assert client.dbfs.mock_calls == []
+
+    kinds = [kind for kind, _ in files_api.calls]
+    assert kinds.count("list") == 4
+    assert sorted(p for kind, p in files_api.calls if kind == "delete") == [
+        f"{root}/sub/b.bin",
+        f"{root}/sub/c.bin",
+    ]
+    deleted_directories = [
+        p for kind, p in files_api.calls if kind == "delete_directory"
+    ]
+    assert deleted_directories.index(f"{root}/sub/empty") < deleted_directories.index(
+        f"{root}/sub"
+    )
+    assert deleted_directories[-1] == root
+
+    with pytest.raises(FileNotFoundError):
+        fs.rm(root, recursive=True)
+
 
 def test_volume_fs_workspace_id(client: WorkspaceClient):
     with VolumeFileSystem(client=client) as fs:
@@ -418,6 +595,9 @@ def test_volume_fs_ls(
         stripped_dbfs_url(f"{test_dir}/2_directory"),
     ]
 
+    # A file lists as itself
+    assert fs.ls(dbfs_url(data_path), detail=False) == [stripped_dbfs_url(data_path)]
+
     with pytest.raises(FileNotFoundError):
         fs.ls(dbfs_url(f"{test_dir}/nonexistent"))
 
@@ -545,6 +725,17 @@ def test_volume_fs_cp_file(
 
     assert data == copied_data
 
+    # Test copying a file larger than one block, with a multipart upload
+    large_data = randbytes(12 * 1024 * 1024)
+    client.files.upload(f"{test_dir}/3_large.bin", BytesIO(large_data))
+    fs.cp_file(
+        dbfs_url(f"{test_dir}/3_large.bin"),
+        dbfs_url(f"{test_dir}/3_large_copied.bin"),
+    )
+    assert bytes_sig(download(client, f"{test_dir}/3_large_copied.bin")) == bytes_sig(
+        large_data
+    )
+
     # Test copying a directory
     fs.cp_file(
         dbfs_url(f"{test_dir}/2_directory"),
@@ -661,7 +852,12 @@ def test_volume_fs_rm(
             "1_data.bin": data,
             "2_directory": {
                 "3_data.bin": data,
+                "4_directory": {
+                    "5_data.bin": data,
+                    "6_empty_directory": {},
+                },
             },
+            "7_data.bin": data,
         },
     )
 
@@ -676,10 +872,59 @@ def test_volume_fs_rm(
     with pytest.raises(ResourceDoesNotExist):
         client.workspace.get_status(f"{test_dir}/2_directory")
 
+    # A file, recursively
+    fs.rm(dbfs_url(f"{test_dir}/7_data.bin"), recursive=True)
+    with pytest.raises(NotFound):
+        client.files.get_metadata(f"{test_dir}/7_data.bin")
+
     with pytest.raises(FileNotFoundError):
         fs.rm(dbfs_url(f"{test_dir}/nonexistent"))
 
 
+@pytest.mark.parametrize(
+    "fs_class",
+    [
+        VolumeFileSystem,
+        DatabricksFileSystem,
+    ],
+)
+def test_volume_fs_mv(
+    client: WorkspaceClient,
+    fs_class: type[VolumeFileSystem | DatabricksFileSystem],
+    volume_test_root: str,
+):
+    fs = fs_class(client=client, verbose_debug_log=True)
+
+    data1 = randbytes(12345)
+    data2 = randbytes(12345)
+
+    test_dir = init_test_dir(
+        client,
+        volume_test_root,
+        fs_class,
+        items={
+            "1_data.bin": data1,
+            "2_directory": {
+                "3_data.bin": data2,
+            },
+        },
+    )
+
+    fs.mv(dbfs_url(f"{test_dir}/1_data.bin"), dbfs_url(f"{test_dir}/1_moved.bin"))
+    assert download(client, f"{test_dir}/1_moved.bin") == data1
+    with pytest.raises(NotFound):
+        client.files.get_metadata(f"{test_dir}/1_data.bin")
+
+    fs.mv(
+        dbfs_url(f"{test_dir}/2_directory"),
+        dbfs_url(f"{test_dir}/2_moved"),
+        recursive=True,
+    )
+    assert download(client, f"{test_dir}/2_moved/3_data.bin") == data2
+    with pytest.raises(NotFound):
+        client.files.get_directory_metadata(f"{test_dir}/2_directory")
+
+
 @pytest.mark.parametrize(
     "fs_class",
     [