feature(logging): include keys in logging store logs

Author: jhamman

src/zarr/storage/logging.py

@@ -5,7 +5,7 @@
 import time
 from collections import defaultdict
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Self
+from typing import TYPE_CHECKING, Any, Self
 
 from zarr.abc.store import AccessMode, ByteRangeRequest, Store
 from zarr.core.buffer import Buffer
@@ -57,17 +57,24 @@ def _default_handler(self) -> logging.Handler:
         return handler
 
     @contextmanager
-    def log(self) -> Generator[None, None, None]:
+    def log(self, hint: Any = "") -> Generator[None, None, None]:
+        """Context manager to log method calls
+
+        Each call to the wrapped store is logged to the configured logger and added to
+        the counter dict.
+        """
         method = inspect.stack()[2].function
         op = f"{type(self._store).__name__}.{method}"
+        if hint:
+            op += f"({hint})"
         self.logger.info(f"Calling {op}")
         start_time = time.time()
         try:
             self.counter[method] += 1
             yield
         finally:
             end_time = time.time()
-            self.logger.info(f"Finished {op} in {end_time - start_time:.2f} seconds")
+            self.logger.info(f"Finished {op} [{end_time - start_time:.2f}s]")
 
     @property
     def supports_writes(self) -> bool:
@@ -95,10 +102,15 @@ def _mode(self) -> AccessMode:  # type: ignore[override]
             return self._store._mode
 
     @property
-    def _is_open(self) -> bool:  # type: ignore[override]
+    def _is_open(self) -> bool:
         with self.log():
             return self._store._is_open
 
+    @_is_open.setter
+    def _is_open(self, value: bool) -> None:
+        with self.log(value):
+            self._store._is_open = value
+
     async def _open(self) -> None:
         with self.log():
             return await self._store._open()
@@ -122,7 +134,7 @@ def __repr__(self) -> str:
         return f"LoggingStore({repr(self._store)!r})"
 
     def __eq__(self, other: object) -> bool:
-        with self.log():
+        with self.log(other):
             return self._store == other
 
     async def get(
@@ -131,56 +143,69 @@ async def get(
         prototype: BufferPrototype,
         byte_range: tuple[int | None, int | None] | None = None,
     ) -> Buffer | None:
-        with self.log():
+        # docstring inherited
+        with self.log(key):
             return await self._store.get(key=key, prototype=prototype, byte_range=byte_range)
 
     async def get_partial_values(
         self,
         prototype: BufferPrototype,
         key_ranges: Iterable[tuple[str, ByteRangeRequest]],
     ) -> list[Buffer | None]:
-        with self.log():
+        # docstring inherited
+        keys = ",".join([k[0] for k in key_ranges])
+        with self.log(keys):
             return await self._store.get_partial_values(prototype=prototype, key_ranges=key_ranges)
 
     async def exists(self, key: str) -> bool:
-        with self.log():
+        # docstring inherited
+        with self.log(key):
             return await self._store.exists(key)
 
     async def set(self, key: str, value: Buffer) -> None:
-        with self.log():
+        # docstring inherited
+        with self.log(key):
             return await self._store.set(key=key, value=value)
 
     async def set_if_not_exists(self, key: str, value: Buffer) -> None:
-        with self.log():
+        # docstring inherited
+        with self.log(key):
             return await self._store.set_if_not_exists(key=key, value=value)
 
     async def delete(self, key: str) -> None:
-        with self.log():
+        # docstring inherited
+        with self.log(key):
             return await self._store.delete(key=key)
 
     async def set_partial_values(
         self, key_start_values: Iterable[tuple[str, int, bytes | bytearray | memoryview]]
     ) -> None:
-        with self.log():
+        # docstring inherited
+        keys = ",".join([k[0] for k in key_start_values])
+        with self.log(keys):
             return await self._store.set_partial_values(key_start_values=key_start_values)
 
     async def list(self) -> AsyncGenerator[str, None]:
+        # docstring inherited
         with self.log():
             async for key in self._store.list():
                 yield key
 
     async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
-        with self.log():
+        # docstring inherited
+        with self.log(prefix):
             async for key in self._store.list_prefix(prefix=prefix):
                 yield key
 
     async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
-        with self.log():
+        # docstring inherited
+        with self.log(prefix):
             async for key in self._store.list_dir(prefix=prefix):
                 yield key
 
     def with_mode(self, mode: AccessModeLiteral) -> Self:
-        with self.log():
+        # docstring inherited
+        with self.log(mode):
             return type(self)(
                 self._store.with_mode(mode),
                 log_level=self.log_level,

src/zarr/storage/remote.py

@@ -82,7 +82,35 @@ def from_url(
         mode: AccessModeLiteral = "r",
         allowed_exceptions: tuple[type[Exception], ...] = ALLOWED_EXCEPTIONS,
     ) -> RemoteStore:
-        fs, path = fsspec.url_to_fs(url, **storage_options)
+        """
+        Create a RemoteStore from a URL.
+
+        Parameters
+        ----------
+        url : str
+            The URL to the root of the store.
+        storage_options : dict, optional
+            The options to pass to fsspec when creating the filesystem.
+        mode : str, optional
+            The mode of the store. Defaults to "r".
+        allowed_exceptions : tuple, optional
+            The exceptions that are allowed to be raised when accessing the
+            store. Defaults to ALLOWED_EXCEPTIONS.
+
+        Returns
+        -------
+        RemoteStore
+        """
+        opts = storage_options or {}
+        opts = {"asynchronous": True, **opts}
+
+        fs, path = fsspec.url_to_fs(url, **opts)
+
+        # fsspec is not consistent about removing the scheme from the path, so check and strip it here
+        # https://github.com/fsspec/filesystem_spec/issues/1722
+        if "://" in path:
+            _, path = path.split("://", maxsplit=1)
+
         return cls(fs=fs, path=path, mode=mode, allowed_exceptions=allowed_exceptions)
 
     async def clear(self) -> None: