feature(store): add LoggingStore wrapper

Author: jhamman

src/zarr/store/logging.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+import inspect
+import logging
+import time
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+from zarr.abc.store import Store
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Generator
+
+    from zarr.core.buffer import Buffer, BufferPrototype
+
+
+class LoggingStore(Store):
+    _store: Store
+
+    def __init__(
+        self,
+        store: Store,
+        log_level: str = "DEBUG",
+        log_handler: logging.Handler | None = None,
+    ):
+        self._store = store
+
+        self._configure_logger(log_level, log_handler)
+
+    def _configure_logger(
+        self, log_level: str = "DEBUG", log_handler: logging.Handler | None = None
+    ) -> None:
+        self.log_level = log_level
+        self.logger = logging.getLogger(f"LoggingStore({self._store!s})")
+        self.logger.setLevel(log_level)
+
+        if not self.logger.hasHandlers():
+            if not log_handler:
+                log_handler = self._default_handler()
+            # Add handler to logger
+            self.logger.addHandler(log_handler)
+
+    def _default_handler(self) -> logging.Handler:
+        """Define a default log handler"""
+        handler = logging.StreamHandler()
+        handler.setLevel(self.log_level)
+        formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+        handler.setFormatter(formatter)
+        return handler
+
+    @contextmanager
+    def log(self) -> Generator[None, None, None]:
+        op = f"{type(self._store).__name__}.{inspect.stack()[2].function}"
+        self.logger.info(f"Calling {op}")
+        start_time = time.time()
+        try:
+            yield
+        finally:
+            end_time = time.time()
+            self.logger.info(f"Finished {op} in {end_time - start_time:.2f} seconds")
+
+    @property
+    def supports_writes(self) -> bool:
+        with self.log():
+            return self._store.supports_writes
+
+    @property
+    def supports_deletes(self) -> bool:
+        with self.log():
+            return self._store.supports_deletes
+
+    @property
+    def supports_partial_writes(self) -> bool:
+        with self.log():
+            return self._store.supports_partial_writes
+
+    @property
+    def supports_listing(self) -> bool:
+        with self.log():
+            return self._store.supports_listing
+
+    async def empty(self) -> bool:
+        with self.log():
+            return await self._store.empty()
+
+    async def clear(self) -> None:
+        with self.log():
+            return await self._store.clear()
+
+    def __str__(self) -> str:
+        return f"logging-{self._store!s}"
+
+    def __repr__(self) -> str:
+        return f"LoggingStore({repr(self._store)!r})"
+
+    def __eq__(self, other: object) -> bool:
+        with self.log():
+            return self._store == other
+
+    async def get(
+        self,
+        key: str,
+        prototype: BufferPrototype,
+        byte_range: tuple[int | None, int | None] | None = None,
+    ) -> Buffer | None:
+        with self.log():
+            return await self._store.get(key=key, prototype=prototype, byte_range=byte_range)
+
+    async def get_partial_values(
+        self,
+        prototype: BufferPrototype,
+        key_ranges: list[tuple[str, tuple[int | None, int | None]]],
+    ) -> list[Buffer | None]:
+        with self.log():
+            return await self._store.get_partial_values(prototype=prototype, key_ranges=key_ranges)
+
+    async def exists(self, key: str) -> bool:
+        with self.log():
+            return await self._store.exists(key)
+
+    async def set(self, key: str, value: Buffer) -> None:
+        with self.log():
+            return await self._store.set(key=key, value=value)
+
+    async def delete(self, key: str) -> None:
+        with self.log():
+            return await self._store.delete(key=key)
+
+    async def set_partial_values(self, key_start_values: list[tuple[str, int, bytes]]) -> None:
+        with self.log():
+            return await self._store.set_partial_values(key_start_values=key_start_values)
+
+    async def list(self) -> AsyncGenerator[str, None]:
+        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():
+            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():
+            async for key in self._store.list_dir(prefix=prefix):
+                yield key

tests/v3/test_store/test_logging.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+from zarr.core.buffer import default_buffer_prototype
+from zarr.store.logging import LoggingStore
+
+if TYPE_CHECKING:
+    from zarr.abc.store import Store
+
+
+@pytest.mark.parametrize("store", ("local", "memory", "zip"), indirect=["store"])
+async def test_logging_store(store: Store, caplog) -> None:
+    wrapped = LoggingStore(store=store, log_level="DEBUG")
+    buffer = default_buffer_prototype().buffer
+
+    caplog.clear()
+    res = await wrapped.set("foo/bar/c/0", buffer.from_bytes(b"\x01\x02\x03\x04"))
+    assert res is None
+    assert len(caplog.record_tuples) == 2
+    for tup in caplog.record_tuples:
+        assert str(store) in tup[0]
+    assert f"Calling {type(store).__name__}.set" in caplog.record_tuples[0][2]
+    assert f"Finished {type(store).__name__}.set" in caplog.record_tuples[1][2]
+
+    caplog.clear()
+    keys = [k async for k in wrapped.list()]
+    assert keys == ["foo/bar/c/0"]
+    assert len(caplog.record_tuples) == 2
+    for tup in caplog.record_tuples:
+        assert str(store) in tup[0]
+    assert f"Calling {type(store).__name__}.list" in caplog.record_tuples[0][2]
+    assert f"Finished {type(store).__name__}.list" in caplog.record_tuples[1][2]