Brian Bjarke JensenandCursor
a768474892
Provide a generic disk-backed FIFO queue with configurable path, retention, and dedup keys so consumers can buffer items across restarts without optional extras. Co-authored-by: Cursor <[email protected]>
127 lines
4.4 KiB
Python
127 lines
4.4 KiB
Python
"""In-memory queue adapter with optional dedup and age-based eviction."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections import OrderedDict
|
|
from collections.abc import Hashable
|
|
from datetime import UTC, datetime
|
|
import threading
|
|
from typing import Any
|
|
|
|
from python_repositories.interfaces.queue_repository_interface import (
|
|
QueueRepositoryInterface,
|
|
)
|
|
|
|
|
|
def _parse_age(value: object) -> datetime:
|
|
"""Normalize an age field to a timezone-aware UTC datetime."""
|
|
if isinstance(value, datetime):
|
|
if value.tzinfo is None:
|
|
return value.replace(tzinfo=UTC)
|
|
return value.astimezone(UTC)
|
|
if isinstance(value, str):
|
|
normalized = value.replace("Z", "+00:00")
|
|
parsed = datetime.fromisoformat(normalized)
|
|
if parsed.tzinfo is None:
|
|
return parsed.replace(tzinfo=UTC)
|
|
return parsed.astimezone(UTC)
|
|
raise ValueError(f"age value must be datetime or ISO-8601 str, got {type(value)!r}")
|
|
|
|
|
|
class MemoryQueueAdapter(QueueRepositoryInterface):
|
|
"""Thread-safe in-process FIFO queue with optional deduplication."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
dedup_keys: tuple[str, ...] = (),
|
|
age_key: str | None = None,
|
|
) -> None:
|
|
self._dedup_keys = dedup_keys
|
|
self._age_key = age_key
|
|
self._lock = threading.RLock()
|
|
self._items: OrderedDict[Hashable, dict[str, Any]] = OrderedDict()
|
|
self._seq = 0
|
|
|
|
def enqueue(self, items: list[dict[str, Any]]) -> None:
|
|
"""Append items; skip duplicates when ``dedup_keys`` is configured."""
|
|
self.enqueue_and_return_added(items)
|
|
|
|
def enqueue_and_return_added(
|
|
self, items: list[dict[str, Any]]
|
|
) -> list[dict[str, Any]]:
|
|
"""Enqueue items and return the subset that was newly stored."""
|
|
if not items:
|
|
return []
|
|
added: list[dict[str, Any]] = []
|
|
with self._lock:
|
|
for raw in items:
|
|
item = self._normalize_item(raw)
|
|
key = self._make_key(item)
|
|
if self._dedup_keys and key in self._items:
|
|
continue
|
|
self._items[key] = item
|
|
added.append(item)
|
|
return added
|
|
|
|
def dequeue_batch(self, *, max_items: int = 1000) -> list[dict[str, Any]]:
|
|
"""Remove and return up to ``max_items`` items in FIFO order."""
|
|
if max_items < 0:
|
|
raise ValueError("max_items must be >= 0")
|
|
with self._lock:
|
|
batch: list[dict[str, Any]] = []
|
|
for _ in range(min(max_items, len(self._items))):
|
|
_key, item = self._items.popitem(last=False)
|
|
batch.append(item)
|
|
return batch
|
|
|
|
def size(self) -> int:
|
|
"""Return the number of items currently in the queue."""
|
|
with self._lock:
|
|
return len(self._items)
|
|
|
|
def evict_older_than(self, cutoff: datetime) -> int:
|
|
"""Remove items whose age field is strictly older than ``cutoff``."""
|
|
if self._age_key is None:
|
|
return 0
|
|
cutoff_utc = _parse_age(cutoff)
|
|
removed = 0
|
|
with self._lock:
|
|
to_remove = [
|
|
key
|
|
for key, item in self._items.items()
|
|
if _parse_age(item[self._age_key]) < cutoff_utc
|
|
]
|
|
for key in to_remove:
|
|
del self._items[key]
|
|
removed += 1
|
|
return removed
|
|
|
|
def clear(self) -> None:
|
|
"""Remove all items from the queue."""
|
|
with self._lock:
|
|
self._items.clear()
|
|
|
|
def snapshot(self) -> list[dict[str, Any]]:
|
|
"""Return a shallow copy of queued items in FIFO order."""
|
|
with self._lock:
|
|
return [dict(item) for item in self._items.values()]
|
|
|
|
def _normalize_item(self, raw: dict[str, Any]) -> dict[str, Any]:
|
|
item = dict(raw)
|
|
if self._dedup_keys:
|
|
missing = [key for key in self._dedup_keys if key not in item]
|
|
if missing:
|
|
raise ValueError(f"item missing dedup key(s): {missing}")
|
|
if self._age_key is not None:
|
|
if self._age_key not in item:
|
|
raise ValueError(f"item missing age key: {self._age_key!r}")
|
|
item[self._age_key] = _parse_age(item[self._age_key])
|
|
return item
|
|
|
|
def _make_key(self, item: dict[str, Any]) -> Hashable:
|
|
if not self._dedup_keys:
|
|
self._seq += 1
|
|
return self._seq
|
|
return tuple(item[key] for key in self._dedup_keys)
|