diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index 61989470d..ed70f3f66 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -7,6 +7,36 @@ This project adheres to [Semantic Versioning](http://semver.org/). ### What's New +#### Faster `get_data` pagination + +Up to a ~80x speedup for some get_data calls. + +#### Shared on-disk cache (opt-out, on by default) + +`client.channels.get_data(...)` now caches the channel windows it returns to disk by default. Subsequent calls covering the same channel/time range — including from a fresh process — read straight out of the cache instead of going to the wire. This also bounds memory: nothing is held in process after the call returns. + +The cache lives on the `SiftClient` as a single shared store: every cache-aware resource writes to one global byte budget at one path, with one LRU policy. The default location is `/sift-data-cache`, capped at 4 GiB with LRU eviction. If the default path can't be opened (read-only filesystem, restricted container, etc.), the client logs a warning and continues with caching disabled — `get_data` still works, it just always goes to the wire. + +`ignore_cache=True` on `client.channels.get_data(...)` now skips writing into the cache as well as reading from it. Previously a "non-caching" workload still appended to the shared cache on every call. + +Configuration lives on the new `client.cache` namespace — knobs are global because the store is shared: + +```python +# Opt out — no data persisted to disk; every get_data call goes to the wire. +client.cache.disable() + +# Reconfigure the location or byte cap. +client.cache.enable(path="/data/sift-cache", max_bytes=2 * 1024 ** 3) + +# Remove a stale or corrupted cache directory. +client.cache.clear() # default tmp path +client.cache.clear("/data/sift-cache") # custom path +``` + +`enable` is also the way to turn the cache back on after a prior `disable` call. + +The cache is powered by [`diskcache`](https://grantjenks.com/docs/diskcache/) (pure-Python, SQLite-backed) with LRU eviction. + #### Report templates Added a full report template API, available as a nested resource of reports at `client.reports.templates` (and asynchronously via `client.async_.reports.templates`). Report templates can now be created, fetched, listed, updated, archived, and unarchived directly from `sift_client`: diff --git a/python/lib/sift_client/__init__.py b/python/lib/sift_client/__init__.py index e98ee783d..54fb62d14 100644 --- a/python/lib/sift_client/__init__.py +++ b/python/lib/sift_client/__init__.py @@ -139,8 +139,11 @@ async def main(): from sift_client.client import SiftClient from sift_client.config import config from sift_client.transport import SiftConnectionConfig +from sift_client.util.cache import CacheNamespace, CacheStats __all__ = [ + "CacheNamespace", + "CacheStats", "SiftClient", "SiftConnectionConfig", "config", diff --git a/python/lib/sift_client/_internal/disk_cache.py b/python/lib/sift_client/_internal/disk_cache.py new file mode 100644 index 000000000..40bdb4606 --- /dev/null +++ b/python/lib/sift_client/_internal/disk_cache.py @@ -0,0 +1,409 @@ +"""Shared on-disk key/value store used by every resource that wants to cache results. + +One :class:`DiskCache` instance lives on the :class:`SiftClient` (see +``client._disk_cache``). Resources don't construct their own — they receive +a reference and wrap it in a typed adapter that namespaces keys (e.g. +``ChannelDataCache`` in ``low_level_wrappers/data.py``). The store itself +is deliberately value-agnostic: callers hand in ``size_bytes`` for the +oversize guard, ``diskcache`` pickles whatever object the caller supplied, +and the store never needs to know what's inside. + +This module is the sibling of :mod:`._disk_cache_config` — the config +holds user intent (enabled / path / max_bytes) and the store is the live +handle keyed off that intent. + +Key behaviours pinned here so the adapter layer can stay thin: + +* Default path lives under :func:`tempfile.gettempdir` and is shared + across processes, so a fresh session reads previously-cached entries. +* The byte cap is one global budget; LRU eviction spans all resources + sharing the store (channels, calculated channels, exports, ...). +* :meth:`clear_disk` (classmethod) refuses to delete a directory that + doesn't look like a sift cache (no diskcache marker), so a typo'd + path can't take out the user's documents. +* Oversized entries are skipped with a one-shot warning per key — + otherwise diskcache's eviction loop would drain every other row + trying to fit an unfittable entry. +* Construction with ``disk_path=None`` (or after :meth:`disable`) + is a silent no-op store. Callers don't need to branch on disabled + state; reads always miss and writes are dropped. +""" + +from __future__ import annotations + +import logging +import os +import shutil +import tempfile +from pathlib import Path +from typing import TYPE_CHECKING, Any, Iterator, cast + +if TYPE_CHECKING: + import diskcache + +logger = logging.getLogger(__name__) + + +class DiskCache: + """Process-wide disk-backed key/value store. + + Wraps a :class:`diskcache.Cache` with the lifecycle management and + safety rails sift resources rely on. The instance is shared — each + resource adapter namespaces its keys (e.g. ``channel:``) so multiple + resources can write to the same store without colliding. + + When ``disk_path`` is ``None``, the instance is a silent no-op: every + ``get`` misses, every ``put`` is dropped, and ``__contains__`` is + always ``False``. This lets callers treat "caching disabled" the same + as a cold cache, with no branching needed at the read/write site. + + Attributes: + DEFAULT_DISK_PATH: Default directory for the shared cache. Lives + under :func:`tempfile.gettempdir` so it survives across + sessions of the same user but doesn't pollute the home + directory. The suffix is fixed so multiple ``SiftClient`` + instances naturally share the same store and pick up each + other's prior sessions. + DEFAULT_DISK_MAX_BYTES: Byte cap seeded on **fresh** caches + when the caller doesn't pass ``disk_max_bytes``. 4 GiB is + generous for the typical ``/tmp`` filesystem; ``diskcache`` + enforces the cap with its own SQLite-backed LRU eviction + once the bound is reached. An existing cache keeps its + previously-persisted cap on reopen — only an explicit + ``disk_max_bytes`` overrides it — so two clients pointing + at the same shared directory don't quietly resize each + other's stores. + + Args: + disk_path: Directory to persist to. ``None`` keeps the store + disabled. A previously-populated directory is reused, so a + fresh process reading the same path sees existing entries. + disk_max_bytes: Byte cap on the store. ``None`` reuses the + cache's persisted cap when the directory already exists; + for a fresh directory it seeds :attr:`DEFAULT_DISK_MAX_BYTES`. + An explicit value always overrides any persisted setting. + Ignored when ``disk_path`` is ``None``. + """ + + DEFAULT_DISK_PATH: str = os.path.join(tempfile.gettempdir(), "sift-data-cache") + DEFAULT_DISK_MAX_BYTES: int = 4 * 1024 * 1024 * 1024 + + # Marker file ``diskcache`` writes inside every cache directory. The + # classmethod :meth:`clear_disk` checks for this before any + # ``shutil.rmtree`` so a typo'd path can't wipe out an unrelated + # directory. Underscore-prefixed because it's an implementation + # detail of the safety guard, not a knob. + _DISKCACHE_MARKER: str = "cache.db" + + def __init__( + self, + *, + disk_path: str | os.PathLike[str] | None = None, + disk_max_bytes: int | None = None, + ): + # Keys we've already logged an "entry exceeds disk cap" warning + # for. Tracks the full namespaced key (e.g. ``channel:foo``), not + # the resource-side id, so two adapters that happen to share an + # id space don't collide on dedup. A successful normal put + # clears the bit so a future regression re-warns. + self._oversized_warned: set[str] = set() + self._disk: diskcache.Cache | None = None + self._disk_path: str | None = None + self._disk_max_bytes: int | None = None + if disk_path is not None: + self._open_disk(str(disk_path), disk_max_bytes) + + @classmethod + def clear_disk(cls, path: str | os.PathLike[str] | None = None) -> None: + """Delete a previously-persisted on-disk cache directory. + + Use this to drop stale caches from previous sessions, recover + from a corrupt cache, or reclaim disk space. The directory is + removed entirely; a future :meth:`enable` call at the same + path opens a fresh empty cache. + + Args: + path: Directory of the cache to clear. ``None`` (the default) + targets :attr:`DEFAULT_DISK_PATH`. + + Raises: + ValueError: If ``path`` exists but does not look like a sift + cache directory (missing the ``diskcache`` marker file). + The guard makes accidental misuse a hard error rather + than silent data loss. + """ + target = Path(path) if path is not None else Path(cls.DEFAULT_DISK_PATH) + if not target.exists(): + return + if not (target / cls._DISKCACHE_MARKER).exists(): + raise ValueError( + f"{str(target)!r} does not look like a sift data cache " + f"directory (missing {cls._DISKCACHE_MARKER!r} marker). " + f"Refusing to delete." + ) + shutil.rmtree(target) + + @property + def disk_enabled(self) -> bool: + """Whether a disk handle is currently open.""" + return self._disk is not None + + @property + def disk_path(self) -> str | None: + """Filesystem path of the cache when enabled, else ``None``.""" + return self._disk_path + + @property + def disk_max_bytes(self) -> int | None: + """Configured byte cap on disk usage, or ``None`` when disabled.""" + return self._disk_max_bytes + + def volume(self) -> int: + """Estimated bytes currently on disk for this cache. ``0`` when disabled. + + ``diskcache`` tracks this against its size cap, so the number here + is the same one its LRU eviction loop reasons about. Includes + SQLite overhead, not just raw value sizes, so the figure trends + slightly higher than the sum of caller-supplied ``size_bytes``. + """ + if self._disk is None: + return 0 + try: + return int(self._disk.volume()) + except Exception: + return 0 + + def __len__(self) -> int: + """Total entries across all adapter prefixes. ``0`` when disabled.""" + if self._disk is None: + return 0 + try: + return len(self._disk) + except Exception: + return 0 + + def __contains__(self, key: str) -> bool: + """True if ``key`` is cached. Always ``False`` when disabled.""" + if self._disk is None: + return False + return key in self._disk + + def __iter__(self) -> Iterator[str]: + """Yield cached keys. Lets adapters scope a clear to their prefix. + + Yields nothing when disabled. The underlying diskcache iterator + is snapshot-style, but callers that intend to mutate during + iteration should still wrap with ``list(...)`` to be safe. + + ``diskcache.Cache`` is typed as yielding ``bytes | str | ...`` + because it supports arbitrary key types; the cast narrows to the + ``str`` contract this layer enforces. Adapters never write + non-string keys. + """ + if self._disk is None: + return + for key in self._disk: + yield cast("str", key) + + def enable( + self, + *, + path: str | os.PathLike[str] | None = None, + max_bytes: int | None = None, + ) -> None: + """Open the disk handle, replacing any previous one. + + Reconfiguring to a different ``path`` or ``max_bytes`` closes the + prior handle first. Existing entries at the new path become + available via :meth:`get` without further setup. + + Args: + path: Directory to persist to. ``None`` uses + :attr:`DEFAULT_DISK_PATH`. + max_bytes: Byte cap. ``None`` keeps whatever the cache + already had — the persisted cap for an existing + directory, or :attr:`DEFAULT_DISK_MAX_BYTES` when the + directory is fresh. An explicit value always overrides + any persisted setting. + """ + target_path = str(path) if path is not None else self.DEFAULT_DISK_PATH + if ( + self._disk is not None + and self._disk_path == target_path + and (max_bytes is None or self._disk_max_bytes == max_bytes) + ): + return + self._close_disk() + self._open_disk(target_path, max_bytes) + + def disable(self) -> None: + """Close the disk handle (if open). Does not touch on-disk contents. + + Use :meth:`clear_disk` to remove a directory from disk. + """ + self._close_disk() + + def get(self, key: str) -> Any | None: + """Return the cached value for ``key`` or ``None`` on a miss. + + Returns ``None`` for misses, decoded values for hits, and ``None`` + (after self-invalidating the row) for corrupt entries surfaced + by ``diskcache`` as ``sqlite3.DatabaseError`` or similar. The + caller is expected to ``isinstance``-check the result against + the type they wrote. + """ + if self._disk is None: + return None + try: + return self._disk.get(key, default=None, retry=True) + except Exception: + # diskcache surfaces ``sqlite3.DatabaseError`` (and friends) + # for corrupt or partially-written entries from a prior + # session. Treat as a miss and force-drop the bad row so + # we don't repeatedly trip the same path. + logger.warning("disk cache read failed for %s; invalidating", key) + try: + del self._disk[key] + except Exception: + pass + return None + + def put(self, key: str, value: Any, *, size_bytes: int) -> None: + """Write ``value`` under ``key``. No-op when disabled. + + Entries whose ``size_bytes`` exceeds :attr:`disk_max_bytes` are + skipped with a one-shot warning per key, since diskcache's + eviction loop would otherwise drain every other row trying — and + failing — to fit an oversized entry. Callers are responsible + for measuring the size; the store stays value-agnostic. + + Args: + key: Namespaced key (e.g. ``"channel:"``). Adapters are + responsible for picking a prefix that won't collide with + other adapters writing to the same store. + value: Anything ``diskcache`` can pickle. + size_bytes: Caller-measured size used for the oversize guard. + """ + if self._disk is None: + return + if self._disk_max_bytes is not None and size_bytes > self._disk_max_bytes: + if key not in self._oversized_warned: + logger.warning( + "Entry for %s (%d bytes) is larger than the disk " + "cache cap (%d bytes); skipping disk cache for this " + "entry so other entries aren't evicted. Raise the " + "cap via ``client.cache.enable(max_bytes=...)`` " + "to cache this entry on disk.", + key, + size_bytes, + self._disk_max_bytes, + ) + self._oversized_warned.add(key) + try: + self._disk.delete(key, retry=True) + except Exception: + pass + return + try: + self._disk.set(key, value, retry=True) + self._oversized_warned.discard(key) + except Exception: + # Best-effort persistence: keep going on disk errors so the + # caller's request still succeeds. Drop the (possibly + # partial) disk row. + logger.warning("disk cache write failed for %s; invalidating", key) + try: + self._disk.delete(key, retry=True) + except Exception: + pass + + def invalidate(self, key: str) -> None: + """Remove ``key`` from the cache. Safe to call when absent.""" + # Invalidation is a fresh start for this key; the next put should + # re-evaluate against the current cap and re-warn if still too big. + self._oversized_warned.discard(key) + if self._disk is not None: + try: + self._disk.delete(key, retry=True) + except Exception: + pass + + def clear(self) -> None: + """Wipe every entry from the store. The directory itself remains. + + Spans all adapters sharing the store — typically used at test + teardown or for full reset. Adapters that want to wipe only their + own namespace should iterate ``self`` and call :meth:`invalidate` + on matching keys. + """ + self._oversized_warned.clear() + if self._disk is not None: + self._disk.clear() + + def close(self) -> None: + """Release the disk file handle. Safe to call when disabled.""" + self._close_disk() + + def __del__(self) -> None: + """Best-effort teardown for callers that don't call :meth:`close`. + + ``diskcache.Cache`` holds a SQLite handle (plus a small + connection pool) that only closes on an explicit call. A + service that builds many transient :class:`~SiftClient` + instances would otherwise leak one handle per client against + the shared cache directory — enough of them, and SQLite's + connection-per-writer limits start to bite. + + This ``__del__`` runs when the owning client is + garbage-collected. Explicit :meth:`close` (or + :meth:`CacheNamespace.disable`) remains the preferred + teardown; this is a safety net. + + Guarded because ``__del__`` also runs during interpreter + shutdown, where module globals and attributes may already be + gone; a raised exception here becomes an unraisable-error + printout, not a real failure to fix. + """ + try: + self._close_disk() + except Exception: + pass + + def _open_disk(self, path: str, max_bytes: int | None) -> None: + import diskcache + + os.makedirs(path, exist_ok=True) + # Only assert ``size_limit`` when either (a) the caller gave an + # explicit override or (b) this is a brand-new cache directory + # (no diskcache marker yet) and we need to seed the default. On + # reopen without an explicit cap, we omit ``size_limit`` so + # ``diskcache`` reuses whatever it persisted in its Settings + # table — otherwise two clients pointing at the same shared + # path would silently resize each other's stores. + kwargs: dict[str, Any] = { + "directory": path, + "eviction_policy": "least-recently-used", + "statistics": 0, + "tag_index": 0, + } + if max_bytes is not None: + kwargs["size_limit"] = max_bytes + elif not (Path(path) / self._DISKCACHE_MARKER).exists(): + kwargs["size_limit"] = self.DEFAULT_DISK_MAX_BYTES + self._disk = diskcache.Cache(**kwargs) + self._disk_path = path + # Read back the *effective* cap (persisted or newly set) so + # ``stats()`` and ``put``'s oversize guard reason about the + # real bound rather than what the caller wished for. + self._disk_max_bytes = int(self._disk.size_limit) + + def _close_disk(self) -> None: + if self._disk is None: + return + try: + self._disk.close() + except Exception: + pass + self._disk = None + self._disk_path = None + self._disk_max_bytes = None diff --git a/python/lib/sift_client/_internal/disk_cache_config.py b/python/lib/sift_client/_internal/disk_cache_config.py new file mode 100644 index 000000000..c49eaf442 --- /dev/null +++ b/python/lib/sift_client/_internal/disk_cache_config.py @@ -0,0 +1,87 @@ +"""User-expressed configuration for a resource's optional disk-cache tier.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import os + + +class DiskCacheConfig: + """Holds a resource's disk-cache enable/path/max-bytes intent. + + Resources own one instance, mutate it via :meth:`enable` / :meth:`disable` + in response to user calls, and read the properties at lazy-init time to + decide what kwargs to forward to their cache-aware wrapper. + + The :attr:`using_default_path` property is the key invariant for the + silent-fallback-vs-loud-raise distinction in resource lazy-init code: + if the user picked a specific path and opening fails, the failure + surfaces; if the user left the default and opening fails, the resource + falls back to memory-only without disrupting the call. + + Args: + enabled: Initial enabled state. Pass ``True`` for opt-out (the disk + tier is on by default and users call ``disable`` to turn it off); + pass ``False`` for opt-in (users call ``enable`` to turn it on). + """ + + def __init__(self, *, enabled: bool = True) -> None: + self._enabled = enabled + self._path: str | None = None + self._max_bytes: int | None = None + + @property + def enabled(self) -> bool: + """Whether the disk tier should be opened on the next lazy init.""" + return self._enabled + + @property + def path(self) -> str | None: + """User-supplied disk-cache path, or ``None`` to defer to the cache's default.""" + return self._path + + @property + def max_bytes(self) -> int | None: + """User-supplied disk-cache byte cap, or ``None`` to defer to the cache's default.""" + return self._max_bytes + + @property + def using_default_path(self) -> bool: + """``True`` when the disk tier is enabled *and* the path is the cache's default. + + Resources use this to decide whether to silently fall back to memory + on a disk-open failure (default path: the user didn't ask for it + specifically, so degrade gracefully) or to re-raise (explicit path: + the user asked for it, so failure must surface). + """ + return self._enabled and self._path is None + + def enable( + self, + *, + path: str | os.PathLike[str] | None = None, + max_bytes: int | None = None, + ) -> None: + """Mark the disk tier as enabled, optionally with a custom path or byte cap. + + Args: + path: Directory to persist to. ``None`` leaves the cache's + default in effect. + max_bytes: Byte cap on the disk tier. ``None`` leaves the + cache's default in effect. + """ + self._enabled = True + self._path = str(path) if path is not None else None + self._max_bytes = max_bytes + + def disable(self) -> None: + """Mark the disk tier as disabled and clear any custom path / byte cap. + + Subsequent :meth:`enable` calls re-enable at the cache's defaults + unless overrides are supplied. + """ + self._enabled = False + self._path = None + self._max_bytes = None diff --git a/python/lib/sift_client/_internal/low_level_wrappers/data.py b/python/lib/sift_client/_internal/low_level_wrappers/data.py index 57b24e398..e16cc750c 100644 --- a/python/lib/sift_client/_internal/low_level_wrappers/data.py +++ b/python/lib/sift_client/_internal/low_level_wrappers/data.py @@ -3,7 +3,7 @@ import asyncio import logging from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Tuple, cast import pandas as pd from pydantic import BaseModel, ConfigDict @@ -16,6 +16,7 @@ ) from sift.data.v2.data_pb2_grpc import DataServiceStub +from sift_client._internal.disk_cache import DiskCache from sift_client._internal.low_level_wrappers.base import LowLevelClientBase from sift_client._internal.time import to_timestamp_nanos from sift_client.sift_types.channel import Channel, ChannelDataType @@ -35,16 +36,485 @@ REQUEST_BATCH_SIZE = 1 -class ChannelCacheEntry(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True) - data: pd.DataFrame +TimeRange = Tuple[datetime, datetime] + + +class SegmentRef(BaseModel): + """Pointer to one cached segment, stored on the per-bucket index. + + ``seg_id`` is ``None`` for an **empty ref** — a record of "we + queried this range and the wire returned no data". Empty refs + contribute to coverage (so a repeat of a known-empty range + doesn't hit the wire). + """ + + model_config = ConfigDict(frozen=True) + seg_id: int | None start_time: datetime end_time: datetime -class ChannelCache(BaseModel): - name_id_map: dict[str, str] - channels: dict[str, ChannelCacheEntry] +class SegmentIndex(BaseModel): + """Per-(channel, run) index of cached segments.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + next_seg_id: int = 0 + segments: list[SegmentRef] = [] + + +class ChannelDataCache: + """Channel-side adapter over the shared :class:`DiskCache` store. + + Each ``(channel_id, run_id)`` bucket is split across two key shapes + in the underlying store: + + * One **index** entry (``channel:v1:::idx``) holding a + :class:`SegmentIndex` — a tiny list of :class:`SegmentRef` ptrs + describing every segment that exists for the bucket. Some refs + are *empty* (``seg_id is None``) — they record "we queried this + range and the wire returned no data" without a backing segment. + * One **segment** entry per fetch with data + (``channel:v1:::seg:``) holding the + :class:`pandas.DataFrame` for that fetch's slice. The + :class:`SegmentRef` on the index already carries the claimed + time range and is the source of truth for coverage, so segment + bodies are stored as raw frames. + + Reads stitch the relevant segments together via :meth:`get_range`, + which also reports which sub-ranges of the requested window have no + cached coverage (``gaps``) so the caller can fetch only the holes. + + Eviction tolerance: ``diskcache``'s LRU can drop a segment while + keeping the index. The reader treats a missing segment as a gap, + which forces a wire fetch for that range. The opposite (index + evicted, segments orphaned) is also fine — orphans are unreachable + and will eventually LRU-evict. + + Attributes: + KEY_PREFIX: Versioned namespace prefix for every key this + adapter writes to the shared :class:`DiskCache`. Picked at + class scope so adapters in other resources can pick + distinct prefixes without runtime negotiation. The trailing + ``v`` component is the schema version: bump it (e.g. + from ``"channel:v1:"`` to ``"channel:v2:"``) when either + entry shape changes incompatibly so old keys are silently + unreachable rather than mis-deserialized. + MAX_SEGMENTS_PER_BUCKET: Per-bucket cap on segment count + before :meth:`_compact_bucket` folds them into one. + Without compaction a long incrementally-pulled run grows + ``len(idx.segments)`` without bound, so every full-range + ``get_range`` would load and concat all segments — + swapping the old O(n^2) write for unbounded read fan-out. + The cap bounds that fan-out at a known constant. + + """ + + KEY_PREFIX: str = "channel:v1:" + MAX_SEGMENTS_PER_BUCKET: int = 16 + + def __init__(self, store: DiskCache): + """Wrap ``store`` with channel-data semantics. + + Args: + store: The shared :class:`DiskCache` instance owned by the + :class:`SiftClient`. Multiple adapters may share one store. + """ + self._store = store + self.name_id_map: dict[str, str] = {} + + # --- key helpers --- + + def _bucket_prefix(self, channel_id: str, run_id: str | None) -> str: + # Stem all index/segment keys for a bucket share. Empty run + # segment is safe because real run ids are UUIDs (never empty) + # so there's no collision between the run-scoped and unscoped + # buckets. + return f"{self.KEY_PREFIX}{run_id or ''}:{channel_id}" + + def _index_key(self, channel_id: str, run_id: str | None) -> str: + return f"{self._bucket_prefix(channel_id, run_id)}:idx" + + def _segment_key(self, channel_id: str, run_id: str | None, seg_id: int) -> str: + return f"{self._bucket_prefix(channel_id, run_id)}:seg:{seg_id}" + + @property + def store(self) -> DiskCache: + """The shared underlying store. Tests reach in for store-level state.""" + return self._store + + # --- public API --- + + def has_any(self, channel_id: str, run_id: str | None = None) -> bool: + """True if at least one segment is cached for this bucket. + + Cheap check (one index read) — does NOT touch segment data. + Useful for "should I consult the cache at all?" gates. + """ + idx = self._load_index(channel_id, run_id) + return idx is not None and bool(idx.segments) + + def get_range( + self, + channel_id: str, + run_id: str | None, + start_time: datetime, + end_time: datetime, + ) -> tuple[pd.DataFrame | None, list[TimeRange]]: + """Return cached data covering ``[start_time, end_time]`` plus gaps. + + Walks the bucket's segments, slices each one to the query range, + stitches them together, and reports which sub-ranges still need + a wire fetch. + + Returns: + ``(stitched_data, gaps)`` where ``stitched_data`` is the + concat of every overlapping segment sliced to the query + range (or ``None`` if no rows were found), and ``gaps`` is + the list of ``(gap_start, gap_end)`` sub-ranges within the + query window not covered by any present segment. + ``gaps == []`` means the cache fully covers the request. + """ + idx = self._load_index(channel_id, run_id) + if idx is None or not idx.segments: + return None, [(start_time, end_time)] + + # Sort by start_time so the stitch order is deterministic. + sorted_refs = sorted(idx.segments, key=lambda r: r.start_time) + + frames: list[pd.DataFrame] = [] + present_ranges: list[TimeRange] = [] + for ref in sorted_refs: + # Skip non-overlapping segments cheaply (no segment load). + if ref.end_time < start_time or ref.start_time > end_time: + continue + + clamped = (max(ref.start_time, start_time), min(ref.end_time, end_time)) + + if ref.seg_id is None: + # Empty ref: the wire was asked about this range and + # returned no data. Count it toward coverage so the + # caller doesn't refetch. + present_ranges.append(clamped) + continue + + frame = self._load_segment(channel_id, run_id, ref.seg_id) + if frame is None: + # Evicted by diskcache LRU (or index/segment got out of + # sync). Treat the segment's range as uncovered so the + # caller refetches it. + continue + + present_ranges.append(clamped) + sliced = frame[start_time:end_time] # type: ignore[misc] + if len(sliced) > 0: + frames.append(sliced) + + gaps = self._compute_gaps(start_time, end_time, present_ranges) + if not frames: + return None, gaps + if len(frames) == 1: + return frames[0], gaps + # Multiple segments → stitch. ``groupby.last()`` dedups any + # boundary-overlapping timestamps and keeps the later segment's + # value on conflict (sorted by start_time above). + return pd.concat(frames).groupby(level=0).last(), gaps + + def put_segment( + self, + channel_id: str, + run_id: str | None, + data: pd.DataFrame, + start_time: datetime, + end_time: datetime, + ) -> None: + """Write a new segment and update the index. + + Args: + channel_id: Parent channel id (bitfield elements all share + the same id — group them upstream and pass one wide + frame here). + run_id: Per-run cache dimension; ``None`` for the unscoped + bucket. + data: This fetch's rows, indexed by tz-aware + ``DatetimeIndex``. + start_time: Claimed lower bound of cache coverage for this + segment. Callers can claim more than ``data`` actually + spans to record "we asked the wire about this range" and + avoid re-fetching empty sub-ranges. + end_time: Claimed upper bound; same semantics as start. + + Per-fetch disk write is O(this segment) — no rewrite of any + already-cached segment. + + Write order is segment-first, then index: an interrupted update + leaves an unreachable orphan segment (harmless; LRU-evicts) but + never leaves the index pointing at a missing segment. Empty + puts skip the segment write entirely and update the index + directly. + """ + idx = self._load_index(channel_id, run_id) or SegmentIndex() + + if len(data) == 0: + # Empty ref: coverage-only entry, no segment body. Doesn't + # bump ``next_seg_id`` because there's no segment to key. + idx.segments.append(SegmentRef(seg_id=None, start_time=start_time, end_time=end_time)) + self._write_index(channel_id, run_id, idx) + else: + if not data.index.is_monotonic_increasing: + data = data.sort_index() + size_bytes = int(data.memory_usage(deep=True).sum()) + seg_id = idx.next_seg_id + + # Segment first. + seg_key = self._segment_key(channel_id, run_id, seg_id) + self._store.put(seg_key, data, size_bytes=size_bytes) + if seg_key not in self._store: + # The store rejected the segment (oversize, disabled, + # etc.). Skipping the index update keeps us from + # leaving a dangling reference to a never-written + # segment. + return + + idx.segments.append(SegmentRef(seg_id=seg_id, start_time=start_time, end_time=end_time)) + idx.next_seg_id += 1 + self._write_index(channel_id, run_id, idx) + + # Bound read fan-out: collapse all segments into one once the + # bucket would carry more than the cap. + if len(idx.segments) > self.MAX_SEGMENTS_PER_BUCKET: + self._compact_bucket(channel_id, run_id) + + def invalidate(self, channel_id: str, run_id: str | None = None) -> None: + """Drop every segment in a bucket plus the index. Safe when absent. + + Only touches the one ``(channel_id, run_id)`` bucket — segments + under other runs survive. Empty refs carry no segment key, so + only data refs (``seg_id is not None``) hit the store here. + """ + idx = self._load_index(channel_id, run_id) + if idx is None: + return + for ref in idx.segments: + if ref.seg_id is None: + continue + self._store.invalidate(self._segment_key(channel_id, run_id, ref.seg_id)) + self._store.invalidate(self._index_key(channel_id, run_id)) + + def clear(self) -> None: + """Wipe every channel entry (all buckets, all runs).""" + for key in list(self._store): + if key.startswith(self.KEY_PREFIX): + self._store.invalidate(key) + + # --- internal helpers --- + + def _load_index(self, channel_id: str, run_id: str | None) -> SegmentIndex | None: + raw = self._store.get(self._index_key(channel_id, run_id)) + if not isinstance(raw, SegmentIndex): + return None + return raw + + def _load_segment( + self, channel_id: str, run_id: str | None, seg_id: int + ) -> pd.DataFrame | None: + raw = self._store.get(self._segment_key(channel_id, run_id, seg_id)) + if not isinstance(raw, pd.DataFrame): + return None + return raw + + def _write_index(self, channel_id: str, run_id: str | None, idx: SegmentIndex) -> None: + # Pydantic indexes serialize small (refs are scalars), but we + # still pass a size hint that scales with segment count so the + # store's per-entry-size eviction sees an honest weight. + self._store.put( + self._index_key(channel_id, run_id), + idx, + size_bytes=max(1024, 128 * len(idx.segments)), + ) + + def _compact_bucket(self, channel_id: str, run_id: str | None) -> None: + """Fold every ref in a bucket into a single merged ref. + + Called by :meth:`put_segment` once the per-bucket ref count + crosses :attr:`MAX_SEGMENTS_PER_BUCKET`. The motivation is read + fan-out: without compaction ``get_range`` on a long + incrementally-pulled run would walk one ref per prior fetch, + so the per-read cost climbs linearly in the number of writes. + After compaction the bucket holds at most one data ref (or one + empty ref, if every prior fetch returned no data) until the + next put. + + Empty refs (``seg_id is None``) carry coverage but no body; + their claimed ``[start, end]`` is folded into the merged ref's + claim along with the data refs. If every loadable body went + away (all-empty or every data ref LRU-evicted), the bucket + compacts to a single empty ref covering the union of all + prior claims so the "no data here" coverage is preserved. + + Write ordering (crash safety): + + 1. Write the merged segment under a fresh ``seg_id`` so a + crash here leaves the old segments + old index intact + (the orphan merged seg LRU-evicts). + 2. Rewrite the index to point at *only* the merged seg. A + crash here leaves the old segs as orphans (LRU-evicts). + 3. Delete the old segment keys. Any crash here is harmless + since the index already points elsewhere. + + Cost note: under a flat cap M, total work across N fetches is + O(N^2 / M) — quadratic in fetches with the constant reduced by + the cap. That's the documented first-cut; if compaction shows + up as a hot spot in stats, the next step is LSM-style + geometric levels rather than tuning M. + """ + idx = self._load_index(channel_id, run_id) + if idx is None or len(idx.segments) <= self.MAX_SEGMENTS_PER_BUCKET: + # Bucket already cleaned up (concurrent invalidate, LRU, + # or a parallel compactor) — nothing to do. + return + + refs = sorted(idx.segments, key=lambda r: r.start_time) + + # Empty refs have no body to load — their contribution is + # purely coverage, which we fold into the merged claim below. + # LRU may also have dropped a data ref's body; those refs + # degrade to gaps (we don't merge missing data). + frames: list[pd.DataFrame] = [] + for ref in refs: + if ref.seg_id is None: + continue + frame = self._load_segment(channel_id, run_id, ref.seg_id) + if frame is not None: + frames.append(frame) + + old_seg_ids: list[int] = [r.seg_id for r in refs if r.seg_id is not None] + + # Claimed coverage of the merged ref spans *every* ref's claim + # (data + empty + evicted) — the index already represents "we + # asked the wire about this range" and compaction must preserve + # that, even if the underlying data went away. + merged_start = min(r.start_time for r in refs) + merged_end = max(r.end_time for r in refs) + + if not frames: + # No usable data bodies left — bucket collapses to a single + # empty ref covering the union of all prior claims. Future + # reads in that range are cache hits returning no rows. + self._replace_index_and_drop_segments( + channel_id, + run_id, + new_segments=[ + SegmentRef(seg_id=None, start_time=merged_start, end_time=merged_end) + ], + next_seg_id=idx.next_seg_id, + drop_seg_ids=old_seg_ids, + ) + logger.debug( + "compacted bucket %r:%r from %d refs to 1 empty ref", + run_id or "", + channel_id, + len(refs), + ) + return + + # Mirror :meth:`get_range`'s dedup: later segments win on + # boundary-overlapping timestamps (refs are already sorted by + # start_time above). + merged_data = frames[0] if len(frames) == 1 else pd.concat(frames).groupby(level=0).last() + merged_size = int(merged_data.memory_usage(deep=True).sum()) + new_seg_id = idx.next_seg_id + + # Step 1: write the merged segment first. + new_seg_key = self._segment_key(channel_id, run_id, new_seg_id) + self._store.put(new_seg_key, merged_data, size_bytes=merged_size) + if new_seg_key not in self._store: + # Store rejected (typically oversize). Leave the bucket + # alone so the per-fetch segments remain readable; the + # next ``put_segment`` will retry the cap check. + logger.debug( + "compaction skipped for bucket %r:%r (merged size %d exceeded store cap)", + run_id or "", + channel_id, + merged_size, + ) + return + + # Steps 2 + 3: rewrite the index to point at only the merged + # seg, then drop superseded segment keys. + self._replace_index_and_drop_segments( + channel_id, + run_id, + new_segments=[ + SegmentRef(seg_id=new_seg_id, start_time=merged_start, end_time=merged_end) + ], + next_seg_id=new_seg_id + 1, + drop_seg_ids=old_seg_ids, + ) + logger.debug( + "compacted bucket %r:%r from %d refs to 1 (merged size %d)", + run_id or "", + channel_id, + len(refs), + merged_size, + ) + + def _replace_index_and_drop_segments( + self, + channel_id: str, + run_id: str | None, + *, + new_segments: list[SegmentRef], + next_seg_id: int, + drop_seg_ids: list[int], + ) -> None: + """Rewrite the index, then invalidate superseded segment keys. + + Order matters for crash safety: the new index goes down first + so that an interrupt before the deletes leaves the old segs as + unreachable orphans (which LRU-evict) rather than leaving the + index pointing at deleted keys. + """ + self._write_index( + channel_id, + run_id, + SegmentIndex(next_seg_id=next_seg_id, segments=new_segments), + ) + for seg_id in drop_seg_ids: + self._store.invalidate(self._segment_key(channel_id, run_id, seg_id)) + + @staticmethod + def _compute_gaps( + query_start: datetime, + query_end: datetime, + covered: list[TimeRange], + ) -> list[TimeRange]: + """Sub-ranges of ``[query_start, query_end]`` not in ``covered``. + + ``covered`` is the list of segment ranges (already clamped to + the query window). Algorithm: merge overlapping/adjacent + intervals, then sweep emitting gaps between them. + """ + if not covered: + return [(query_start, query_end)] + + sorted_ranges = sorted(covered, key=lambda r: r[0]) + merged: list[TimeRange] = [sorted_ranges[0]] + for seg_start, seg_end in sorted_ranges[1:]: + last_start, last_end = merged[-1] + if seg_start <= last_end: + merged[-1] = (last_start, max(last_end, seg_end)) + else: + merged.append((seg_start, seg_end)) + + gaps: list[TimeRange] = [] + cursor = query_start + for seg_start, seg_end in merged: + if seg_start > cursor: + gaps.append((cursor, seg_start)) + cursor = max(cursor, seg_end) + if cursor < query_end: + gaps.append((cursor, query_end)) + return gaps class DataLowLevelClient(LowLevelClientBase, WithGrpcClient): @@ -53,15 +523,29 @@ class DataLowLevelClient(LowLevelClientBase, WithGrpcClient): This class provides a thin wrapper around the autogenerated bindings for the DataAPI. """ - channel_cache: ChannelCache = ChannelCache(name_id_map={}, channels={}) - - def __init__(self, grpc_client: GrpcClient): + def __init__( + self, + grpc_client: GrpcClient, + *, + channel_cache: ChannelDataCache | None = None, + ): """Initialize the DataLowLevelClient. Args: grpc_client: The gRPC client to use for making API calls. + channel_cache: Adapter wrapping the shared :class:`DiskCache` the + :class:`SiftClient` owns. When ``None`` (only the unit-test + construction path), the wrapper falls back to a no-op store + so cache reads/writes are silent. Production callers always + pass an adapter built from ``client._get_disk_cache()``. """ super().__init__(grpc_client) + # Production wires the shared store in via the resource. The fallback + # here lets a bare ``DataLowLevelClient(MagicMock())`` keep working + # in unit tests without forcing every site to plumb a store. + if channel_cache is None: + channel_cache = ChannelDataCache(DiskCache()) + self.channel_cache = channel_cache def _update_name_id_map(self, channels: list[Channel]): """Update the name id map with the new channels.""" @@ -105,122 +589,129 @@ async def _get_data_impl( response = cast("GetDataResponse", response) return response.data, response.next_page_token # type: ignore # mypy doesn't know RepeatedCompositeFieldContainer can be treated like a list - def _filter_cached_channels(self, channel_ids: list[str]) -> tuple[list[str], list[str]]: - cached_channels = [] - not_cached_channels = [] - for channel_id in channel_ids: - if self.channel_cache.channels.get(channel_id): - cached_channels.append(channel_id) - else: - not_cached_channels.append(channel_id) - return cached_channels, not_cached_channels - - def _check_cache( - self, - *, - channel_id: str, - start_time: datetime, - end_time: datetime, - run_id: str | None = None, - ) -> tuple[pd.DataFrame | None, datetime | None, datetime | None]: - """Check if the data for a channel during a run is cached and return how to query remaining data if so. - - There are a variety of requested start/end time vs cached start/end time cases to consider. - Below diagram represents time aligned ranges for each case: - - Cache interval: |-------------------------------| - Case 1: |---------------------------| - Case 2: |--------------------------------| - Case 3: |----------| - Case 4: |--------------------------------| - Case 5: |------| or |-----------------------------------------| - - Returns: - A tuple of (data, start_time, end_time) - where data is a pandas dataframe and start and end times are what should be used for the next call based on what is not covered by the cached data. - """ - cached_data = self.channel_cache.channels.get(channel_id) - ret_start_time = start_time - ret_end_time = end_time - ret_data = None - if cached_data: - start_time_cached = cached_data.start_time - end_time_cached = cached_data.end_time - ret_data = cached_data.data - # Filter data to desiredtime range - ret_data = ret_data[start_time:end_time] # type: ignore # mypy doesn't understand pandas that well seemingly - - if start_time_cached <= start_time: - if start_time < end_time_cached: - if end_time <= end_time_cached: - # Case 1 - ret_start_time = None # type: ignore - ret_end_time = None # type: ignore - else: - # Case 2 - ret_start_time = end_time_cached - ret_end_time = end_time - else: - # Case 3 - return (None, start_time, end_time) - else: - if start_time_cached < end_time and end_time <= end_time_cached: - # Case 4 - ret_start_time = start_time - ret_end_time = start_time_cached - else: - # Case 5 - return (None, start_time, end_time) - - return (ret_data, ret_start_time, ret_end_time) - def _update_cache( self, *, channel_data: dict[str, pd.DataFrame], + fetched_ranges_per_channel: dict[str, list[TimeRange]], start_time: datetime, end_time: datetime, run_id: str | None = None, ): - """Update the cache with the new data and start/end times.""" + """Write each channel's fresh data or empty-ref as a new segment. + + Per-fetch disk write is O(this segment) — no merging with prior + segments and no re-pickle of accumulated data, so n sequential + incremental pulls cost O(n) total disk write instead of O(n²). + + Bitfield grouping: ``try_deserialize_channel_data`` returns one + dotted-name DataFrame per bitfield element, all mapping to the + same parent ``channel_id`` via :attr:`name_id_map`. We group + them by parent id and concat into one wide frame so each fetch + produces exactly one segment per channel, regardless of how + many elements that channel exposes. + + Empty results are recorded as empty refs (no segment body) — + one ref per fetched range — so a repeat of a known-empty range + is a cache hit returning no rows instead of another wire call. + Only the **unscoped** path records empty refs: run-scoped + absence isn't assertable and would permanently mask data an + ongoing run ingests after the empty query (subsequent same- + range queries would hit the cache and never refetch). + + Args: + channel_data: Merged per-name frames coming out of + :meth:`_merge_pages`. A channel absent from this dict + returned zero rows across every page; a present channel + with a zero-row frame had every page return zero rows + (bitfield elements that all came back empty). + fetched_ranges_per_channel: Per-channel-id list of the + ``[start, end]`` windows we actually asked the wire + about this call. Drives empty-ref recording: any cid + in this dict whose data dropped out for the empty case + gets one empty ref per range. Pure cache hits don't + reach this method (the caller skips when nothing was + fetched). + """ assert start_time is not None assert end_time is not None name_id_map = self.channel_cache.name_id_map + # Group dotted-name frames by parent channel id so bitfield + # elements land in one segment. + by_channel_id: dict[str, list[pd.DataFrame]] = {} for channel_name, data in channel_data.items(): channel_id = name_id_map.get(channel_name) if not channel_id: raise ValueError( - f"{channel_name} not found in name_id_map. Not sure got data for this channel without a call that should've updated the map." + f"{channel_name} not found in name_id_map. Not sure got " + f"data for this channel without a call that should've " + f"updated the map." ) + by_channel_id.setdefault(channel_id, []).append(data) - suggested_start_time = start_time + ids_with_data: set[str] = set() + for channel_id, frames in by_channel_id.items(): + if len(frames) == 1: + combined = frames[0] + else: + # Bitfield: per-element single-column frames → one wide + # frame. ``groupby.last`` dedups any boundary overlaps. + combined = pd.concat(frames).groupby(level=0).last() + + if len(combined) == 0: + # Bitfield with every element empty, or other edge + # cases — handled by the empty-ref loop below. + continue + + ids_with_data.add(channel_id) + + # Segment coverage range. For run-scoped queries, claim + # only what the data actually spans (we can't assert + # absence outside the data — the run might not have + # started yet). For unscoped queries, claim the full + # requested range so a follow-up of the same range hits. + seg_end = end_time if run_id: - if len(data) > 0: - suggested_start_time = data.index[0] - else: - # Because we didn't get any data, we can't know what the start time should be. - # And because this was queried w/ a run ID, we can't say there's no data before the run started. - # So we just don't update the cache. - continue - - if channel_id in self.channel_cache.channels: - self.channel_cache.channels[channel_id].data = ( - pd.concat([self.channel_cache.channels[channel_id].data, data]) - .groupby(level=0) - .last() - ) - self.channel_cache.channels[channel_id].start_time = min( - suggested_start_time, self.channel_cache.channels[channel_id].start_time - ) - self.channel_cache.channels[channel_id].end_time = max( - end_time, self.channel_cache.channels[channel_id].end_time - ) + # ``combined.index`` is a ``DatetimeIndex`` (built from + # the wire's nanosecond timestamps), so ``index[0]`` is + # always ``pd.Timestamp`` at runtime; pandas-stubs types + # it as the wider ``Scalar`` union. + seg_start = cast("pd.Timestamp", combined.index[0]).to_pydatetime() else: - self.channel_cache.channels[channel_id] = ChannelCacheEntry( - data=data, - start_time=suggested_start_time, - end_time=end_time, + seg_start = start_time + + self.channel_cache.put_segment( + channel_id=channel_id, + run_id=run_id, + data=combined, + start_time=seg_start, + end_time=seg_end, + ) + + # Empty refs for channels that were queried this call but came + # back with zero rows. One ref per fetched range so the + # coverage claim matches what we actually asked the wire about + # — overclaiming would let a future query for an adjacent + # range hit the cache and miss data that actually exists. + # + # Skipped for run-scoped queries: absence at query time + # doesn't imply future absence (the run may still be + # ingesting), and an empty ref would permanently mask data + # that arrives later on the same window. + if run_id: + return + empty_df = pd.DataFrame() + for cid, ranges in fetched_ranges_per_channel.items(): + if cid in ids_with_data: + continue + for fetched_start, fetched_end in ranges: + self.channel_cache.put_segment( + channel_id=cid, + run_id=run_id, + data=empty_df, + start_time=fetched_start, + end_time=fetched_end, ) async def get_channel_data( @@ -235,23 +726,56 @@ async def get_channel_data( ignore_cache: bool = False, ) -> dict[str, pd.DataFrame]: """Get the data for a channel during a run.""" - ret_data = {} + ret_data: dict[str, pd.DataFrame] = {} # No data will be returned if end_time is not provided. start_time = start_time or datetime.fromtimestamp(0, tz=timezone.utc) end_time = end_time or datetime.now(timezone.utc) self._update_name_id_map(channels) - channel_ids = [c.id_ for c in channels] - cached_channels, not_cached_channels = ( - ([], channel_ids) if ignore_cache else self._filter_cached_channels(channel_ids) # type: ignore - ) + + # Two work queues. Fully uncached channels share the full range + # and get batched; partial-hit channels carry per-gap ranges + # and go one fetch at a time. ``fetched_ranges_per_channel`` + # records the exact ``[s, e]`` windows we asked the wire about, + # per channel — used downstream to record empty refs (a "no + # data here" coverage claim) so a repeat of a known-empty + # query is a cache hit instead of another wire call. + fully_uncached: list[str] = [] + partial_gaps: list[tuple[str, list[TimeRange]]] = [] + fetched_ranges_per_channel: dict[str, list[TimeRange]] = {} + + for channel in channels: + cid = channel.id_ + assert cid is not None + if ignore_cache: + cached_data: pd.DataFrame | None = None + gaps: list[TimeRange] = [(start_time, end_time)] + else: + cached_data, gaps = self.channel_cache.get_range(cid, run_id, start_time, end_time) + + if cached_data is not None: + # Slice per column so each result key carries only its + # own element frame (matches the per-element shape + # ``try_deserialize_channel_data`` produces; without + # this slice, a bitfield's wide cached frame would land + # under every dotted key). + for name in cached_data.columns: + ret_data[name] = cached_data[[name]] + + if not gaps: + continue + fetched_ranges_per_channel.setdefault(cid, []).extend(gaps) + if len(gaps) == 1 and gaps[0] == (start_time, end_time): + fully_uncached.append(cid) + else: + partial_gaps.append((cid, gaps)) tasks = [] - # Queue up calls for non-cached channels in batches. + # Batch fully-uncached channels (sharing the full requested + # range) into one wire call each. batch_size = REQUEST_BATCH_SIZE - for i in range(0, len(not_cached_channels), batch_size): # type: ignore - batch = not_cached_channels[i : i + batch_size] # type: ignore - + for i in range(0, len(fully_uncached), batch_size): + batch = fully_uncached[i : i + batch_size] task = asyncio.create_task( self._handle_pagination( self._get_data_impl, @@ -267,51 +791,75 @@ async def get_channel_data( ) tasks.append(task) - # Handling cached channels 1 by 1 instead of in batches to account for channels that may have been cached from calls with different start/end times. - for channel_id in cached_channels: - cached_data, new_start_time, new_end_time = self._check_cache( - channel_id=channel_id, + # Partial gaps: one fetch per (channel, gap). + for cid, gaps in partial_gaps: + for gap_start, gap_end in gaps: + task = asyncio.create_task( + self._handle_pagination( + self._get_data_impl, + kwargs={ + "channel_ids": [cid], + "run_id": run_id, + "start_time": gap_start, + "end_time": gap_end, + }, + page_size=page_size, + max_results=max_results, + ) + ) + tasks.append(task) + + pages = await asyncio.gather(*tasks) + ret_data = self._merge_pages(pages, initial=ret_data) + + # Pure cache hits never reach ``_update_cache`` because nothing + # was fetched (``fetched_ranges_per_channel`` is empty). When + # we did fetch, ``_update_cache`` writes both data segments and + # empty refs — the latter covers the "asked the wire, got + # nothing" case so a repeat doesn't refetch. + if not ignore_cache and fetched_ranges_per_channel: + self._update_cache( + channel_data=ret_data, + fetched_ranges_per_channel=fetched_ranges_per_channel, start_time=start_time, end_time=end_time, run_id=run_id, ) - if cached_data is not None: - for name in cached_data.columns: - ret_data[name] = cached_data - if new_start_time is None: - # Cache fully encompassed the desired time range so don't queue a call. - continue - task = asyncio.create_task( - self._handle_pagination( - self._get_data_impl, - kwargs={ - "channel_ids": [channel_id], - "run_id": run_id, - "start_time": new_start_time, - "end_time": new_end_time or end_time, - }, - page_size=page_size, - max_results=max_results, - ) - ) - tasks.append(task) + return ret_data - pages = await asyncio.gather(*tasks) - # Flatten the data + def _merge_pages( + self, + pages: list[list[Any]], + *, + initial: dict[str, pd.DataFrame], + ) -> dict[str, pd.DataFrame]: + """Flatten paged channel data + any cached slices into one DataFrame per channel. + + ``initial`` carries the cached slices ``get_channel_data`` + stitched inline via :meth:`ChannelDataCache.get_range` before + dispatching wire fetches for the gaps. Cached entries are + folded in as the first frame for their channel so they + participate in the same final concat; ``groupby(level=0).last()`` + preserves the previous behavior of letting a later-positioned + (fresher) value win on duplicate timestamps. + """ + per_channel_frames: dict[str, list[pd.DataFrame]] = {} for page in pages: for data in page: - page_results = self.try_deserialize_channel_data(data) - for name, df in page_results.items(): - if name not in ret_data: - ret_data[name] = df - else: - ret_data[name] = pd.concat([ret_data[name], df]).groupby(level=0).last() - - self._update_cache( - channel_data=ret_data, start_time=start_time, end_time=end_time, run_id=run_id - ) - + for name, df in self.try_deserialize_channel_data(data).items(): + per_channel_frames.setdefault(name, []).append(df) + + ret_data: dict[str, pd.DataFrame] = dict(initial) + for name, frames in per_channel_frames.items(): + if name in ret_data: + # Cached slice goes first so fresher pages (positioned later + # in the list) win on overlapping timestamps after groupby. + frames.insert(0, ret_data[name]) + if len(frames) == 1: + ret_data[name] = frames[0] + else: + ret_data[name] = pd.concat(frames).groupby(level=0).last() return ret_data @staticmethod diff --git a/python/lib/sift_client/_tests/_internal/low_level_wrappers/test_data.py b/python/lib/sift_client/_tests/_internal/low_level_wrappers/test_data.py new file mode 100644 index 000000000..ba5a8f98c --- /dev/null +++ b/python/lib/sift_client/_tests/_internal/low_level_wrappers/test_data.py @@ -0,0 +1,1619 @@ +"""Tests for :mod:`sift_client._internal.low_level_wrappers.data`. + +Five classes, narrowest scope first: + +* :class:`TestChannelDataCache` — the typed adapter over the shared + :class:`DiskCache`. Covers key namespacing, eviction-tolerant + stitching, gap computation, and the prefix-scoped ``clear``. +* :class:`TestMergePages` — ``DataLowLevelClient._merge_pages``, the + per-channel concat helper. +* :class:`TestDataLowLevelClient` — constructor wiring and per-instance + isolation. +* :class:`TestGetChannelData` — end-to-end on the public + ``get_channel_data`` API against a mocked ``_get_data_impl``. +* :class:`TestBitFieldChannels` — the channel-id-vs-channel-name seam + where bitfield elements share one parent id but surface as multiple + dotted-name keys in the result. + +Storage-layer behaviour (oversize guards, marker-checked clear, +cross-session reload) lives in ``_tests/_internal/test_disk_cache.py``; +this file stays focused on the channel-data path. + +The OOM regression that motivated this code happened because the cache +was a class attribute that grew without bound. ``test_per_instance_isolation`` +is the canary that catches anyone re-introducing that pattern, even though +ownership has since moved to the client. + +The cache stores data as per-fetch *segments* under one index per +``(channel_id, run_id)`` bucket. ``put_segment`` writes a new segment ++ updates the index; ``get_range`` walks the index, slices each +overlapping segment to the query window, and reports any uncovered +sub-ranges as gaps. The segment model is what eliminates the prior +write amplification on incremental pulls (where every fetch re-pickled +the entire accumulated frame). +""" + +from __future__ import annotations + +from contextlib import contextmanager +from datetime import datetime, timedelta, timezone +from typing import Any, Iterator, cast +from unittest.mock import MagicMock, patch + +import pandas as pd +import pytest + +from sift_client._internal.disk_cache import DiskCache +from sift_client._internal.low_level_wrappers.data import ( + ChannelDataCache, + DataLowLevelClient, +) +from sift_client.sift_types.channel import ( + Channel, + ChannelBitFieldElement, + ChannelDataType, +) + +_NOW = datetime(2025, 1, 1, tzinfo=timezone.utc) +_WINDOW_END = _NOW + timedelta(days=1) + + +# ---------- shared helpers ----------- + + +def _frame( + cid: str = "value", + *, + rows: int = 5, + start: datetime = _NOW, + offset: int = 0, + freq: str = "ms", + value_dtype: str = "float64", +) -> pd.DataFrame: + """DataFrame indexed by a tz-aware DatetimeIndex with ``rows`` rows.""" + index = pd.date_range(start, periods=rows, freq=freq, tz=timezone.utc) + return pd.DataFrame( + {cid: [(offset + i) * 1.0 for i in range(rows)]}, + index=index, + ).astype({cid: value_dtype}) + + +def _channel(cid: str) -> Channel: + """Minimal ``Channel`` with required fields populated.""" + return Channel( + id_=cid, + name=cid, + data_type=ChannelDataType.DOUBLE, + description="", + unit="", + asset_id="a1", + is_archived=False, + created_date=_NOW, + modified_date=_NOW, + created_by_user_id="u1", + modified_by_user_id="u1", + ) + + +def _bitfield_channel(*, cid: str, name: str, elements: list[str]) -> Channel: + """Bitfield ``Channel`` with the named elements (8-bit, indexed in order). + + ``cid`` and ``name`` are kept distinct so test assertions can distinguish + "is the cache keyed by id?" from "is the result dict keyed by name?". + """ + return Channel( + id_=cid, + name=name, + data_type=ChannelDataType.BIT_FIELD, + description="", + unit="", + asset_id="a1", + is_archived=False, + bit_field_elements=[ + ChannelBitFieldElement(name=el, index=i, bit_count=8) for i, el in enumerate(elements) + ], + created_date=_NOW, + modified_date=_NOW, + created_by_user_id="u1", + modified_by_user_id="u1", + ) + + +def _client_with_cache(tmp_path, subdir: str = "cache") -> DataLowLevelClient: + """Build a ``DataLowLevelClient`` whose adapter points at ``tmp_path``. + + Tests that exercise cache behaviour (hits/misses) need an actual + disk-backed adapter, so the store has to be opened explicitly. A + plain ``DataLowLevelClient(MagicMock())`` defaults to a no-op store + and would silently turn every cache test into a wire-path test. + """ + store = DiskCache(disk_path=tmp_path / subdir) + return DataLowLevelClient(MagicMock(), channel_cache=ChannelDataCache(store)) + + +def _patch_deserializer(sentinel_to_frames: dict[str, dict[str, pd.DataFrame]]) -> Any: + """Patch ``try_deserialize_channel_data`` to translate string sentinels. + + Lets tests pass strings in lieu of protos. Returned object is a context + manager; callers use ``with _patch_deserializer(...):``. + """ + return patch.object( + DataLowLevelClient, + "try_deserialize_channel_data", + staticmethod(lambda s: sentinel_to_frames[s]), + ) + + +@contextmanager +def _fake_grpc( + client: DataLowLevelClient, + channel_to_pages: dict[str, list[pd.DataFrame | dict[str, pd.DataFrame]]], +) -> Iterator[list[dict[str, Any]]]: + """Mock the gRPC boundary so each "page" is a sentinel string. + + ``_get_data_impl`` is replaced with a coroutine that pops one page off + ``channel_to_pages[cid]`` per call per channel, until exhausted. + ``try_deserialize_channel_data`` is patched to map the sentinel back + to the corresponding ``{channel: DataFrame}`` dict. + + A page entry can be either: + + * ``pd.DataFrame`` — wrapped as ``{cid: df}`` (the single-channel + shape ``try_deserialize_channel_data`` returns for non-bitfield + channels). + * ``dict[str, pd.DataFrame]`` — used as-is (the multi-name shape + ``BitFieldValues`` produces, with keys like ``"."`` + per bitfield element). + + Yields a ``call_log`` list so tests can assert which channels actually + hit the wire. The patch is torn down and ``_get_data_impl`` restored + on exit. + """ + sentinel_to_frames: dict[str, dict[str, pd.DataFrame]] = {} + next_page_index: dict[str, int] = dict.fromkeys(channel_to_pages, 0) + call_log: list[dict[str, Any]] = [] + + async def fake_impl( + *, + channel_ids: list[str], + page_size: int | None = None, + page_token: str | None = None, + order_by: str | None = None, + **kwargs: Any, + ) -> tuple[list[str], str]: + call_log.append({"channel_ids": list(channel_ids), **kwargs}) + data: list[str] = [] + more_remaining = False + for cid in channel_ids: + i = next_page_index[cid] + if i >= len(channel_to_pages[cid]): + continue # this channel is exhausted; just emit nothing + sentinel = f"{cid}|{i}" + page = channel_to_pages[cid][i] + sentinel_to_frames[sentinel] = dict(page) if isinstance(page, dict) else {cid: page} + data.append(sentinel) + next_page_index[cid] += 1 + if next_page_index[cid] < len(channel_to_pages[cid]): + more_remaining = True + # ``_handle_pagination`` loops until it sees ``page_token == ""``. + return data, ("next" if more_remaining else "") + + original_impl = client._get_data_impl + client._get_data_impl = fake_impl # type: ignore[method-assign] + try: + with _patch_deserializer(sentinel_to_frames): + yield call_log + finally: + client._get_data_impl = original_impl # type: ignore[method-assign] + + +# ---------- tests ----------- + + +def _put( + adapter: ChannelDataCache, + channel_id: str, + *, + data: pd.DataFrame | None = None, + rows: int = 5, + start: datetime = _NOW, + offset: int = 0, + freq: str = "ms", + seg_start: datetime | None = None, + seg_end: datetime | None = None, + run_id: str | None = None, +) -> pd.DataFrame: + """Convenience: build a frame, write it as one segment, return it. + + ``seg_start`` / ``seg_end`` default to the data's actual range so + tests get tightly-bounded segments unless they specifically want to + claim extra coverage. + """ + if data is None: + data = _frame(channel_id, rows=rows, start=start, offset=offset, freq=freq) + if seg_start is None: + seg_start = cast("pd.Timestamp", data.index[0]).to_pydatetime() + if seg_end is None: + seg_end = cast("pd.Timestamp", data.index[-1]).to_pydatetime() + adapter.put_segment( + channel_id=channel_id, + run_id=run_id, + data=data, + start_time=seg_start, + end_time=seg_end, + ) + return data + + +class TestChannelDataCache: + """The typed adapter over the shared :class:`DiskCache`. + + Five invariants get pinned across the per-segment shape: + + 1. Every operation routes through the namespaced key + (``channel:v1:::{idx,seg:N}``), so two adapters sharing + one store don't collide on bare resource ids. + 2. Run id is part of the cache dimension: the same ``channel_id`` + under two different runs is two cache buckets, not one. + 3. :meth:`ChannelDataCache.get_range` stitches multiple segments + and reports uncovered sub-ranges as gaps. Missing (evicted) + segments degrade to gaps, never to errors. + 4. :meth:`ChannelDataCache.invalidate` drops every segment in a + bucket and the index, leaving other buckets untouched. + 5. :meth:`ChannelDataCache.clear` wipes only the adapter's namespace + — entries belonging to other adapters survive. + + Store-level behaviour (oversized guards, cross-session reload, + marker-checked clear_disk) is exercised in ``test_disk_cache.py``. + """ + + def test_miss_returns_none_and_full_gap(self, tmp_path): + adapter = ChannelDataCache(DiskCache(disk_path=tmp_path / "miss")) + try: + assert not adapter.has_any("c1") + data, gaps = adapter.get_range("c1", None, _NOW, _WINDOW_END) + assert data is None + assert gaps == [(_NOW, _WINDOW_END)] + finally: + adapter.store.close() + + def test_round_trip_single_segment(self, tmp_path): + """Put one segment, get_range back covers the whole frame and reports no gap.""" + adapter = ChannelDataCache(DiskCache(disk_path=tmp_path / "rt")) + try: + df = _put(adapter, "c1", rows=8) + assert adapter.has_any("c1") + got, gaps = adapter.get_range("c1", None, df.index[0], df.index[-1]) + assert got is not None + pd.testing.assert_frame_equal(got, df) + assert gaps == [] + finally: + adapter.store.close() + + def test_writes_use_namespaced_index_and_segment_keys(self, tmp_path): + """The raw store sees ``channel:v1:::idx`` + ``...:seg:0``. + + Pins the per-segment key shape: one index plus one segment key + per fetch. Without the prefix, a second adapter that happens to + share an id with the channel adapter would clobber the rows. + The ``v1`` is the schema version baked into the prefix so a + bump silently retires the entire old keyspace. + """ + store = DiskCache(disk_path=tmp_path / "ns") + adapter = ChannelDataCache(store) + try: + _put(adapter, "c1", rows=4) + assert "channel:v1::c1:idx" in store + assert "channel:v1::c1:seg:0" in store + assert "c1" not in store + assert "channel:c1" not in store # never the bare-id shape + assert "channel::c1:idx" not in store # never the unversioned shape + finally: + store.close() + + def test_run_id_is_part_of_the_key(self, tmp_path): + """Same channel under two runs is two cache buckets, not one. + + Regression guard for the run-scoping bug: a bare ``channel:`` + key conflated runs and served run-A's data to a query for run-B. + With segments, that turns into "the two buckets share an index" + — the test below pins them as fully independent. + """ + adapter = ChannelDataCache(DiskCache(disk_path=tmp_path / "runs")) + try: + df_a = _put(adapter, "c1", rows=4, run_id="run-A") + df_b = _put(adapter, "c1", rows=8, run_id="run-B") + + assert adapter.has_any("c1", "run-A") + assert adapter.has_any("c1", "run-B") + assert not adapter.has_any("c1") # unscoped bucket stays empty + assert not adapter.has_any("c1", "run-C") # unknown run still misses + + got_a, _ = adapter.get_range("c1", "run-A", df_a.index[0], df_a.index[-1]) + got_b, _ = adapter.get_range("c1", "run-B", df_b.index[0], df_b.index[-1]) + assert got_a is not None + assert len(got_a) == 4 + assert got_b is not None + assert len(got_b) == 8 + finally: + adapter.store.close() + + def test_unscoped_and_scoped_buckets_are_independent(self, tmp_path): + """An unscoped put (``run_id=None``) doesn't satisfy a run-scoped get.""" + adapter = ChannelDataCache(DiskCache(disk_path=tmp_path / "indep")) + try: + _put(adapter, "c1", rows=4) # no run + assert adapter.has_any("c1") + assert not adapter.has_any("c1", "run-A") + data, gaps = adapter.get_range("c1", "run-A", _NOW, _WINDOW_END) + assert data is None + assert gaps == [(_NOW, _WINDOW_END)] + finally: + adapter.store.close() + + def test_get_range_isinstance_filters_foreign_segments(self, tmp_path): + """A segment row with the wrong shape reads as a miss → gap. + + Models a corrupt entry or a key collision from another writer. + ``_load_segment`` isinstance-checks before handing back; the + evicted-segment-as-gap fallback covers this too. + """ + store = DiskCache(disk_path=tmp_path / "foreign") + adapter = ChannelDataCache(store) + try: + _put(adapter, "c1", rows=4) + # Overwrite the segment's payload with foreign data; the + # index still claims it exists, so the read should treat + # the segment range as a gap. + store.put("channel:v1::c1:seg:0", {"not": "an entry"}, size_bytes=64) + data, gaps = adapter.get_range("c1", None, _NOW, _WINDOW_END) + assert data is None + # The whole query range is uncovered (one merged gap). + assert gaps == [(_NOW, _WINDOW_END)] + finally: + store.close() + + def test_invalidate_is_run_scoped(self, tmp_path): + """``invalidate`` only drops the named ``(channel, run)`` bucket.""" + adapter = ChannelDataCache(DiskCache(disk_path=tmp_path / "inval")) + try: + adapter.invalidate("never_added") # safe before any puts + _put(adapter, "c1", rows=4, run_id="run-A") + _put(adapter, "c1", rows=8, run_id="run-B") + adapter.invalidate("c1", "run-A") + assert not adapter.has_any("c1", "run-A") + assert adapter.has_any("c1", "run-B") # run-B survives + finally: + adapter.store.close() + + def test_clear_is_prefix_scoped(self, tmp_path): + """``clear`` drops channel rows across all runs, leaves other adapters alone.""" + store = DiskCache(disk_path=tmp_path / "scoped") + adapter = ChannelDataCache(store) + try: + _put(adapter, "c1", rows=4) # unscoped + _put(adapter, "c2", rows=4, run_id="run-A") + store.put("other:1", "foreign-value", size_bytes=64) + adapter.clear() + assert not adapter.has_any("c1") + assert not adapter.has_any("c2", "run-A") + assert "other:1" in store + finally: + store.close() + + def test_size_bytes_propagates_to_store(self, tmp_path): + """Oversized segments are skipped by the store; index/segment-write order matters. + + The segment is written first, then the index. When the store + refuses the segment (oversize), the index stays empty and + ``has_any`` reports false. Without this ordering you'd get an + index entry pointing at a missing segment. + """ + big = _frame("c1", rows=10_000) + size_bytes = int(big.memory_usage(deep=True).sum()) + store = DiskCache(disk_path=tmp_path / "size", disk_max_bytes=size_bytes // 2) + adapter = ChannelDataCache(store) + try: + adapter.put_segment( + "c1", None, big, big.index[0].to_pydatetime(), big.index[-1].to_pydatetime() + ) + assert not adapter.has_any("c1") + finally: + store.close() + + def test_no_op_store_keeps_adapter_silent(self): + """An adapter on a disabled store behaves like a cold cache.""" + adapter = ChannelDataCache(DiskCache()) + assert not adapter.store.disk_enabled + _put(adapter, "c1", rows=4) + assert not adapter.has_any("c1") + data, gaps = adapter.get_range("c1", None, _NOW, _WINDOW_END) + assert data is None + assert gaps == [(_NOW, _WINDOW_END)] + adapter.invalidate("c1") + adapter.clear() + + # --- stitching + gap behaviour (new with the per-segment shape) --- + + def test_get_range_stitches_multiple_segments(self, tmp_path): + """Two segments whose claimed ranges together cover the query → one stitched frame. + + Segments claim ``[seg_start, seg_end]`` boundaries that abut at + the query midpoint so gap math reports zero gaps. The stitch + path concats and dedups via ``groupby.last``. + """ + adapter = ChannelDataCache(DiskCache(disk_path=tmp_path / "stitch")) + try: + df1 = _frame("c1", rows=5, start=_NOW, freq="ms", offset=0) + df1_end = df1.index[-1].to_pydatetime() + df2 = _frame( + "c1", + rows=5, + start=df1_end + timedelta(milliseconds=1), + freq="ms", + offset=100, + ) + query_start = df1.index[0].to_pydatetime() + query_end = df2.index[-1].to_pydatetime() + # Claim seg2 starts right where seg1 ends so the abutting + # ranges leave no gap. Real callers do this via + # ``_update_cache`` claiming the requested window. + adapter.put_segment("c1", None, df1, query_start, df1_end) + adapter.put_segment("c1", None, df2, df1_end, query_end) + got, gaps = adapter.get_range("c1", None, query_start, query_end) + assert got is not None + assert len(got) == 10 + expected = pd.concat([df1, df2]).groupby(level=0).last() + pd.testing.assert_frame_equal(got.sort_index(), expected.sort_index(), check_freq=False) + assert gaps == [] # abutting claimed ranges → no gap + finally: + adapter.store.close() + + def test_get_range_reports_internal_gap_between_segments(self, tmp_path): + """Query window wider than the two segments → one gap between them.""" + adapter = ChannelDataCache(DiskCache(disk_path=tmp_path / "midgap")) + try: + seg_a_start = _NOW + seg_a_end = _NOW + timedelta(seconds=5) + seg_b_start = _NOW + timedelta(seconds=10) + seg_b_end = _NOW + timedelta(seconds=15) + adapter.put_segment( + "c1", + None, + _frame("c1", rows=2, start=seg_a_start, freq="s"), + seg_a_start, + seg_a_end, + ) + adapter.put_segment( + "c1", + None, + _frame("c1", rows=2, start=seg_b_start, freq="s"), + seg_b_start, + seg_b_end, + ) + _, gaps = adapter.get_range("c1", None, seg_a_start, seg_b_end) + assert gaps == [(seg_a_end, seg_b_start)] + finally: + adapter.store.close() + + def test_get_range_reports_outer_gaps(self, tmp_path): + """Query wider than the cached segment on both sides → two gaps.""" + adapter = ChannelDataCache(DiskCache(disk_path=tmp_path / "outer")) + try: + seg_start = _NOW + timedelta(seconds=5) + seg_end = _NOW + timedelta(seconds=10) + adapter.put_segment( + "c1", + None, + _frame("c1", rows=2, start=seg_start, freq="s"), + seg_start, + seg_end, + ) + query_end = _NOW + timedelta(seconds=15) + _, gaps = adapter.get_range("c1", None, _NOW, query_end) + assert gaps == [(_NOW, seg_start), (seg_end, query_end)] + finally: + adapter.store.close() + + def test_put_segment_normalizes_non_monotonic_index_before_write(self, tmp_path): + """A shuffled-index frame stored and re-read: no ``KeyError``, sorted result. + + Pandas raises ``KeyError`` on value-based partial slicing of a + non-monotonic ``DatetimeIndex``, so :meth:`get_range`'s + ``frame[start:end]`` label slice crashes if a segment ever lands + on disk unsorted. The SDK can't request descending order today, + but the wire could theoretically return interleaved pages; + :meth:`put_segment` sorts on store so the read path is safe + without a per-slice ``try``/``except``. + """ + adapter = ChannelDataCache(DiskCache(disk_path=tmp_path / "unsorted")) + try: + sorted_frame = _frame("c1", rows=5, freq="s") + # Shuffle to a definitely non-monotonic order. + shuffled = sorted_frame.iloc[[2, 0, 4, 1, 3]] + assert not shuffled.index.is_monotonic_increasing + + seg_start = sorted_frame.index[0].to_pydatetime() + seg_end = sorted_frame.index[-1].to_pydatetime() + adapter.put_segment("c1", None, shuffled, seg_start, seg_end) + + data, gaps = adapter.get_range("c1", None, seg_start, seg_end) + assert data is not None + assert data.index.is_monotonic_increasing + pd.testing.assert_frame_equal(data, sorted_frame, check_freq=False) + assert gaps == [] + finally: + adapter.store.close() + + def test_evicted_segment_degrades_to_gap(self, tmp_path): + """If the store loses a segment (LRU), the reader treats its range as a gap. + + Pins the eviction-tolerance contract. The index can outlive + its segments under memory pressure; reads must not error and + must surface the uncovered range so the caller refetches. + """ + store = DiskCache(disk_path=tmp_path / "evict") + adapter = ChannelDataCache(store) + try: + df = _put(adapter, "c1", rows=5) + store.invalidate("channel:v1::c1:seg:0") # simulate eviction + data, gaps = adapter.get_range("c1", None, df.index[0], df.index[-1]) + assert data is None + assert gaps == [(df.index[0], df.index[-1])] + finally: + store.close() + + def test_segment_count_grows_with_each_put(self, tmp_path): + """``put_segment`` writes a new key per call — no merging on the write path. + + Pins the "no rewrite of accumulated bytes per fetch" contract. + ``next_seg_id`` advances; raw store key count grows by one + index update + one new segment per put. + """ + store = DiskCache(disk_path=tmp_path / "grow") + adapter = ChannelDataCache(store) + try: + for i in range(3): + _put( + adapter, + "c1", + rows=2, + start=_NOW + timedelta(seconds=i * 10), + freq="s", + ) + # Expect: 1 index + 3 segments. + channel_keys = sorted(k for k in store if k.startswith("channel:")) + assert channel_keys == [ + "channel:v1::c1:idx", + "channel:v1::c1:seg:0", + "channel:v1::c1:seg:1", + "channel:v1::c1:seg:2", + ] + finally: + store.close() + + def test_empty_put_records_empty_ref_and_skips_refetch(self, tmp_path): + """``put_segment`` with empty data records a ``seg_id=None`` ref. + + Empty ref behavior: + + * ``has_any`` returns True (the bucket has a coverage claim, + just no body). + * ``get_range`` over the claimed range returns ``(None, [])`` + — no data, no gaps, so the caller doesn't refetch. + * No segment key lands in the store (only the index gets a + row); the empty ref lives entirely on the index. + """ + store = DiskCache(disk_path=tmp_path / "empty") + adapter = ChannelDataCache(store) + try: + empty = pd.DataFrame( + {"c1": []}, + index=pd.DatetimeIndex([], tz=timezone.utc), + ) + adapter.put_segment("c1", None, empty, _NOW, _WINDOW_END) + + assert adapter.has_any("c1") + data, gaps = adapter.get_range("c1", None, _NOW, _WINDOW_END) + assert data is None + assert gaps == [] + + channel_keys = sorted(k for k in store if k.startswith("channel:")) + assert channel_keys == ["channel:v1::c1:idx"] + finally: + store.close() + + def test_empty_ref_only_covers_its_claimed_range(self, tmp_path): + """Reads outside the empty ref's claimed range still report gaps. + + Pins the "claim only what you queried" contract: an empty ref + covering ``[T0, T5]`` doesn't suppress a refetch for + ``[T5, T10]``. + """ + adapter = ChannelDataCache(DiskCache(disk_path=tmp_path / "narrow")) + try: + claim_end = _NOW + timedelta(seconds=5) + adapter.put_segment("c1", None, pd.DataFrame(), _NOW, claim_end) + + _, gaps = adapter.get_range("c1", None, _NOW, _WINDOW_END) + # Empty ref covers [_NOW, claim_end]; remainder is a gap. + assert gaps == [(claim_end, _WINDOW_END)] + finally: + adapter.store.close() + + def test_empty_ref_alongside_data_segment(self, tmp_path): + """Stitching data + empty ref returns the data and reports no gap. + + Mirrors a real partial-hit scenario: cache already has data + for ``[T0, T5]``, the caller queries ``[T0, T10]``, the gap + ``[T5, T10]`` fetch returns no rows and lands as an empty + ref. Subsequent ``get_range(T0, T10)`` returns the cached + data with ``gaps == []``. + """ + adapter = ChannelDataCache(DiskCache(disk_path=tmp_path / "mix")) + try: + data = _frame("c1", rows=5, freq="s") + mid = data.index[-1].to_pydatetime() + tail = mid + timedelta(seconds=5) + + adapter.put_segment("c1", None, data, data.index[0].to_pydatetime(), mid) + adapter.put_segment("c1", None, pd.DataFrame(), mid, tail) + + got, gaps = adapter.get_range("c1", None, data.index[0].to_pydatetime(), tail) + assert got is not None + pd.testing.assert_frame_equal(got, data) + assert gaps == [] + finally: + adapter.store.close() + + # --- gap-math unit tests --- + + @pytest.mark.parametrize( + ("covered", "expected"), + [ + # Fully uncovered. + ([], [(_NOW, _NOW + timedelta(seconds=10))]), + # Fully covered. + ([(_NOW, _NOW + timedelta(seconds=10))], []), + # Left-only gap. + ( + [(_NOW + timedelta(seconds=5), _NOW + timedelta(seconds=10))], + [(_NOW, _NOW + timedelta(seconds=5))], + ), + # Right-only gap. + ( + [(_NOW, _NOW + timedelta(seconds=5))], + [(_NOW + timedelta(seconds=5), _NOW + timedelta(seconds=10))], + ), + # Two segments, one internal gap. + ( + [ + (_NOW, _NOW + timedelta(seconds=2)), + (_NOW + timedelta(seconds=7), _NOW + timedelta(seconds=10)), + ], + [(_NOW + timedelta(seconds=2), _NOW + timedelta(seconds=7))], + ), + # Overlapping covered ranges merge before gap math. + ( + [ + (_NOW, _NOW + timedelta(seconds=6)), + (_NOW + timedelta(seconds=4), _NOW + timedelta(seconds=8)), + ], + [(_NOW + timedelta(seconds=8), _NOW + timedelta(seconds=10))], + ), + ], + ids=[ + "fully_uncovered", + "fully_covered", + "left_gap", + "right_gap", + "internal_gap", + "overlapping_covered_merges", + ], + ) + def test_compute_gaps(self, covered, expected): + query_start = _NOW + query_end = _NOW + timedelta(seconds=10) + assert ChannelDataCache._compute_gaps(query_start, query_end, covered) == expected + + +class TestCompaction: + """``ChannelDataCache.MAX_SEGMENTS_PER_BUCKET`` triggers a merge. + + The cap bounds read fan-out: without it, every incremental + ``put_segment`` would add another segment that ``get_range`` would + have to load. Compaction folds the bucket into a single segment + once the count crosses the cap, so the next ``get_range`` only + walks one segment. + + The interesting invariants: + + 1. **Below cap, no merge.** Compaction is a write-time side effect, + not a constant background tax; under the cap the bucket layout + is unchanged. + 2. **At cap+1, fires inline.** ``put_segment`` returns with the + bucket already collapsed — the next reader doesn't pay for a + fragmented bucket. + 3. **Data preserved through merge.** Concrete row count and column + contents survive; LRU-evicted bodies degrade to absence (not + errors) but their refs are still dropped from the index so the + index never points at nothing. + 4. **Old segment keys are deleted.** Otherwise the store carries + the merged seg + every superseded seg until LRU catches up, + which would defeat the bucket's byte-budget accounting. + + A ``monkeypatch`` lowers the cap so tests stay fast — the prod + constant is large enough that a literal 17-fetch test would dwarf + the rest of the suite. + """ + + def test_below_cap_no_compaction(self, tmp_path, monkeypatch): + """Three segments under a cap of four stay separate.""" + monkeypatch.setattr(ChannelDataCache, "MAX_SEGMENTS_PER_BUCKET", 4) + store = DiskCache(disk_path=tmp_path / "below") + adapter = ChannelDataCache(store) + try: + for i in range(3): + _put(adapter, "c1", rows=2, start=_NOW + timedelta(seconds=i * 10)) + + channel_keys = sorted(k for k in store if k.startswith("channel:")) + assert channel_keys == [ + "channel:v1::c1:idx", + "channel:v1::c1:seg:0", + "channel:v1::c1:seg:1", + "channel:v1::c1:seg:2", + ] + finally: + store.close() + + def test_crossing_cap_collapses_to_single_segment(self, tmp_path, monkeypatch): + """One write past the cap triggers an inline compaction. + + Pins all three observable effects in one place: + * Index ends up with exactly one ref. + * The merged segment's ``seg_id`` is fresh (no overlap with + any pre-compaction id). + * Old segment keys are gone from the raw store. + """ + monkeypatch.setattr(ChannelDataCache, "MAX_SEGMENTS_PER_BUCKET", 3) + store = DiskCache(disk_path=tmp_path / "cross") + adapter = ChannelDataCache(store) + try: + for i in range(4): + _put(adapter, "c1", rows=2, start=_NOW + timedelta(seconds=i * 10)) + + idx = adapter._load_index("c1", None) + assert idx is not None + assert len(idx.segments) == 1 + merged_id = idx.segments[0].seg_id + assert merged_id is not None + assert merged_id >= 4 + + channel_keys = sorted(k for k in store if k.startswith("channel:")) + assert channel_keys == [ + "channel:v1::c1:idx", + f"channel:v1::c1:seg:{merged_id}", + ] + finally: + store.close() + + def test_compaction_preserves_data(self, tmp_path, monkeypatch): + """``get_range`` after compaction returns the same rows as before. + + Builds disjoint, time-sorted frames so the dedup path is a + no-op; correctness of the dedup itself is covered by the + overlapping-timestamps test below. + """ + monkeypatch.setattr(ChannelDataCache, "MAX_SEGMENTS_PER_BUCKET", 2) + store = DiskCache(disk_path=tmp_path / "preserve") + adapter = ChannelDataCache(store) + try: + frames = [] + for i in range(3): + frame = _put(adapter, "c1", rows=2, start=_NOW + timedelta(seconds=i * 10)) + frames.append(frame) + + idx = adapter._load_index("c1", None) + assert idx is not None + assert len(idx.segments) == 1 + + expected = pd.concat(frames) + got, gaps = adapter.get_range( + "c1", + None, + frames[0].index[0].to_pydatetime(), + frames[-1].index[-1].to_pydatetime(), + ) + assert gaps == [] + assert got is not None + pd.testing.assert_frame_equal(got, expected) + finally: + store.close() + + def test_compaction_dedups_overlapping_timestamps(self, tmp_path, monkeypatch): + """When two segments share a timestamp, the later put wins. + + Matches :meth:`get_range`'s ``groupby(level=0).last()`` dedup, + so a refetch that overwrites a stale cached value still ends + up with the fresh value after compaction. + """ + monkeypatch.setattr(ChannelDataCache, "MAX_SEGMENTS_PER_BUCKET", 1) + store = DiskCache(disk_path=tmp_path / "dedup") + adapter = ChannelDataCache(store) + try: + ts = _NOW + stale = pd.DataFrame( + {"c1": [1.0]}, + index=pd.DatetimeIndex([ts], tz=timezone.utc), + ) + fresh = pd.DataFrame( + {"c1": [2.0]}, + index=pd.DatetimeIndex([ts], tz=timezone.utc), + ) + adapter.put_segment("c1", None, stale, ts, ts) + adapter.put_segment("c1", None, fresh, ts, ts) + + idx = adapter._load_index("c1", None) + assert idx is not None + assert len(idx.segments) == 1 + + got, _ = adapter.get_range("c1", None, ts, ts) + assert got is not None + assert got.iloc[0, 0] == 2.0 + finally: + store.close() + + def test_empty_refs_fold_into_merged_claim(self, tmp_path, monkeypatch): + """Empty refs alongside data refs collapse into one data ref. + + After compaction the bucket carries a single data segment + whose claimed range spans all prior refs (data + empty). The + empty refs are dropped from the index, their coverage now + represented by the merged ref's wider claim. + """ + monkeypatch.setattr(ChannelDataCache, "MAX_SEGMENTS_PER_BUCKET", 2) + store = DiskCache(disk_path=tmp_path / "fold_empty") + adapter = ChannelDataCache(store) + try: + data = _put(adapter, "c1", rows=4, start=_NOW, freq="s") + data_end = data.index[-1].to_pydatetime() + empty_end = data_end + timedelta(seconds=10) + # Second put (empty) crosses the cap of 2 → compacts. + adapter.put_segment("c1", None, pd.DataFrame(), data_end, empty_end) + adapter.put_segment( + "c1", None, pd.DataFrame(), empty_end, empty_end + timedelta(seconds=10) + ) + + idx = adapter._load_index("c1", None) + assert idx is not None + assert len(idx.segments) == 1 + sole = idx.segments[0] + assert sole.seg_id is not None # data ref survives + assert sole.start_time == data.index[0].to_pydatetime() + assert sole.end_time == empty_end + timedelta(seconds=10) + + # Query the wider claim — data slice returns the actual + # rows, and the empty tail counts as covered (no gap). + got, gaps = adapter.get_range("c1", None, data.index[0].to_pydatetime(), sole.end_time) + assert got is not None + pd.testing.assert_frame_equal(got, data) + assert gaps == [] + finally: + store.close() + + def test_all_empty_bucket_collapses_to_single_empty_ref(self, tmp_path, monkeypatch): + """Compacting a bucket with only empty refs leaves one empty ref. + + The merged claim is the union of every prior empty ref's + claim, so the "no data anywhere in this window" coverage is + preserved while the per-fetch refs collapse. + """ + monkeypatch.setattr(ChannelDataCache, "MAX_SEGMENTS_PER_BUCKET", 2) + store = DiskCache(disk_path=tmp_path / "all_empty") + adapter = ChannelDataCache(store) + try: + t0 = _NOW + t1 = _NOW + timedelta(seconds=10) + t2 = _NOW + timedelta(seconds=20) + t3 = _NOW + timedelta(seconds=30) + for s, e in [(t0, t1), (t1, t2), (t2, t3)]: + adapter.put_segment("c1", None, pd.DataFrame(), s, e) + + idx = adapter._load_index("c1", None) + assert idx is not None + assert len(idx.segments) == 1 + sole = idx.segments[0] + assert sole.seg_id is None + assert sole.start_time == t0 + assert sole.end_time == t3 + + data, gaps = adapter.get_range("c1", None, t0, t3) + assert data is None + assert gaps == [] + + # No segment keys at all — index-only bucket. + seg_keys = [k for k in store if k.startswith("channel:") and ":seg:" in k] + assert seg_keys == [] + finally: + store.close() + + +class TestMergePages: + """Behaviour of :meth:`DataLowLevelClient._merge_pages`. + + The helper replaces a previously inline O(N²) per-page concat loop with + a single batched concat per channel. These tests pin the merge + semantics so a future refactor can't silently drift: + + * Single-frame channels skip the concat entirely (cheap identity path). + * Multi-frame channels concat in collected order; ``groupby.last`` + makes the latest frame win on overlapping timestamps. + * Cached slices from the segment-stitched read are folded in as + the *first* frame so fresh pages still win on overlap. + """ + + @pytest.mark.parametrize("pages", [[], [[]]], ids=["no_tasks_queued", "task_returned_empty"]) + def test_no_fresh_data_returns_initial(self, pages: list) -> None: + """No fresh pages → initial dict passes through by identity.""" + client = DataLowLevelClient(MagicMock()) + initial_df = _frame("chan", rows=5) + with _patch_deserializer({}): + result = client._merge_pages(pages=pages, initial={"chan": initial_df}) + assert result["chan"] is initial_df + + def test_single_frame_skips_concat(self) -> None: + """One frame for a channel → returned by identity, no concat call.""" + only_df = _frame("chan", rows=5) + client = DataLowLevelClient(MagicMock()) + with _patch_deserializer({"p1": {"chan": only_df}}): + result = client._merge_pages(pages=[["p1"]], initial={}) + assert result["chan"] is only_df + + def test_disjoint_pages_concat_in_order(self) -> None: + """Multiple disjoint pages for one channel → single concat result.""" + df1 = _frame("chan", rows=10, start=_NOW, offset=0, freq="s") + df2 = _frame("chan", rows=10, start=_NOW + timedelta(minutes=1), offset=10, freq="s") + df3 = _frame("chan", rows=10, start=_NOW + timedelta(minutes=2), offset=20, freq="s") + client = DataLowLevelClient(MagicMock()) + sentinels = {"p1": {"chan": df1}, "p2": {"chan": df2}, "p3": {"chan": df3}} + with _patch_deserializer(sentinels): + result = client._merge_pages(pages=[["p1", "p2"], ["p3"]], initial={}) + expected = pd.concat([df1, df2, df3]).groupby(level=0).last() + pd.testing.assert_frame_equal(result["chan"].sort_index(), expected.sort_index()) + assert len(result["chan"]) == 30 + + def test_overlapping_timestamps_later_page_wins(self) -> None: + """On overlap, the later page's value survives ``groupby.last``. + + Pins the old inline ``concat([acc, new]).groupby(level=0).last()`` + semantic: latest concat position wins on conflict. + """ + index = pd.date_range(_NOW, periods=5, freq="ms", tz=timezone.utc) + df_first = pd.DataFrame({"chan": [0] * 5}, index=index) + df_second = pd.DataFrame({"chan": [99] * 5}, index=index) + client = DataLowLevelClient(MagicMock()) + with _patch_deserializer({"p1": {"chan": df_first}, "p2": {"chan": df_second}}): + result = client._merge_pages(pages=[["p1", "p2"]], initial={}) + assert (result["chan"]["chan"] == 99).all() + + def test_cached_slice_folded_in_first_and_loses_on_overlap(self) -> None: + """Cached slice from the segment read is the first frame in the merge. + + Fresh pages must overwrite cached values on duplicate timestamps, + matching the pre-existing "latest fetch wins" semantic. + """ + index = pd.date_range(_NOW, periods=5, freq="ms", tz=timezone.utc) + cached = pd.DataFrame({"chan": [-1] * 5}, index=index) + fresh = pd.DataFrame({"chan": [42] * 5}, index=index) + client = DataLowLevelClient(MagicMock()) + with _patch_deserializer({"p1": {"chan": fresh}}): + result = client._merge_pages(pages=[["p1"]], initial={"chan": cached}) + assert (result["chan"]["chan"] == 42).all() + + def test_multiple_channels_independent(self) -> None: + """Per-channel grouping is independent: one channel's pages don't bleed.""" + a1 = _frame("a", rows=5, start=_NOW, offset=0, freq="s") + a2 = _frame("a", rows=5, start=_NOW + timedelta(minutes=1), offset=5, freq="s") + b1 = _frame("b", rows=5, start=_NOW, offset=100, freq="s") + client = DataLowLevelClient(MagicMock()) + sentinels = {"p_a1": {"a": a1}, "p_a2": {"a": a2}, "p_b1": {"b": b1}} + with _patch_deserializer(sentinels): + result = client._merge_pages(pages=[["p_a1", "p_b1"], ["p_a2"]], initial={}) + assert len(result["a"]) == 10 + assert len(result["b"]) == 5 + assert (result["b"]["b"] >= 100).all() + + def test_does_not_mutate_initial(self) -> None: + """``initial`` is a defensive copy; caller's dict isn't mutated.""" + cached = _frame("chan", rows=5) + initial = {"chan": cached} + fresh = _frame("chan", rows=5, start=_NOW + timedelta(seconds=1), offset=10) + client = DataLowLevelClient(MagicMock()) + with _patch_deserializer({"p1": {"chan": fresh}}): + client._merge_pages(pages=[["p1"]], initial=initial) + assert initial["chan"] is cached + + +class TestDataLowLevelClient: + """Constructor wiring and per-instance isolation. + + Per-call behaviour (cache hits, ``ignore_cache``, pagination) lives in + :class:`TestGetChannelData`. + """ + + def test_default_construction_uses_no_op_store(self) -> None: + """Default construction leaves the adapter wrapping a disabled store. + + Resources wire the shared store in via the keyword arg; the + ``MagicMock()``-only path here keeps unit tests free of disk I/O. + """ + client = DataLowLevelClient(MagicMock()) + assert isinstance(client.channel_cache, ChannelDataCache) + assert not client.channel_cache.store.disk_enabled + + def test_per_instance_isolation(self, tmp_path) -> None: + """Two clients with distinct stores must not share cache state. + + Regression test for the original OOM bug: ``channel_cache`` was a + class attribute, so every ``SiftClient`` in the process appended + to the same dict. Two fresh adapters over independent stores must + stay independent — even now that store ownership has moved to the + client, the contract is the same. + """ + client_a = _client_with_cache(tmp_path, "a") + client_b = _client_with_cache(tmp_path, "b") + try: + _put(client_a.channel_cache, "c1", rows=10) + assert client_a.channel_cache.has_any("c1") + assert not client_b.channel_cache.has_any("c1") + finally: + client_a.channel_cache.store.close() + client_b.channel_cache.store.close() + + def test_adapter_kwarg_propagates(self, tmp_path) -> None: + """The constructor honours an externally-constructed adapter. + + Mirrors the production wiring where ``ChannelsAPIAsync`` builds + the adapter from ``client._get_disk_cache()`` and hands it in. + """ + store = DiskCache(disk_path=tmp_path / "external", disk_max_bytes=8_192) + adapter = ChannelDataCache(store) + client = DataLowLevelClient(MagicMock(), channel_cache=adapter) + try: + assert client.channel_cache is adapter + assert client.channel_cache.store is store + assert client.channel_cache.store.disk_max_bytes == 8_192 + finally: + store.close() + + +class TestGetChannelData: + """End-to-end assertions on the public ``get_channel_data`` return shape.""" + + @pytest.mark.asyncio + async def test_single_page_per_channel(self) -> None: + """Result is keyed by channel name; single-page frames pass through unchanged.""" + client = DataLowLevelClient(MagicMock()) + c1_df, c2_df = _frame("c1"), _frame("c2", offset=100) + with _fake_grpc(client, {"c1": [c1_df], "c2": [c2_df]}): + result = await client.get_channel_data( + channels=[_channel("c1"), _channel("c2")], + start_time=_NOW, + end_time=_WINDOW_END, + ignore_cache=True, + ) + assert set(result.keys()) == {"c1", "c2"} + pd.testing.assert_frame_equal(result["c1"], c1_df) + pd.testing.assert_frame_equal(result["c2"], c2_df) + + @pytest.mark.asyncio + async def test_multi_page_response_concatenated_per_channel(self) -> None: + """Three disjoint pages for one channel → single merged frame. + + Catches regressions in the ``_handle_pagination`` + ``_merge_pages`` + interaction (the perf fix's batched concat must still produce the + full 30-row contiguous result). + """ + client = DataLowLevelClient(MagicMock()) + p1 = _frame("c1", rows=10, start=_NOW, offset=0) + p2 = _frame("c1", rows=10, start=_NOW + timedelta(seconds=1), offset=10) + p3 = _frame("c1", rows=10, start=_NOW + timedelta(seconds=2), offset=20) + with _fake_grpc(client, {"c1": [p1, p2, p3]}): + result = await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + ignore_cache=True, + ) + assert set(result.keys()) == {"c1"} + assert len(result["c1"]) == 30 + expected = pd.concat([p1, p2, p3]).groupby(level=0).last() + pd.testing.assert_frame_equal(result["c1"].sort_index(), expected.sort_index()) + + @pytest.mark.asyncio + async def test_cache_hit_short_circuits_grpc(self, tmp_path) -> None: + """Second request for the same channel + window skips ``_get_data_impl``. + + Stages two pages-worth of data so a faulty cache that falls through + wouldn't silently pass by hitting EOF — any second-call invocation + would consume the second page and bump ``len(call_log)``. + """ + client = _client_with_cache(tmp_path) + df = _frame("c1") + try: + with _fake_grpc(client, {"c1": [df, df]}) as call_log: + first = await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + ) + calls_after_first = len(call_log) + assert calls_after_first >= 1 + + second = await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + ) + assert len(call_log) == calls_after_first, ( + "second call should be served from cache without invoking _get_data_impl" + ) + pd.testing.assert_frame_equal(first["c1"].sort_index(), second["c1"].sort_index()) + finally: + client.channel_cache.store.close() + + @pytest.mark.asyncio + async def test_partial_cache_hit_merges_cached_and_fresh(self, tmp_path) -> None: + """Cached + uncached channels resolved together in one return dict. + + Only the uncached channel triggers ``_get_data_impl``. + """ + client = _client_with_cache(tmp_path) + c1_df, c2_df = _frame("c1"), _frame("c2", offset=100) + try: + with _fake_grpc(client, {"c1": [c1_df], "c2": [c2_df]}) as call_log: + await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + ) + calls_after_warmup = len(call_log) + + result = await client.get_channel_data( + channels=[_channel("c1"), _channel("c2")], + start_time=_NOW, + end_time=_WINDOW_END, + ) + new_calls = call_log[calls_after_warmup:] + + assert new_calls, "c2 should hit the wire on the second call" + for call in new_calls: + assert call["channel_ids"] == ["c2"], f"only c2 should hit the wire, saw {call!r}" + assert set(result.keys()) == {"c1", "c2"} + pd.testing.assert_frame_equal(result["c1"].sort_index(), c1_df.sort_index()) + pd.testing.assert_frame_equal(result["c2"].sort_index(), c2_df.sort_index()) + finally: + client.channel_cache.store.close() + + @pytest.mark.asyncio + async def test_run_id_keeps_cache_buckets_separate(self, tmp_path) -> None: + """A cached entry under run-A must NOT satisfy a query for run-B. + + End-to-end version of the run-scope bug fix. Two sequential + ``_fake_grpc`` blocks stage distinct single-page data for each + run; the cache state persists across blocks because it lives on + the adapter. With the bug unfixed (bare ``channel:`` keys), + the run-B query would short-circuit on run-A's cache entry + instead of hitting the wire — the ``len(log_b) >= 1`` assertion + is the canary. + """ + client = _client_with_cache(tmp_path) + df_a = _frame("c1", rows=4, offset=0) + df_b = _frame("c1", rows=8, offset=100) + try: + with _fake_grpc(client, {"c1": [df_a]}) as log_a: + first = await client.get_channel_data( + channels=[_channel("c1")], + run_id="run-A", + start_time=_NOW, + end_time=_WINDOW_END, + ) + assert len(log_a) >= 1, "run-A should hit the wire on first call" + + with _fake_grpc(client, {"c1": [df_b]}) as log_b: + second = await client.get_channel_data( + channels=[_channel("c1")], + run_id="run-B", + start_time=_NOW, + end_time=_WINDOW_END, + ) + assert len(log_b) >= 1, "run-B should still hit the wire even though run-A cached c1" + + pd.testing.assert_frame_equal(first["c1"].sort_index(), df_a.sort_index()) + pd.testing.assert_frame_equal(second["c1"].sort_index(), df_b.sort_index()) + # Both runs end up cached independently; the unscoped bucket + # stays empty because every query named a run. + assert client.channel_cache.has_any("c1", "run-A") + assert client.channel_cache.has_any("c1", "run-B") + assert not client.channel_cache.has_any("c1") + finally: + client.channel_cache.store.close() + + @pytest.mark.asyncio + async def test_pure_cache_hit_does_not_rewrite_disk(self, tmp_path) -> None: + """A full cache hit skips ``_update_cache`` instead of rewriting bytes. + + Mitigates one face of the ``_update_cache`` write amplification: + without the ``had_fresh_data`` gate, repeating the same query + would re-merge and re-pickle the entry back to disk every + call even though nothing changed. Spies on the adapter's + ``put`` to assert zero writes on the second call. + """ + client = _client_with_cache(tmp_path) + df = _frame("c1") + try: + with _fake_grpc(client, {"c1": [df]}): + await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + ) + + put_calls: list[Any] = [] + original_put = client.channel_cache.put_segment + + def spy_put_segment(*args: Any, **kwargs: Any) -> None: + put_calls.append((args, kwargs)) + original_put(*args, **kwargs) + + client.channel_cache.put_segment = spy_put_segment # type: ignore[method-assign] + with _fake_grpc(client, {"c1": []}): + await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + ) + assert put_calls == [], "second call is a full cache hit; no disk write should occur" + finally: + client.channel_cache.store.close() + + @pytest.mark.asyncio + async def test_disjoint_forward_paging_lands_in_separate_segments(self, tmp_path) -> None: + """Two disjoint forward pulls write two segments and stitch on read. + + Under the segment shape, each fetch lands in its own cache key + — no merging on the write path. A subsequent read stitches the + segments back together via ``get_range``. + """ + client = _client_with_cache(tmp_path) + df1 = _frame("c1", rows=5, start=_NOW, offset=0) + df2 = _frame("c1", rows=5, start=_NOW + timedelta(minutes=1), offset=100) + window1_end = _NOW + timedelta(seconds=30) + try: + with _fake_grpc(client, {"c1": [df1]}): + await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=window1_end, + ) + with _fake_grpc(client, {"c1": [df2]}): + await client.get_channel_data( + channels=[_channel("c1")], + start_time=window1_end, + end_time=_WINDOW_END, + ) + + # Two segments under the parent id (no rewrite on the second + # fetch — that's the whole point of the per-segment shape). + idx = client.channel_cache._load_index("c1", None) + assert idx is not None + assert len(idx.segments) == 2 + + stitched, gaps = client.channel_cache.get_range("c1", None, _NOW, _WINDOW_END) + assert stitched is not None + assert len(stitched) == 10 + assert set(stitched.columns) == {"c1"} + # Reads still report the truly-uncovered window between the + # two segments (segments only claim what their data spans + # when run-scoped, but here we used the requested ranges). + assert gaps == [] + finally: + client.channel_cache.store.close() + + @pytest.mark.asyncio + async def test_empty_wire_response_records_empty_ref_and_skips_refetch(self, tmp_path) -> None: + """A query that returns no rows caches an empty ref; the repeat hits cache. + + End-to-end version of the empty-ref contract: a wire call that + returns zero rows lands an empty :class:`SegmentRef` in the + bucket so a repeat of the same query is a full cache hit (no + wire call, no rows). Pins the "known-empty range doesn't + refetch" behavior the segment-cache write amplification fix + depended on for correctness. + """ + client = _client_with_cache(tmp_path) + try: + # First call: zero pages → empty wire response. + with _fake_grpc(client, {"c1": []}) as log1: + result1 = await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + ) + assert log1, "first call should hit the wire" + assert "c1" not in result1 + + # Bucket carries an empty ref for the queried window. + assert client.channel_cache.has_any("c1") + idx = client.channel_cache._load_index("c1", None) + assert idx is not None + assert len(idx.segments) == 1 + assert idx.segments[0].seg_id is None + assert idx.segments[0].start_time == _NOW + assert idx.segments[0].end_time == _WINDOW_END + + # Second identical call: cache hit, no wire traffic. + with _fake_grpc(client, {"c1": []}) as log2: + result2 = await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + ) + assert log2 == [], "second call should be a full cache hit" + assert "c1" not in result2 + finally: + client.channel_cache.store.close() + + @pytest.mark.asyncio + async def test_run_scoped_empty_wire_response_does_not_mask_future_data(self, tmp_path) -> None: + """Run-scoped empty responses must not cache; later ingest still shows up. + + Absence at query time doesn't imply future absence for a run + that may still be ingesting. Caching an empty ref over the + queried window would permanently mask data the run writes + later — the same window becomes a full cache hit (gaps == + ``[]``) and never refetches. This pins the run-scoped + exception to the empty-ref rule. + """ + client = _client_with_cache(tmp_path) + df = _frame("c1", rows=4) + try: + with _fake_grpc(client, {"c1": []}) as log1: + result1 = await client.get_channel_data( + channels=[_channel("c1")], + run_id="run-A", + start_time=_NOW, + end_time=_WINDOW_END, + ) + assert log1, "first (empty) call should hit the wire" + assert "c1" not in result1 + assert not client.channel_cache.has_any("c1", "run-A"), ( + "run-scoped empty response must not land an empty ref " + "(would mask data the run ingests later)" + ) + + with _fake_grpc(client, {"c1": [df]}) as log2: + result2 = await client.get_channel_data( + channels=[_channel("c1")], + run_id="run-A", + start_time=_NOW, + end_time=_WINDOW_END, + ) + assert log2, ( + "second call must refetch — the previous empty response " + "cannot become a cache hit for an ongoing run" + ) + pd.testing.assert_frame_equal(result2["c1"].sort_index(), df.sort_index()) + finally: + client.channel_cache.store.close() + + @pytest.mark.asyncio + async def test_ignore_cache_true_returns_fresh_and_skips_write(self, tmp_path) -> None: + """``ignore_cache=True`` returns mock data and leaves the cache empty. + + End-to-end version of the latent bug that compounded the customer's + OOM: pre-fix, ``_update_cache`` ran even when the caller had asked + the cache to be ignored. + """ + client = _client_with_cache(tmp_path) + df = _frame("c1") + try: + with _fake_grpc(client, {"c1": [df]}): + result = await client.get_channel_data( + channels=[_channel("c1")], + start_time=_NOW, + end_time=_WINDOW_END, + ignore_cache=True, + ) + pd.testing.assert_frame_equal(result["c1"], df) + assert not client.channel_cache.has_any("c1") + finally: + client.channel_cache.store.close() + + +class TestBitFieldChannels: + """Bitfield channels exercise the channel-id-vs-channel-name seam. + + A bitfield channel has one parent ``channel_id`` and N elements that + surface in result dicts as dotted names (``parent.element``). The + data path uses ids for the cache key (one entry per channel) but + names for the result dict (one frame per element), so any test that + only covers single-channel non-bitfield flows misses the join. + + These tests pin three things: + + 1. ``_update_name_id_map`` populates an entry for the parent *and* + every element, all pointing at the parent id — so when + ``_update_cache`` looks up a dotted name, it lands on the + parent's cache row. + 2. Element frames merge into a single wide-DataFrame segment + under the parent id, and ``get_range`` correctly treats the + bitfield as cached as a unit (zero gaps for the warmed window). + 3. Fresh-fetch and cache-hit return shapes match — both produce + one single-column frame per dotted element name. The cache-hit + branch slices the wide cached frame per column to keep this + symmetry; without that slice, a cached bitfield would hand the + full wide frame back under every dotted key. + """ + + def test_update_name_id_map_populates_dotted_and_parent_names(self) -> None: + """Parent name and every dotted element name map to the parent id. + + This is the contract :meth:`_update_cache` relies on when it + looks up ``name_id_map.get(channel_name)`` for the dotted names + returned by ``try_deserialize_channel_data``. + """ + client = DataLowLevelClient(MagicMock()) + ch = _bitfield_channel(cid="abc", name="ch1", elements=["lo", "hi"]) + client._update_name_id_map([ch]) + assert client.channel_cache.name_id_map == { + "ch1": "abc", + "ch1.lo": "abc", + "ch1.hi": "abc", + } + + def test_update_cache_merges_elements_into_one_segment(self, tmp_path) -> None: + """All element frames land in one segment keyed by the parent id. + + ``_update_cache`` groups frames by their parent id before + writing, so bitfield elements that came in one wire fetch + produce exactly one wide segment — not one segment per element. + Nothing should land under a dotted-name key. + """ + client = _client_with_cache(tmp_path) + ch = _bitfield_channel(cid="abc", name="ch1", elements=["lo", "hi"]) + client._update_name_id_map([ch]) + df_lo = _frame("ch1.lo", rows=5) + df_hi = _frame("ch1.hi", rows=5) + try: + client._update_cache( + channel_data={"ch1.lo": df_lo, "ch1.hi": df_hi}, + fetched_ranges_per_channel={"abc": [(_NOW, _WINDOW_END)]}, + start_time=_NOW, + end_time=_WINDOW_END, + ) + assert client.channel_cache.has_any("abc") + assert not client.channel_cache.has_any("ch1.lo") + assert not client.channel_cache.has_any("ch1.hi") + + # Exactly one segment under the parent id holding both + # element columns. ``get_range`` returns the wide frame; + # downstream slicing happens in ``get_channel_data``. + idx = client.channel_cache._load_index("abc", None) + assert idx is not None + assert len(idx.segments) == 1 + + data, _ = client.channel_cache.get_range("abc", None, _NOW, _WINDOW_END) + assert data is not None + assert set(data.columns) == {"ch1.lo", "ch1.hi"} + assert len(data) == 5 + finally: + client.channel_cache.store.close() + + def test_get_range_treats_bitfield_as_one_bucket(self, tmp_path) -> None: + """A bitfield warmed by a prior fetch reports zero gaps for the same range. + + Pins the read side of the id/name asymmetry: a bitfield warmed + by a prior fetch is "cached" as a unit under the parent id. + Without this, a second call would re-fetch the whole window + even though the elements are in the cache. + """ + client = _client_with_cache(tmp_path) + ch = _bitfield_channel(cid="abc", name="ch1", elements=["lo", "hi"]) + client._update_name_id_map([ch]) + try: + client._update_cache( + channel_data={ + "ch1.lo": _frame("ch1.lo", rows=5), + "ch1.hi": _frame("ch1.hi", rows=5), + }, + fetched_ranges_per_channel={"abc": [(_NOW, _WINDOW_END)]}, + start_time=_NOW, + end_time=_WINDOW_END, + ) + data, gaps = client.channel_cache.get_range("abc", None, _NOW, _WINDOW_END) + assert data is not None + assert gaps == [] # parent id covers the whole window + finally: + client.channel_cache.store.close() + + @pytest.mark.asyncio + async def test_get_channel_data_fresh_bitfield_returns_per_element_frames( + self, tmp_path + ) -> None: + """Fresh fetch: result has dotted keys, each a single-column frame. + + Mirrors ``try_deserialize_channel_data``'s output shape — one + DataFrame per bitfield element, named after the element column. + The cache is left warmed under the parent id so a follow-up + test can compare the cache-hit shape against this baseline. + """ + client = _client_with_cache(tmp_path) + ch = _bitfield_channel(cid="abc", name="ch1", elements=["lo", "hi"]) + df_lo = _frame("ch1.lo", rows=5) + df_hi = _frame("ch1.hi", rows=5) + try: + with _fake_grpc(client, {"abc": [{"ch1.lo": df_lo, "ch1.hi": df_hi}]}): + result = await client.get_channel_data( + channels=[ch], + start_time=_NOW, + end_time=_WINDOW_END, + ) + assert set(result.keys()) == {"ch1.lo", "ch1.hi"} + assert list(result["ch1.lo"].columns) == ["ch1.lo"] + assert list(result["ch1.hi"].columns) == ["ch1.hi"] + pd.testing.assert_frame_equal(result["ch1.lo"].sort_index(), df_lo.sort_index()) + pd.testing.assert_frame_equal(result["ch1.hi"].sort_index(), df_hi.sort_index()) + # Cache holds one wide segment under the parent id. + data, _ = client.channel_cache.get_range("abc", None, _NOW, _WINDOW_END) + assert data is not None + assert set(data.columns) == {"ch1.lo", "ch1.hi"} + finally: + client.channel_cache.store.close() + + @pytest.mark.asyncio + async def test_get_channel_data_cached_bitfield_returns_per_element_frames( + self, tmp_path + ) -> None: + """Cache-hit shape matches fresh-fetch shape (per-element single-column). + + Same query as :func:`test_get_channel_data_fresh_bitfield_returns_per_element_frames` + run twice — the second call must NOT hit the wire (full cache + hit) AND must return the same per-element single-column shape. + Regression guard for the cache-hit branch's ``cached_data[[name]]`` + slice — without it, a cached bitfield would return the wide + cached frame under every dotted key. + """ + client = _client_with_cache(tmp_path) + ch = _bitfield_channel(cid="abc", name="ch1", elements=["lo", "hi"]) + df_lo = _frame("ch1.lo", rows=5) + df_hi = _frame("ch1.hi", rows=5) + try: + with _fake_grpc(client, {"abc": [{"ch1.lo": df_lo, "ch1.hi": df_hi}]}): + await client.get_channel_data( + channels=[ch], + start_time=_NOW, + end_time=_WINDOW_END, + ) + with _fake_grpc(client, {"abc": []}) as log: + cached = await client.get_channel_data( + channels=[ch], + start_time=_NOW, + end_time=_WINDOW_END, + ) + assert log == [], "second call must be a full cache hit (no wire calls)" + assert set(cached.keys()) == {"ch1.lo", "ch1.hi"} + assert list(cached["ch1.lo"].columns) == ["ch1.lo"] + assert list(cached["ch1.hi"].columns) == ["ch1.hi"] + # ``check_freq=False``: the cached frame loses its DatetimeIndex + # ``freq`` attribute through the concat/groupby roundtrip in + # ``_update_cache``. Values and timestamps still match — only + # the (non-semantic) freq metadata diverges. + pd.testing.assert_frame_equal( + cached["ch1.lo"].sort_index(), df_lo.sort_index(), check_freq=False + ) + pd.testing.assert_frame_equal( + cached["ch1.hi"].sort_index(), df_hi.sort_index(), check_freq=False + ) + finally: + client.channel_cache.store.close() diff --git a/python/lib/sift_client/_tests/_internal/test_disk_cache.py b/python/lib/sift_client/_tests/_internal/test_disk_cache.py new file mode 100644 index 000000000..29bcd6546 --- /dev/null +++ b/python/lib/sift_client/_tests/_internal/test_disk_cache.py @@ -0,0 +1,413 @@ +"""Tests for :mod:`sift_client._internal.disk_cache`. + +Two classes, narrowest scope first: + +* :class:`TestDiskCache` — direct unit tests on :class:`DiskCache`: + the disabled-when-no-path no-op, fresh writes/reads, cross-session + reload, oversize guard + dedup keyed on the full namespaced key, and + the marker-guarded :meth:`DiskCache.clear_disk` classmethod. +* :class:`TestClearDisk` — the classmethod's defensive guards. + +The store is intentionally key/value-agnostic — every test treats it as +a plain ``str``-keyed dict that happens to persist across handles, with +``size_bytes`` supplied by the caller. The channel-specific adapter +(:class:`ChannelDataCache`) is exercised separately in ``test_data.py``. +""" + +from __future__ import annotations + +import logging +from contextlib import contextmanager +from typing import Iterator +from unittest.mock import patch + +import diskcache +import pytest + +from sift_client._internal.disk_cache import DiskCache + +# Snapshot of the production constant captured at import time. The autouse +# ``_isolate_default_disk_cache_path`` fixture in ``conftest.py`` overrides +# the class attribute per test; the constant-shape test still needs the +# real value to assert against. +_PRODUCTION_DEFAULT_DISK_PATH = DiskCache.DEFAULT_DISK_PATH + + +@contextmanager +def _capture_disk_cache_warnings() -> Iterator[list[logging.LogRecord]]: + """Capture warnings emitted by the disk-cache logger directly. + + Pytest's ``caplog`` reads from the root logger, but the Sift pytest + plugin sets ``propagate=False`` on the ``sift_client`` logger when + audit logging is active, so records emitted from any descendant don't + reach the root. Attaching a list-backed handler at the leaf logger + bypasses that. + """ + target = logging.getLogger("sift_client._internal.disk_cache") + records: list[logging.LogRecord] = [] + + class _ListHandler(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + records.append(record) + + handler = _ListHandler(level=logging.WARNING) + target.addHandler(handler) + try: + yield records + finally: + target.removeHandler(handler) + + +class TestDiskCache: + """End-to-end behaviour of the shared on-disk store.""" + + def test_disabled_when_no_path(self) -> None: + """``DiskCache()`` with no ``disk_path`` is a silent no-op.""" + cache = DiskCache() + assert cache.disk_enabled is False + assert cache.disk_path is None + assert cache.disk_max_bytes is None + # Every operation no-ops; no AttributeError, no warning. + cache.put("k", "v", size_bytes=4) + assert "k" not in cache + assert cache.get("k") is None + assert list(iter(cache)) == [] + cache.invalidate("k") + cache.clear() + cache.close() + + def test_fresh_cache_writes_and_reads(self, tmp_path) -> None: + """A fresh disk directory accepts writes and serves them back.""" + cache = DiskCache(disk_path=tmp_path / "fresh") + try: + assert cache.disk_enabled + assert cache.disk_path == str(tmp_path / "fresh") + assert cache.disk_max_bytes == DiskCache.DEFAULT_DISK_MAX_BYTES + cache.put("k", {"hello": "world"}, size_bytes=64) + assert "k" in cache + assert cache.get("k") == {"hello": "world"} + finally: + cache.close() + + def test_reopen_existing_dir_sees_prior_session_entries(self, tmp_path) -> None: + """Closing then reopening at the same path surfaces prior entries. + + This is the cold-start reuse guarantee: a fresh process pointing + at a populated directory reads back what an earlier process wrote. + """ + path = tmp_path / "prev-session" + session1 = DiskCache(disk_path=path) + session1.put("k", [1, 2, 3], size_bytes=24) + session1.close() + + session2 = DiskCache(disk_path=path) + try: + assert "k" in session2 + assert session2.get("k") == [1, 2, 3] + finally: + session2.close() + + def test_repeated_put_overwrites(self, tmp_path) -> None: + cache = DiskCache(disk_path=tmp_path / "overwrite") + try: + cache.put("k", "first", size_bytes=8) + cache.put("k", "second", size_bytes=8) + assert cache.get("k") == "second" + finally: + cache.close() + + def test_invalidate_removes_entry(self, tmp_path) -> None: + cache = DiskCache(disk_path=tmp_path / "inval") + try: + cache.invalidate("never_added") # safe before any puts + cache.put("k", "v", size_bytes=4) + cache.invalidate("k") + assert "k" not in cache + assert cache.get("k") is None + finally: + cache.close() + + def test_clear_wipes_store(self, tmp_path) -> None: + cache = DiskCache(disk_path=tmp_path / "clear") + try: + cache.put("a", 1, size_bytes=8) + cache.put("b", 2, size_bytes=8) + cache.clear() + assert "a" not in cache + assert "b" not in cache + finally: + cache.close() + + def test_iter_yields_keys(self, tmp_path) -> None: + """``__iter__`` exposes the keyspace so adapters can prefix-clear.""" + cache = DiskCache(disk_path=tmp_path / "iter") + try: + cache.put("alpha:1", 1, size_bytes=8) + cache.put("beta:1", 2, size_bytes=8) + cache.put("alpha:2", 3, size_bytes=8) + assert set(cache) == {"alpha:1", "alpha:2", "beta:1"} + finally: + cache.close() + + def test_disable_closes_handle(self, tmp_path) -> None: + """Turning off disk closes the handle and silences subsequent ops.""" + cache = DiskCache(disk_path=tmp_path / "disable") + try: + cache.put("k", "v", size_bytes=4) + cache.disable() + assert not cache.disk_enabled + assert cache.disk_path is None + assert "k" not in cache + assert cache.get("k") is None + cache.put("new", "x", size_bytes=4) # silently dropped + assert "new" not in cache + finally: + cache.close() + + def test_del_releases_handle_when_close_not_called(self, tmp_path) -> None: + """Dropping a :class:`DiskCache` without :meth:`close` still releases its handle. + + Guards against the specific leak pattern of transient callers + (short-lived :class:`SiftClient` instances in a service loop): + each unclosed cache holds a SQLite handle against the shared + directory. ``__del__`` calls into ``diskcache.Cache.close`` on + GC so the handle is returned even when the owner forgets. + Verified via ``mock.patch.object`` on ``diskcache.Cache.close`` + so we're pinning the method call (not just an observable + side-effect that could pass for other reasons). + """ + cache = DiskCache(disk_path=tmp_path / "gc") + cache.put("k", "v", size_bytes=4) + with patch.object(diskcache.Cache, "close", autospec=True) as mock_close: + # ``del`` on the sole reference triggers immediate collection + # under CPython's refcount; the safety net calls ``close``. + del cache + assert mock_close.called, ( + "__del__ must call diskcache.Cache.close so the SQLite " + "handle isn't leaked when the owner is GC'd without an " + "explicit close()" + ) + + def test_reopen_preserves_persisted_size_limit(self, tmp_path) -> None: + """Reopening a cache without ``max_bytes`` honors its persisted cap. + + The cache directory is shared across sessions, but historically + every open re-asserted the caller's ``max_bytes`` (defaulting + to 4 GiB), so one client's 4 GiB default would silently resize + another client's 2 GiB cache the next time they opened the + same path. ``diskcache`` already persists ``size_limit`` in its + Settings table; :meth:`_open_disk` now leans on that unless the + caller passes an explicit override. + """ + path = tmp_path / "sticky-cap" + custom_cap = 2 * 1024 * 1024 * 1024 # 2 GiB + first = DiskCache(disk_path=path, disk_max_bytes=custom_cap) + try: + assert first.disk_max_bytes == custom_cap + finally: + first.close() + + second = DiskCache(disk_path=path) + try: + assert second.disk_max_bytes == custom_cap, ( + "reopen without max_bytes must reuse the persisted cap, " + "not silently reset to DEFAULT_DISK_MAX_BYTES" + ) + finally: + second.close() + + # Explicit override on a subsequent open still wins over the + # persisted value — otherwise callers couldn't ever shrink or + # grow an existing cache. + new_cap = custom_cap // 2 + third = DiskCache(disk_path=path, disk_max_bytes=new_cap) + try: + assert third.disk_max_bytes == new_cap + finally: + third.close() + + def test_fresh_cache_seeds_default_size_limit(self, tmp_path) -> None: + """A brand-new cache directory picks up ``DEFAULT_DISK_MAX_BYTES``.""" + cache = DiskCache(disk_path=tmp_path / "fresh") + try: + assert cache.disk_max_bytes == DiskCache.DEFAULT_DISK_MAX_BYTES + finally: + cache.close() + + def test_enable_reconfigures_path(self, tmp_path) -> None: + """Reconfiguring to a different path closes the old handle. + + The new directory starts empty: ``k`` lived in the old directory + so the lookup at the new path misses. + """ + cache = DiskCache(disk_path=tmp_path / "a") + try: + cache.put("k", "v", size_bytes=4) + cache.enable(path=tmp_path / "b") + assert cache.disk_path == str(tmp_path / "b") + assert "k" not in cache + finally: + cache.close() + + def test_enable_noop_when_same_settings(self, tmp_path) -> None: + """Re-enabling with identical settings doesn't churn the disk handle.""" + cache = DiskCache(disk_path=tmp_path / "noop") + try: + handle_before = cache._disk + cache.enable(path=tmp_path / "noop", max_bytes=DiskCache.DEFAULT_DISK_MAX_BYTES) + assert cache._disk is handle_before + finally: + cache.close() + + def test_oversized_entry_skipped_and_preserves_neighbours(self, tmp_path) -> None: + """An entry larger than the cap is skipped without evicting peers. + + Without this guard, ``diskcache``'s cull would evict every other + row trying to fit an unfittable entry, then drop the entry itself + — the wipe-everything failure mode the cache work originally fixed. + + Cap is sized to leave plenty of room for diskcache's pickle + envelope around the small entries while still being small enough + that the declared oversized ``size_bytes`` (10 MB) trips the + guard. ``size_bytes`` is the caller's contract — the store + compares that, not the actual on-disk size. + """ + cap = 1 * 1024 * 1024 # 1 MiB + cache = DiskCache(disk_path=tmp_path / "oversize", disk_max_bytes=cap) + try: + cache.put("small-1", "value", size_bytes=64) + cache.put("small-2", "value", size_bytes=64) + with _capture_disk_cache_warnings() as records: + cache.put("huge", "value", size_bytes=10 * 1024 * 1024) + assert "small-1" in cache + assert "small-2" in cache + assert "huge" not in cache + assert any("larger than the disk cache cap" in r.getMessage() for r in records) + finally: + cache.close() + + def test_oversized_put_drops_prior_entry(self, tmp_path) -> None: + """An oversized re-insert must drop the prior value, not silently keep it.""" + cap = 1 * 1024 * 1024 + cache = DiskCache(disk_path=tmp_path / "drop-prior", disk_max_bytes=cap) + try: + cache.put("k", "small", size_bytes=64) + assert "k" in cache + cache.put("k", "big", size_bytes=10 * 1024 * 1024) + assert "k" not in cache + finally: + cache.close() + + def test_oversized_put_warns_once_per_key(self, tmp_path) -> None: + """Repeated oversized puts for the same key log once, not every call.""" + cap = 1 * 1024 * 1024 + cache = DiskCache(disk_path=tmp_path / "dedup", disk_max_bytes=cap) + try: + with _capture_disk_cache_warnings() as records: + for _ in range(5): + cache.put("k", "v", size_bytes=10 * 1024 * 1024) + warnings = [r for r in records if "larger than the disk cache cap" in r.getMessage()] + assert len(warnings) == 1 + finally: + cache.close() + + def test_oversized_warning_resets_after_normal_put(self, tmp_path) -> None: + """A successful normal-sized put clears the dedup bit for that key.""" + cap = 1 * 1024 * 1024 + cache = DiskCache(disk_path=tmp_path / "reset-normal", disk_max_bytes=cap) + try: + with _capture_disk_cache_warnings() as records: + cache.put("k", "v", size_bytes=10 * 1024 * 1024) # 1st warning + cache.put("k", "v", size_bytes=64) # resets state + cache.put("k", "v", size_bytes=10 * 1024 * 1024) # 2nd warning + warnings = [r for r in records if "larger than the disk cache cap" in r.getMessage()] + assert len(warnings) == 2 + finally: + cache.close() + + def test_dedup_keys_on_full_namespaced_key(self, tmp_path) -> None: + """Dedup is per-key, so two adapters' colliding bare ids don't share state. + + Pins the design choice that the oversize warning dedup tracks the + full namespaced key handed to ``put`` (e.g. ``channel:foo`` vs + ``calc:foo``) rather than collapsing on the bare id. Two different + prefixes for the same suffix each get their own one-shot warning. + """ + cap = 1 * 1024 * 1024 + cache = DiskCache(disk_path=tmp_path / "two-prefixes", disk_max_bytes=cap) + try: + with _capture_disk_cache_warnings() as records: + cache.put("alpha:foo", "v", size_bytes=10 * 1024 * 1024) + cache.put("beta:foo", "v", size_bytes=10 * 1024 * 1024) + warnings = [r for r in records if "larger than the disk cache cap" in r.getMessage()] + assert len(warnings) == 2 + messages = [r.getMessage() for r in warnings] + assert any("alpha:foo" in m for m in messages) + assert any("beta:foo" in m for m in messages) + finally: + cache.close() + + def test_invalidate_resets_oversized_warning(self, tmp_path) -> None: + cap = 1 * 1024 * 1024 + cache = DiskCache(disk_path=tmp_path / "reset-inval", disk_max_bytes=cap) + try: + with _capture_disk_cache_warnings() as records: + cache.put("k", "v", size_bytes=10 * 1024 * 1024) + cache.invalidate("k") + cache.put("k", "v", size_bytes=10 * 1024 * 1024) + warnings = [r for r in records if "larger than the disk cache cap" in r.getMessage()] + assert len(warnings) == 2 + finally: + cache.close() + + def test_clear_resets_oversized_warning(self, tmp_path) -> None: + cap = 1 * 1024 * 1024 + cache = DiskCache(disk_path=tmp_path / "reset-clear", disk_max_bytes=cap) + try: + with _capture_disk_cache_warnings() as records: + cache.put("a", "v", size_bytes=10 * 1024 * 1024) + cache.put("b", "v", size_bytes=10 * 1024 * 1024) + cache.clear() + cache.put("a", "v", size_bytes=10 * 1024 * 1024) + cache.put("b", "v", size_bytes=10 * 1024 * 1024) + warnings = [r for r in records if "larger than the disk cache cap" in r.getMessage()] + assert len(warnings) == 4 + finally: + cache.close() + + +class TestClearDisk: + """:meth:`DiskCache.clear_disk` removes a cache dir, refuses other dirs.""" + + def test_clear_removes_directory(self, tmp_path) -> None: + path = tmp_path / "victim" + cache = DiskCache(disk_path=path) + cache.put("k", "v", size_bytes=4) + cache.close() + assert path.exists() + DiskCache.clear_disk(path) + assert not path.exists() + + def test_clear_missing_path_is_noop(self, tmp_path) -> None: + DiskCache.clear_disk(tmp_path / "never-existed") # no raise + + def test_clear_refuses_non_diskcache_directory(self, tmp_path) -> None: + """A typo'd path with unrelated contents must not be wiped.""" + target = tmp_path / "user-stuff" + target.mkdir() + (target / "important.txt").write_text("don't delete me") + with pytest.raises(ValueError, match="does not look like a sift data cache"): + DiskCache.clear_disk(target) + assert (target / "important.txt").read_text() == "don't delete me" + + def test_default_path_constant_under_tmp(self) -> None: + """Default lives under the OS tmp dir, not a user directory. + + Reads the module-level snapshot rather than ``DEFAULT_DISK_PATH`` + directly because the autouse fixture monkeypatches that attribute + for every test. + """ + import tempfile + + assert _PRODUCTION_DEFAULT_DISK_PATH.startswith(tempfile.gettempdir()) + assert _PRODUCTION_DEFAULT_DISK_PATH.endswith("sift-data-cache") diff --git a/python/lib/sift_client/_tests/_internal/test_disk_cache_config.py b/python/lib/sift_client/_tests/_internal/test_disk_cache_config.py new file mode 100644 index 000000000..bce8a4ab9 --- /dev/null +++ b/python/lib/sift_client/_tests/_internal/test_disk_cache_config.py @@ -0,0 +1,112 @@ +"""Tests for :class:`sift_client._internal.disk_cache_config.DiskCacheConfig`. + +The class is a small intent holder; the tests pin three things that +resource lazy-init code relies on: + +* Enable / disable round-trips preserve the right state and clear overrides. +* ``using_default_path`` reflects "enabled AND no user override", which + drives the silent-fallback-vs-loud-raise distinction in resources. +* ``enable`` accepts ``os.PathLike`` and stringifies it eagerly so consumers + never need to handle ``pathlib.Path`` vs ``str``. +""" + +from __future__ import annotations + +import pathlib + +import pytest + +from sift_client._internal.disk_cache_config import DiskCacheConfig + + +class TestDiskCacheConfig: + def test_opt_out_initial_state_enabled_no_overrides(self) -> None: + """``enabled=True`` (opt-out) starts on with no overrides.""" + config = DiskCacheConfig(enabled=True) + assert config.enabled + assert config.path is None + assert config.max_bytes is None + assert config.using_default_path + + def test_opt_in_initial_state_disabled(self) -> None: + """``enabled=False`` (opt-in) starts off; ``using_default_path`` is False.""" + config = DiskCacheConfig(enabled=False) + assert not config.enabled + assert config.path is None + assert config.max_bytes is None + assert not config.using_default_path + + def test_enable_with_no_args_keeps_defaults(self) -> None: + """``enable()`` with no args turns on and clears any prior overrides.""" + config = DiskCacheConfig(enabled=False) + config.enable() + assert config.enabled + assert config.path is None + assert config.max_bytes is None + assert config.using_default_path + + def test_enable_with_path_marks_non_default(self) -> None: + """A user-supplied path flips ``using_default_path`` off.""" + config = DiskCacheConfig(enabled=True) + config.enable(path="/custom/path") + assert config.enabled + assert config.path == "/custom/path" + assert not config.using_default_path + + def test_enable_with_max_bytes_keeps_default_path(self) -> None: + """Setting ``max_bytes`` alone doesn't make the path non-default.""" + config = DiskCacheConfig(enabled=True) + config.enable(max_bytes=1024) + assert config.enabled + assert config.path is None + assert config.max_bytes == 1024 + assert config.using_default_path + + def test_enable_stringifies_pathlike(self) -> None: + """``os.PathLike`` inputs are stored as strings so consumers can be dumb.""" + config = DiskCacheConfig(enabled=True) + config.enable(path=pathlib.Path("/some/path")) + assert isinstance(config.path, str) + assert config.path == "/some/path" + + def test_disable_clears_overrides(self) -> None: + """``disable()`` zeroes path and max_bytes so a future re-enable starts clean.""" + config = DiskCacheConfig(enabled=True) + config.enable(path="/custom", max_bytes=4096) + config.disable() + assert not config.enabled + assert config.path is None + assert config.max_bytes is None + assert not config.using_default_path + + def test_reenable_after_disable_returns_to_defaults(self) -> None: + """``disable`` then ``enable()`` (no args) restores the opt-out starting state.""" + config = DiskCacheConfig(enabled=True) + config.enable(path="/custom", max_bytes=4096) + config.disable() + config.enable() + assert config.enabled + assert config.path is None + assert config.max_bytes is None + assert config.using_default_path + + @pytest.mark.parametrize( + ("enabled", "path", "expected"), + [ + (True, None, True), + (True, "/custom", False), + (False, None, False), + (False, "/custom", False), # disabled wins even with a stashed path + ], + ids=["enabled+default", "enabled+custom", "disabled+default", "disabled+custom"], + ) + def test_using_default_path_matrix( + self, enabled: bool, path: str | None, expected: bool + ) -> None: + """``using_default_path`` is the AND of ``enabled`` and ``path is None``.""" + config = DiskCacheConfig(enabled=enabled) + if path is not None: + # Bypass enable() so we can exercise the disabled+custom combo + # without enable() flipping enabled back on. + config._path = path + assert config.using_default_path is expected diff --git a/python/lib/sift_client/_tests/conftest.py b/python/lib/sift_client/_tests/conftest.py index 41469dac5..5790a2f7a 100644 --- a/python/lib/sift_client/_tests/conftest.py +++ b/python/lib/sift_client/_tests/conftest.py @@ -9,6 +9,30 @@ from sift_client.util.util import AsyncAPIs +@pytest.fixture(autouse=True) +def _isolate_default_disk_cache_path(monkeypatch, tmp_path): + """Redirect ``DiskCache.DEFAULT_DISK_PATH`` to a per-test tmp dir. + + On-disk caching is **opt-out** — any test that triggers the lazy + ``DiskCache`` init through ``SiftClient._get_disk_cache`` would + otherwise create the real ``/tmp/sift-data-cache`` directory and leak + state across runs. Redirecting the default to ``tmp_path`` keeps every + test self-contained without each test having to know the cache is on + by default. + + The override preserves the ``sift-data-cache`` suffix so + ``TestClearDisk::test_default_path_constant_under_tmp`` keeps + validating the real shape of the constant. + """ + from sift_client._internal.disk_cache import DiskCache + + monkeypatch.setattr( + DiskCache, + "DEFAULT_DISK_PATH", + str(tmp_path / "sift-data-cache"), + ) + + @pytest.fixture(scope="session") def sift_client() -> SiftClient: """Create a SiftClient instance for testing. diff --git a/python/lib/sift_client/_tests/test_client_cache.py b/python/lib/sift_client/_tests/test_client_cache.py new file mode 100644 index 000000000..d1de3f866 --- /dev/null +++ b/python/lib/sift_client/_tests/test_client_cache.py @@ -0,0 +1,388 @@ +"""Tests for :mod:`sift_client.util.cache`. + +The namespace is the user-facing surface for the shared on-disk store +that lives on the :class:`SiftClient`. Three concerns get pinned here: + +1. Default policy (opt-out: caching on at the default path) lands on + the live store on first use. +2. Pre-init configuration (``client.cache.disable()`` / + ``enable(path=..., max_bytes=...)`` before any resource has touched + the cache) takes effect on the lazy build. +3. Post-init reconfiguration mutates the live :class:`DiskCache` in + place rather than swapping it out — every resource adapter holds a + reference to the same store. + +The single-instance-shared-across-resources invariant is the architectural +linchpin: a future second adapter must see the *same* handle as the channel +adapter so a global byte budget and LRU still apply. +""" + +from __future__ import annotations + +import pytest + +from sift_client._internal.disk_cache import DiskCache + + +def _make_client(): + """Build a SiftClient-like object with the bits the namespace needs. + + Reaching into ``sift_client.SiftClient.__init__`` requires a live gRPC + config; the namespace only touches ``_disk_cache_config`` and + ``_disk_cache``, so a tiny stand-in keeps these tests independent of + transport setup. + """ + from sift_client._internal.disk_cache_config import DiskCacheConfig + from sift_client.util.cache import CacheNamespace + + class _StandinClient: + def __init__(self) -> None: + self._disk_cache_config = DiskCacheConfig(enabled=True) + self._disk_cache: DiskCache | None = None + self.cache = CacheNamespace(self) # type: ignore[arg-type] + + # Matches the real ``SiftClient._get_disk_cache`` so any namespace + # code that goes through this accessor (e.g. ``stats()``) sees + # the same lazy-init semantics in tests. + def _get_disk_cache(self) -> DiskCache: + return _get_disk_cache(self) + + return _StandinClient() + + +# Pull the same lazy-init helper the real client uses so we exercise the +# default-path-fallback path against the live code rather than a mock. +def _get_disk_cache(client) -> DiskCache: + if client._disk_cache is None: + config = client._disk_cache_config + if not config.enabled: + client._disk_cache = DiskCache() + return client._disk_cache + target_path = config.path or DiskCache.DEFAULT_DISK_PATH + try: + client._disk_cache = DiskCache( + disk_path=target_path, + disk_max_bytes=config.max_bytes, + ) + except Exception: + if not config.using_default_path: + raise + client._disk_cache = DiskCache() + return client._disk_cache + + +class TestCacheNamespaceDefaults: + """Opt-out default: the namespace is on, default path, fresh start.""" + + def test_enabled_by_default(self): + """First lazy access lands at ``DiskCache.DEFAULT_DISK_PATH``.""" + client = _make_client() + store = _get_disk_cache(client) + try: + assert store.disk_enabled + assert store.disk_path == DiskCache.DEFAULT_DISK_PATH + finally: + store.close() + + def test_one_store_shared_across_lazy_calls(self): + """Re-entering ``_get_disk_cache`` returns the same handle.""" + client = _make_client() + first = _get_disk_cache(client) + second = _get_disk_cache(client) + try: + assert first is second + finally: + first.close() + + +class TestEnable: + """``client.cache.enable`` configures the store, pre- and post-init.""" + + def test_pre_init_path_lands_on_store(self, tmp_path): + client = _make_client() + client.cache.enable(path=str(tmp_path / "pre"), max_bytes=4096) + store = _get_disk_cache(client) + try: + assert store.disk_enabled + assert store.disk_path == str(tmp_path / "pre") + assert store.disk_max_bytes == 4096 + finally: + store.close() + + def test_post_init_swap_uses_same_store_instance(self, tmp_path): + """Reconfiguring after first use mutates in place rather than re-creating. + + Every resource adapter holds a reference to ``client._disk_cache``; + if a reconfig replaced the handle, those adapters would still see + the stale one. ``DiskCache.enable`` swaps the *contents* on + the same instance. + """ + client = _make_client() + client.cache.disable() # start from off so this is a real on transition + store = _get_disk_cache(client) + try: + assert not store.disk_enabled + client.cache.enable(path=str(tmp_path / "post")) + assert client._disk_cache is store # same instance + assert store.disk_enabled + assert store.disk_path == str(tmp_path / "post") + finally: + store.close() + + def test_enable_with_default_path_lands_on_default(self, monkeypatch, tmp_path): + """``enable()`` with no args uses :attr:`DEFAULT_DISK_PATH`. + + Redirects the constant so the test doesn't create the real + ``/tmp/sift-data-cache`` directory. + """ + fake_default = str(tmp_path / "fake-default") + monkeypatch.setattr(DiskCache, "DEFAULT_DISK_PATH", fake_default) + + client = _make_client() + client.cache.enable() + store = _get_disk_cache(client) + try: + assert store.disk_path == fake_default + finally: + store.close() + + +class TestDisable: + """``client.cache.disable`` turns the live cache off.""" + + def test_disable_closes_live_handle(self, tmp_path): + client = _make_client() + client.cache.enable(path=str(tmp_path / "to-close")) + store = _get_disk_cache(client) + try: + assert store.disk_enabled + client.cache.disable() + assert not store.disk_enabled + assert store.disk_path is None + finally: + store.close() + + def test_disable_before_lazy_init_keeps_store_off(self, tmp_path): + """Calling disable before first use means the lazy build skips the open.""" + client = _make_client() + client.cache.disable() + store = _get_disk_cache(client) + try: + assert not store.disk_enabled + finally: + store.close() + + +class TestClearProxy: + """``client.cache.clear`` proxies through to :meth:`DiskCache.clear_disk`.""" + + def test_clear_removes_directory(self, tmp_path): + path = tmp_path / "to-clear" + # Populate a real cache directory so the marker check passes. + cache = DiskCache(disk_path=path) + cache.close() + assert path.exists() + + client = _make_client() + client.cache.clear(path) + assert not path.exists() + + +class TestLazyInitFallback: + """The default-path-failure fallback used by ``SiftClient._get_disk_cache``.""" + + def test_default_path_failure_falls_back_to_no_cache(self, monkeypatch, tmp_path): + """If the default cache path can't be opened, the lazy init produces + a disabled :class:`DiskCache` rather than raising. + + Simulated by pointing ``DEFAULT_DISK_PATH`` at a path that already + exists as a regular file — ``os.makedirs(..., exist_ok=True)`` + raises ``FileExistsError`` for non-directory targets. + """ + blocker = tmp_path / "not-a-dir" + blocker.write_text("i am a file, not a directory") + monkeypatch.setattr(DiskCache, "DEFAULT_DISK_PATH", str(blocker)) + + client = _make_client() + store = _get_disk_cache(client) # must not raise + try: + assert not store.disk_enabled + finally: + store.close() + + def test_explicit_path_failure_propagates(self, tmp_path): + """An explicit path that can't be opened propagates the OSError. + + Silent fallback would hide a user mistake. + """ + blocker = tmp_path / "not-a-dir" + blocker.write_text("i am a file, not a directory") + + client = _make_client() + client.cache.enable(path=str(blocker)) + with pytest.raises(FileExistsError): + _get_disk_cache(client) + + +class TestStats: + """``client.cache.stats()`` reports the current cache state. + + Three shapes get pinned: + + 1. **Disabled** — every numeric field zeroes out and ``path`` is + ``None``. Matches the cold/silent-store contract. + 2. **Enabled, empty** — ``enabled=True``, ``path`` populated, + sizes/counts at zero. Distinguishes "nothing cached yet" from + "no cache configured". + 3. **Enabled, populated** — bucket and segment counters reflect + what the adapter wrote, and foreign-prefix rows don't bleed + into the channel counters. + """ + + def test_stats_when_disabled(self): + """Disabled cache reports zeros and ``None`` path.""" + client = _make_client() + client.cache.disable() + stats = client.cache.stats() + assert stats.enabled is False + assert stats.path is None + assert stats.max_bytes is None + assert stats.size_bytes == 0 + assert stats.entry_count == 0 + assert stats.channel_buckets == 0 + assert stats.channel_segments == 0 + assert "disabled" in str(stats).lower() + + def test_stats_when_enabled_empty(self, tmp_path): + """Enabled but empty cache reports the path and zero usage.""" + path = str(tmp_path / "empty") + client = _make_client() + client.cache.enable(path=path, max_bytes=8 * 1024 * 1024) + try: + stats = client.cache.stats() + assert stats.enabled is True + assert stats.path == path + assert stats.max_bytes == 8 * 1024 * 1024 + assert stats.entry_count == 0 + assert stats.channel_buckets == 0 + assert stats.channel_segments == 0 + # path appears in the friendly print + assert path in str(stats) + finally: + client._disk_cache.close() # type: ignore[union-attr] + + def test_stats_counts_buckets_and_segments(self, tmp_path): + """Writes increment ``channel_buckets`` and ``channel_segments``. + + Uses ``ChannelDataCache`` directly (rather than the + ``DataLowLevelClient`` end-to-end path) so the test stays + focused on the stats accounting. Two of the three channels + get a second segment so the bucket-vs-segment split is + observable, not just N == N. + """ + from datetime import datetime, timedelta, timezone + + import pandas as pd + + from sift_client._internal.low_level_wrappers.data import ChannelDataCache + + client = _make_client() + client.cache.enable(path=str(tmp_path / "stats")) + try: + store = client._get_disk_cache() + adapter = ChannelDataCache(store) + base = datetime(2025, 1, 1, tzinfo=timezone.utc) + for i, cid in enumerate(("c1", "c2", "c3")): + df = pd.DataFrame( + {cid: [float(i)]}, + index=pd.date_range(base, periods=1, freq="s", tz="UTC"), + ) + adapter.put_segment( + cid, + None, + df, + start_time=base, + end_time=base, + ) + # Append a second, disjoint segment to two of the buckets + # so segments > buckets. + for cid in ("c1", "c2"): + later = base + timedelta(minutes=5) + df = pd.DataFrame( + {cid: [9.0]}, + index=pd.date_range(later, periods=1, freq="s", tz="UTC"), + ) + adapter.put_segment(cid, None, df, start_time=later, end_time=later) + + stats = client.cache.stats() + assert stats.enabled is True + assert stats.channel_buckets == 3 + assert stats.channel_segments == 5 + # 3 index rows + 5 segment rows. + assert stats.entry_count == 8 + assert stats.size_bytes > 0 + finally: + client._disk_cache.close() # type: ignore[union-attr] + + def test_stats_ignores_foreign_adapter_keys_in_channel_count(self, tmp_path): + """Keys outside the channel namespace don't bump channel counters. + + Pins the prefix-scoping so a future second adapter doesn't + double-count here. + """ + client = _make_client() + client.cache.enable(path=str(tmp_path / "foreign")) + try: + store = client._get_disk_cache() + store.put("other:foo", "x", size_bytes=64) + store.put("other:bar", "y", size_bytes=64) + + stats = client.cache.stats() + assert stats.entry_count == 2 + assert stats.channel_buckets == 0 + assert stats.channel_segments == 0 + finally: + client._disk_cache.close() # type: ignore[union-attr] + + +class TestSiftClientIntegration: + """End-to-end through the real :class:`SiftClient.__init__` entry point. + + Asserts the wire-up: the namespace really lives at ``client.cache``, + the config is mutable through it, and the lazy ``_get_disk_cache`` + returns the configured store. + """ + + def _make_real_client(self): + from sift_client import SiftClient, SiftConnectionConfig + + return SiftClient( + connection_config=SiftConnectionConfig( + api_key="x", + grpc_url="disabled.invalid:0", + rest_url="https://disabled.invalid", + use_ssl=False, + ) + ) + + def test_attribute_present_and_uses_real_lazy_init(self, monkeypatch, tmp_path): + fake_default = str(tmp_path / "real-client-default") + monkeypatch.setattr(DiskCache, "DEFAULT_DISK_PATH", fake_default) + + client = self._make_real_client() + store = client._get_disk_cache() + try: + assert client.cache is not None + assert store.disk_enabled + assert store.disk_path == fake_default + finally: + store.close() + + def test_disable_before_first_get_data_keeps_store_off(self): + client = self._make_real_client() + client.cache.disable() + store = client._get_disk_cache() + try: + assert not store.disk_enabled + finally: + store.close() diff --git a/python/lib/sift_client/client.py b/python/lib/sift_client/client.py index 2e2b64ffd..bc03e74aa 100644 --- a/python/lib/sift_client/client.py +++ b/python/lib/sift_client/client.py @@ -1,6 +1,12 @@ from __future__ import annotations +import logging +import warnings +from typing import TYPE_CHECKING + +from sift_client._internal.disk_cache_config import DiskCacheConfig from sift_client._internal.urls import frontend_origin_for_api +from sift_client.errors import SiftWarning from sift_client.resources import ( AssetsAPI, AssetsAPIAsync, @@ -43,8 +49,14 @@ WithGrpcClient, WithRestClient, ) +from sift_client.util.cache import CacheNamespace from sift_client.util.util import AsyncAPIs +if TYPE_CHECKING: + from sift_client._internal.disk_cache import DiskCache + +logger = logging.getLogger(__name__) + class SiftClient( WithGrpcClient, @@ -126,6 +138,9 @@ class SiftClient( data_import: DataImportAPI """Instance of the Data Import API for making synchronous requests.""" + cache: CacheNamespace + """Surface for the shared on-disk cache used by every cache-aware resource.""" + async_: AsyncAPIs """Accessor for the asynchronous APIs. All asynchronous APIs are available as attributes on this accessor.""" @@ -148,6 +163,7 @@ def __init__( Set this for on-prem or custom deployments whose API host can't be mapped to a frontend automatically; see the ``app_url`` property. A value here takes precedence over ``connection_config.app_url``. + """ if not (api_key and grpc_url and rest_url) and not connection_config: raise ValueError( @@ -179,6 +195,14 @@ def __init__( # pytest plugin's ``--sift-disabled`` mode. self._simulate: bool = False + # Shared on-disk cache: user intent in ``_disk_cache_config`` (opt-out + # default), live handle in ``_disk_cache`` (lazy so importing this + # module doesn't pay the diskcache cost up front). The + # ``client.cache`` namespace mutates both. + self._disk_cache_config = DiskCacheConfig(enabled=True) + self._disk_cache: DiskCache | None = None + self.cache = CacheNamespace(self) + self.ping = PingAPI(self) self.assets = AssetsAPI(self) self.calculated_channels = CalculatedChannelsAPI(self) @@ -230,6 +254,52 @@ def rest_client(self) -> RestClient: """The REST client used by the SiftClient for making REST API calls.""" return self._rest_client + def _get_disk_cache(self) -> DiskCache: + """Lazy accessor for the shared on-disk cache. Internal to resources. + + The cache is built on first use so that importing ``sift_client`` + doesn't pay the ``diskcache``/``sqlite`` cost up front. The opt-out + default ("disk caching on at the temp-dir path") is applied here, + along with the silent-fallback-on-default-path failure: if the + user left :class:`DiskCacheConfig` at its defaults and opening + fails (read-only ``/tmp``, restricted container, ...), we log a + warning and return a disabled :class:`DiskCache` so resources can + still serve requests by going to the wire. An explicit user- + supplied path that can't be opened propagates so the caller knows + their request didn't take. + + After the first call this just returns the memoized handle. + Subsequent ``client.cache.enable(...)`` calls mutate the + existing handle in place; this method is not re-entered. + """ + if self._disk_cache is None: + from sift_client._internal.disk_cache import DiskCache + + config = self._disk_cache_config + if not config.enabled: + self._disk_cache = DiskCache() + return self._disk_cache + target_path = config.path or DiskCache.DEFAULT_DISK_PATH + try: + self._disk_cache = DiskCache( + disk_path=target_path, + disk_max_bytes=config.max_bytes, + ) + except Exception: + if not config.using_default_path: + raise + warnings.warn( + f"Could not open the default sift data cache at {target_path}; " + "falling back to no caching. Call " + "``client.cache.disable()`` to silence this " + "warning, or pass an explicit path via " + "``client.cache.enable(path=...)``.", + SiftWarning, + stacklevel=2, + ) + self._disk_cache = DiskCache() + return self._disk_cache + @property def app_url(self) -> str | None: """The Sift web-app origin for this client, or None if it can't be determined. diff --git a/python/lib/sift_client/resources/channels.py b/python/lib/sift_client/resources/channels.py index aa5cdf96e..df5d218a9 100644 --- a/python/lib/sift_client/resources/channels.py +++ b/python/lib/sift_client/resources/channels.py @@ -242,9 +242,20 @@ async def unarchive(self, channels: list[str | Channel]) -> None: def _ensure_data_low_level_client(self): """Ensure that the data low level client is initialized. Separated out like this to not require large dependencies (pandas/pyarrow) for the client if not fetching data.""" if self._data_low_level_client is None: - from sift_client._internal.low_level_wrappers.data import DataLowLevelClient - - self._data_low_level_client = DataLowLevelClient(grpc_client=self.client.grpc_client) + from sift_client._internal.low_level_wrappers.data import ( + ChannelDataCache, + DataLowLevelClient, + ) + + # The shared on-disk store lives on the client; we just wrap it + # in the channel-side adapter. Cache configuration (enable / + # disable / clear / path / max_bytes) is owned by + # ``client.cache`` — there's no resource-level knob anymore. + store = self.client._get_disk_cache() + self._data_low_level_client = DataLowLevelClient( + grpc_client=self.client.grpc_client, + channel_cache=ChannelDataCache(store), + ) async def get_data( self, diff --git a/python/lib/sift_client/util/cache.py b/python/lib/sift_client/util/cache.py new file mode 100644 index 000000000..265e3e863 --- /dev/null +++ b/python/lib/sift_client/util/cache.py @@ -0,0 +1,273 @@ +"""User-facing surface for the shared on-disk cache. + +This module hosts the small bag of methods exposed as ``client.cache``, +plus the :class:`CacheStats` snapshot type returned by +:meth:`CacheNamespace.stats`. Both classes are public — they're +re-exported from the top-level ``sift_client`` package and surfaced in +the generated API docs. + +The cache itself (a :class:`~sift_client._internal.disk_cache.DiskCache`) +lives on :class:`~sift_client.client.SiftClient` so every resource that +wants to persist results across calls can reach into one shared store. +The namespace deliberately mirrors :class:`DiskCache` rather than the +old per-resource API (``client.channels.enable_data_cache_disk(...)``): +since the store is shared, configuration is global. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from sift_client._internal.disk_cache import DiskCache + +if TYPE_CHECKING: + import os + + from sift_client.client import SiftClient + +logger = logging.getLogger(__name__) + + +_BYTE_UNITS = ( + (1024**4, "TiB"), + (1024**3, "GiB"), + (1024**2, "MiB"), + (1024, "KiB"), +) + + +def _format_bytes(n: int) -> str: + """Render ``n`` bytes in the largest unit that doesn't underflow to zero.""" + for threshold, suffix in _BYTE_UNITS: + if n >= threshold: + return f"{n / threshold:.1f} {suffix}" + return f"{n} B" + + +@dataclass(frozen=True) +class CacheStats: + """Snapshot of the shared on-disk cache at call time. + + Returned by :meth:`CacheNamespace.stats`. Frozen dataclass so it + plays well with logging, snapshot tests, and "compare two readings" + diagnostics without surprise mutation. + + Field semantics: + + * **enabled** — whether the disk handle is open. When ``False``, all + the size/count fields are zero regardless of on-disk state. + * **path** — directory the cache is open against, or ``None`` when + disabled. Useful for "where does this cache actually live?". + * **max_bytes** — configured byte cap on disk usage, or ``None`` + when disabled. ``diskcache``'s LRU evicts once usage approaches + this. + * **size_bytes** — current on-disk usage including SQLite overhead. + Tends to trend slightly higher than the sum of per-entry + ``size_bytes`` the resources hand to the store. + * **entry_count** — total cache keys across all adapter prefixes + (channel index + segment rows, plus any future foreign-adapter + rows). + * **channel_buckets** — number of distinct + ``(channel_id, run_id)`` channel buckets currently cached. + Counted by walking the channel adapter's index keys; one index + per bucket regardless of how many segments live under it. + * **channel_segments** — total cached segment frames across every + bucket. Each ``put_segment`` call adds one; ``diskcache`` LRU + evictions can drop them independently of their parent index, so + this can be lower than the sum of segments the adapter has + written over the session. + + ``str(stats)`` prints a multi-line summary suitable for + notebook/REPL display; the structured fields are available for + programmatic checks. + """ + + enabled: bool + path: str | None + max_bytes: int | None + size_bytes: int + entry_count: int + channel_buckets: int + channel_segments: int + + def __str__(self) -> str: + if not self.enabled: + return "Sift cache: disabled" + cap = _format_bytes(self.max_bytes) if self.max_bytes is not None else "no cap" + pct = f" ({self.size_bytes / self.max_bytes * 100:.1f}%)" if self.max_bytes else "" + return ( + "Sift cache:\n" + f" path: {self.path}\n" + f" used: {_format_bytes(self.size_bytes)} / {cap}{pct}\n" + f" channels: {self.channel_buckets} buckets, " + f"{self.channel_segments} segments\n" + f" entries: {self.entry_count} total" + ) + + +class CacheNamespace: + """Resource-agnostic surface for the on-disk cache shared by all resources. + + Exposed as ``client.cache``. The actual handle + (:class:`~sift_client._internal.disk_cache.DiskCache`) is constructed + lazily on first use so importing :mod:`sift_client` doesn't pay the + diskcache cost up front. Configuration changes made before that + first use are recorded against the client's + :class:`~sift_client._internal.disk_cache_config.DiskCacheConfig` + and applied when the store opens; changes after first use are + routed directly to the live store. + + Default policy: disk caching is **opt-out** (the config is + constructed with ``enabled=True``). Users who don't want any state + on disk call :meth:`disable` to silence it; users who want a custom + location or byte cap call :meth:`enable` with arguments. + """ + + def __init__(self, client: SiftClient): + """Bind this namespace to ``client``. Constructed by :class:`SiftClient`.""" + self._client = client + + def enable( + self, + *, + path: str | os.PathLike[str] | None = None, + max_bytes: int | None = None, + ) -> None: + """Enable (or reconfigure) on-disk caching. + + Disk caching is **on by default** at + :attr:`~sift_client._internal.disk_cache.DiskCache.DEFAULT_DISK_PATH`; + use this method to override the path or size, or to turn the + cache back on after a prior :meth:`disable` call. + + Reconfiguring a live cache (different ``path`` or ``max_bytes``) + closes the previous handle and opens a new one. Existing entries + at the new path become available as cache hits. + + An explicit ``path`` that can't be opened (permission denied, + read-only filesystem, ...) raises so the caller knows their + request didn't take. The default-path open does *not* raise — + see :meth:`SiftClient._get_disk_cache` for the silent fall-back. + + Args: + path: Directory to persist to. ``None`` (the default) uses + the cache's :attr:`DEFAULT_DISK_PATH`. + max_bytes: Byte cap on disk usage. ``None`` uses the cache's + :attr:`DEFAULT_DISK_MAX_BYTES` (4 GiB). When the bound + is reached, ``diskcache``'s LRU eviction takes over. + + Example: + client.cache.enable(path="/data/sift-cache") + client.cache.enable(max_bytes=1024 ** 3) # 1 GiB + """ + client = self._client + client._disk_cache_config.enable(path=path, max_bytes=max_bytes) + if client._disk_cache is not None: + client._disk_cache.enable(path=path, max_bytes=max_bytes) + + def disable(self) -> None: + """Opt out of on-disk caching (no reads or writes). + + Caching is on by default; call this when you don't want any + cached data written to or read from disk. Closes any open cache + file handle. The on-disk directory is NOT deleted — use + :meth:`clear` to wipe it. + """ + client = self._client + client._disk_cache_config.disable() + if client._disk_cache is not None: + client._disk_cache.disable() + + def stats(self) -> CacheStats: + """Return a snapshot of the current cache state. + + Calling this triggers the lazy disk-handle open if it hasn't + happened yet (so a stats call on a fresh client doesn't + mis-report as disabled just because no resource has touched + the store). After the open it's read-only — no mutation or + migration. + + Channel-specific counts come from walking the keyspace once + and counting keys under the channel adapter's namespace prefix. + If a second cache-aware resource grows its own key prefix + later, extend this method to surface per-adapter counts; for + now there's only one adapter so the dedicated field stays flat. + + Example: + >>> print(client.cache.stats()) + Sift cache: + path: /tmp/sift-data-cache + used: 142.3 MiB / 4.0 GiB (3.5%) + channels: 12 buckets, 487 segments + entries: 499 total + """ + # Importing here keeps this module light at import time — pulling + # the adapter pulls pandas, which is the whole reason + # ``_ensure_data_low_level_client`` is lazy too. + from sift_client._internal.low_level_wrappers.data import ChannelDataCache + + # ``_get_disk_cache`` is idempotent: it opens the store on first + # use and memoizes, so calling it here for a stats peek is the + # same cost as the first resource call would have paid. + # Without this, a fresh client that hasn't yet used a cache- + # aware resource would mis-report as disabled. + store = self._client._get_disk_cache() + if not store.disk_enabled: + return CacheStats( + enabled=False, + path=None, + max_bytes=None, + size_bytes=0, + entry_count=0, + channel_buckets=0, + channel_segments=0, + ) + + # One pass over the keyspace. Cheap — diskcache keys are SQLite + # rows, and we only touch metadata (no value loads). + # Bucket index keys end in ``:idx`` and segment keys carry + # ``:seg:`` inside the suffix; those two shapes partition the + # channel adapter's namespace cleanly. + channel_prefix = ChannelDataCache.KEY_PREFIX + channel_buckets = 0 + channel_segments = 0 + for key in store: + if not key.startswith(channel_prefix): + continue + if key.endswith(":idx"): + channel_buckets += 1 + elif ":seg:" in key: + channel_segments += 1 + return CacheStats( + enabled=True, + path=store.disk_path, + max_bytes=store.disk_max_bytes, + size_bytes=store.volume(), + entry_count=len(store), + channel_buckets=channel_buckets, + channel_segments=channel_segments, + ) + + def clear(self, path: str | os.PathLike[str] | None = None) -> None: + """Delete a previously-persisted on-disk cache directory. + + Drops stale caches from previous sessions, recovers from a + corrupt cache, or reclaims disk space. Removes the directory + entirely; if disk caching is on, the next access re-opens an + empty cache at the same path. + + Args: + path: Directory of the cache to clear. ``None`` (the default) + targets the cache's + :attr:`~sift_client._internal.disk_cache.DiskCache.DEFAULT_DISK_PATH`. + + Raises: + ValueError: If ``path`` exists but does not look like a sift + data cache directory. + """ + DiskCache.clear_disk(path) + + +__all__ = ["CacheNamespace", "CacheStats"] diff --git a/python/pyproject.toml b/python/pyproject.toml index b435022e7..dfe94c043 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -36,6 +36,7 @@ dependencies = [ "requests~=2.25", "requests-toolbelt~=1.0", "alive-progress~=3.0", + "diskcache~=5.6", # May move these to optional dependencies in the future. "pandas-stubs>=2.0,<4.0", "types-PyYAML~=6.0", @@ -251,6 +252,13 @@ dev-all = ["development", "all", "build"] docs-build = ["dev-all", "docs"] # Note python 3.9+ [tool.uv] +# Pin uv to a version that writes lockfile revision 3 (introduced in uv 0.8.4 +# by astral-sh/uv#14489, which added ``exclude-newer-package`` to the lock +# schema). Older uv silently rolls the lockfile back to revision 2 on the +# next ``uv lock`` / ``uv sync`` (a no-op-looking change), then a teammate +# on a newer uv re-bumps it — churning the revision field in PRs. +# ``required-version`` blocks the older uv up front with a clear error. +required-version = ">=0.8.4" # Fork resolution per Python minor in the support range. Each fork resolves # independently, which lets 3.8 pick numpy 1.24.x + rosbags 0.9.23 without # being constrained by the 3.9+ universe (numpy 2.0 drops 3.8). @@ -352,6 +360,13 @@ module = "nptdms" ignore_missing_imports = true ignore_errors = true +# diskcache ships without inline type hints or PEP 561 marker. Used by the +# channel data cache's optional on-disk tier. +[[tool.mypy.overrides]] +module = "diskcache" +ignore_missing_imports = true +ignore_errors = true + # alive-progress 3.3.0 ships py.typed but its `alive_it` signature is too # tight (declares `Collection[Never]`), which breaks unpacking generators. [[tool.mypy.overrides]] diff --git a/python/uv.lock b/python/uv.lock index d152551a9..7a0c68645 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -638,6 +638,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/02/c3/253a89ee03fc9b9682f1541728eb66db7db22148cd94f89ab22528cd1e1b/deprecation-2.1.0-py2.py3-none-any.whl", hash = "sha256:a10811591210e1fb0e768a8c25517cabeabcba6f0bf96564f8ff45189f90b14a", size = 11178, upload-time = "2020-04-20T14:23:36.581Z" }, ] +[[package]] +name = "diskcache" +version = "5.6.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3f/21/1c1ffc1a039ddcc459db43cc108658f32c57d271d7289a2794e401d0fdb6/diskcache-5.6.3.tar.gz", hash = "sha256:2c3a3fa2743d8535d832ec61c2054a1641f41775aa7c556758a109941e33e4fc", size = 67916, upload-time = "2023-08-31T06:12:00.316Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/27/4570e78fc0bf5ea0ca45eb1de3818a23787af9b390c0b0a0033a1b8236f9/diskcache-5.6.3-py3-none-any.whl", hash = "sha256:5e31b2d5fbad117cc363ebaf6b689474db18a1f6438bc82358b024abd4c2ca19", size = 45550, upload-time = "2023-08-31T06:11:58.822Z" }, +] + [[package]] name = "eval-type-backport" version = "0.3.1" @@ -4334,6 +4343,7 @@ source = { editable = "." } dependencies = [ { name = "alive-progress", version = "3.1.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.9'" }, { name = "alive-progress", version = "3.3.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.9'" }, + { name = "diskcache" }, { name = "eval-type-backport" }, { name = "filelock", version = "3.16.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.9'" }, { name = "filelock", version = "3.19.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.9.*'" }, @@ -4562,6 +4572,7 @@ requires-dist = [ { name = "cffi", marker = "extra == 'dev-all'", specifier = "~=1.14" }, { name = "cffi", marker = "extra == 'docs-build'", specifier = "~=1.14" }, { name = "cffi", marker = "extra == 'openssl'", specifier = "~=1.14" }, + { name = "diskcache", specifier = "~=5.6" }, { name = "eval-type-backport", specifier = "~=0.2" }, { name = "filelock", specifier = "~=3.15" }, { name = "googleapis-common-protos", specifier = ">=1.60" },