Merge pull request #2355 from jhamman/fix/logging_store_improvements

Fix/logging store improvements

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
@@ -75,22 +75,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:
@@ -118,10 +120,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()
@@ -147,7 +154,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(
@@ -157,7 +164,7 @@ async def get(
         byte_range: tuple[int | None, int | None] | None = None,
     ) -> Buffer | None:
         # docstring inherited
-        with self.log():
+        with self.log(key):
             return await self._store.get(key=key, prototype=prototype, byte_range=byte_range)
 
     async def get_partial_values(
@@ -166,34 +173,36 @@ async def get_partial_values(
         key_ranges: Iterable[tuple[str, ByteRangeRequest]],
     ) -> list[Buffer | None]:
         # docstring inherited
-        with self.log():
+        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:
         # docstring inherited
-        with self.log():
+        with self.log(key):
             return await self._store.exists(key)
 
     async def set(self, key: str, value: Buffer) -> None:
         # docstring inherited
-        with self.log():
+        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:
         # docstring inherited
-        with self.log():
+        with self.log(key):
             return await self._store.set_if_not_exists(key=key, value=value)
 
     async def delete(self, key: str) -> None:
         # docstring inherited
-        with self.log():
+        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:
         # docstring inherited
-        with self.log():
+        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]:
@@ -204,19 +213,19 @@ async def list(self) -> AsyncGenerator[str, None]:
 
     async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         # docstring inherited
-        with self.log():
+        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]:
         # docstring inherited
-        with self.log():
+        with self.log(prefix):
             async for key in self._store.list_dir(prefix=prefix):
                 yield key
 
     def with_mode(self, mode: AccessModeLiteral) -> Self:
         # docstring inherited
-        with self.log():
+        with self.log(mode):
             return type(self)(
                 self._store.with_mode(mode),
                 log_level=self.log_level,