import threading
import uuid
from collections.abc import Iterator, Mapping, Sequence
from datetime import datetime, timezone
from typing import Any
import numpy as np
import psycopg
from heros.event import event
from jinja2 import StrictUndefined, Template, UndefinedError
from psycopg.types.json import Jsonb
from herostools.helper import log
from .artifact_storage import ArtifactStorage, InMemoryArtifactStorage, S3ArtifactStorage
from .base import HERODataArchiver
_DROPPED = object()
[docs]
def _daily_artifact_bucket(default_bucket: str | None, now: datetime | None = None) -> str | None:
"""Return the UTC daily bucket derived from a storage default bucket.
Args:
default_bucket: Static bucket used only for legacy-reference fallback.
now: Optional timestamp, used by tests to make routing deterministic.
Returns:
Daily bucket name, or None for storage backends without buckets.
"""
if default_bucket is None:
return None
current = now or datetime.now(timezone.utc)
return f"{default_bucket}-{current:%Y-%m-%d}"
[docs]
def _substitute_artifacts(obj: Any, bucket: str | None = None) -> tuple[Any, list[tuple[uuid.UUID, str, np.ndarray]]]:
"""Recursively replace numpy arrays with artifact refs; drop other non-JSON values.
The artifact name for each array is its dot-separated key path within the
original dict (e.g. "sensor.spectrum" for a nested key, "items.1" for a
list element at index 1). Each artifact_ref dict includes intrinsic array
metadata: ``size_bytes`` (element data only, excluding the npy header),
``shape``, and ``dtype``. ``bucket`` is the storage bucket name at write
time (``None`` when no storage backend is configured).
Args:
obj: Dict or list to walk.
bucket: Bucket name to embed in every artifact_ref produced.
Returns:
Tuple of (cleaned_obj, artifacts) where artifacts is a list of
(artifact_id, artifact_name, array) triples.
"""
artifacts: list[tuple[uuid.UUID, str, np.ndarray]] = []
def _walk(o: Any, path: str = "") -> Any:
if isinstance(o, np.ndarray):
aid = uuid.uuid4()
artifacts.append((aid, path, o))
array = np.asarray(o)
return {
"__type__": "artifact_ref",
"artifact_id": str(aid),
"size_bytes": int(array.nbytes),
"bucket": bucket,
"shape": list(array.shape),
"dtype": array.dtype.str,
}
if isinstance(o, dict):
result = {}
for k, v in o.items():
child_path = f"{path}.{k}" if path else k
walked = _walk(v, child_path)
if walked is _DROPPED:
log.warning(f"dropping non-serializable value for key '{k}' of type {type(v).__name__}")
else:
result[k] = walked
return result
if isinstance(o, list):
result = []
for i, v in enumerate(o):
child_path = f"{path}.{i}" if path else str(i)
walked = _walk(v, child_path)
if walked is not _DROPPED:
result.append(walked)
else:
log.warning(f"dropping non-serializable list element at index {i} of type {type(v).__name__}")
return result
if isinstance(o, (str, int, float, bool)) or o is None:
return o
return _DROPPED
return _walk(obj), artifacts
[docs]
def _collect_artifact_refs(obj: Any) -> list[dict]:
"""Recursively collect all artifact_ref dicts.
Args:
obj: Dict, list, or scalar to walk.
Returns:
List of artifact_ref dicts found.
"""
if isinstance(obj, dict):
if obj.get("__type__") == "artifact_ref":
return [obj]
refs = []
for v in obj.values():
refs.extend(_collect_artifact_refs(v))
return refs
if isinstance(obj, list):
refs = []
for v in obj:
refs.extend(_collect_artifact_refs(v))
return refs
return []
[docs]
def _resolve_artifacts_batch(obj: Any, lookup: dict[str, np.ndarray]) -> Any:
"""Recursively replace artifact_ref dicts using a pre-fetched id->array lookup.
Args:
obj: Dict, list, or scalar to walk.
lookup: Mapping of artifact_id strings to numpy arrays.
Returns:
Object with all artifact_ref dicts replaced by numpy arrays.
"""
if isinstance(obj, dict):
if obj.get("__type__") == "artifact_ref":
return lookup.get(obj["artifact_id"], obj)
return {k: _resolve_artifacts_batch(v, lookup) for k, v in obj.items()}
if isinstance(obj, list):
return [_resolve_artifacts_batch(v, lookup) for v in obj]
return obj
[docs]
class PostgresDatabase:
"""Pure SQL layer for the archiver schema — stores and retrieves clean JSON.
Can be used independently of HEROS to insert, query, and delete records.
Payloads must be clean JSON dicts (no numpy arrays). Artifact references
stored in ``data_json`` are returned as plain dicts by :meth:`get` without
resolution; use :class:`PostgresArchiver` for the full artifact pipeline.
Each identifier maps to exactly one row in ``records``. Successive payloads
for the same identifier are shallow-merged (``||``) at write time; later
events win on key collision.
Args:
db_url: psycopg connection string.
identifier_key: Key used to look up the record identifier in the payload.
allow_purge: Enable the :meth:`purge` method (disabled by default).
retention_days: Default retention window in days used by
:meth:`_find_expired`. ``0`` means keep forever.
"""
def __init__(
self,
db_url: str,
identifier_key: str = "identifier",
allow_purge: bool = False,
retention_days: int = 0,
):
self._db_url = db_url
self._local = threading.local()
self._identifier_key = identifier_key
self._allow_purge = allow_purge
self._retention_days = retention_days
[docs]
def _get_conn(self) -> psycopg.Connection:
"""Return a per-thread connection, opening one if necessary.
psycopg connections are not thread-safe; each thread owns its own connection.
"""
conn = getattr(self._local, "conn", None)
if conn is None or conn.closed:
self._local.conn = psycopg.connect(self._db_url, autocommit=True)
return self._local.conn
[docs]
def close(self) -> None:
"""Close the calling thread's database connection."""
conn = getattr(self._local, "conn", None)
if conn is not None and not conn.closed:
conn.close()
[docs]
def _store(self, payload: dict) -> str | None:
"""Upsert a clean JSON payload into records, merging into any existing row.
The identifier is extracted from the payload using :attr:`_identifier_key`.
Returns ``None`` and logs an error when the identifier key is missing.
Args:
payload: Clean JSON dict to store. Must not contain numpy arrays.
Returns:
The identifier string, or ``None`` if the identifier key is absent.
"""
identifier = payload.get(self._identifier_key)
if identifier is None:
log.error(f"identifier key '{self._identifier_key}' not found in payload, dropping record")
return None
identifier = str(identifier)
conn = self._get_conn()
with conn.transaction(), conn.cursor() as cur:
cur.execute(
"""
INSERT INTO records (identifier, data_json, updated_at)
VALUES (%s, %s, now())
ON CONFLICT (identifier) DO UPDATE
SET data_json = COALESCE(records.data_json, '{}'::jsonb) || EXCLUDED.data_json,
updated_at = EXCLUDED.updated_at
""",
(identifier, Jsonb(payload)),
)
return identifier
[docs]
def get_ids(self, tag: str, values: list[Any]) -> list[str]:
"""Return identifiers matching the given tag/key in either source.
Searches ``identifier_groups`` (post-hoc tags) and ``records.data_json``
(payload fields). Results are deduplicated and sorted.
Args:
tag: Tag key (``identifier_groups``) or top-level ``data_json`` key.
values: Accepted values (cast to str for comparison).
Returns:
Sorted, deduplicated list of matching identifiers.
"""
str_values = [str(v) for v in set(values)]
with self._get_conn().cursor() as cur:
cur.execute(
"""
SELECT identifier FROM (
SELECT identifier FROM identifier_groups
WHERE tag_key = %s AND tag_value = ANY(%s)
UNION
SELECT identifier FROM records
WHERE data_json->>%s = ANY(%s)
) sub
ORDER BY identifier
""",
(tag, str_values, tag, str_values),
)
return [row[0] for row in cur.fetchall()]
[docs]
def get(self, ids: list[str]) -> dict[str, dict]:
"""Fetch the merged record for each identifier, including post-hoc tags.
Tags stored via :meth:`tag_ids` are merged into the returned dict
alongside ``data_json`` fields. Tags overwrite ``data_json`` values on
key collision. Artifact references remain as plain dicts in the result.
Args:
ids: Identifiers to fetch.
Returns:
Dict mapping identifier to the merged record dict.
Identifiers with no record are absent from the result.
"""
with self._get_conn().cursor() as cur:
cur.execute(
"""
SELECT
r.identifier,
COALESCE(r.data_json, '{}'::jsonb)
|| COALESCE(
jsonb_object_agg(ig.tag_key, ig.tag_value)
FILTER (WHERE ig.tag_key IS NOT NULL),
'{}'::jsonb
) AS merged
FROM records r
LEFT JOIN identifier_groups ig ON ig.identifier = r.identifier
WHERE r.identifier = ANY(%s)
GROUP BY r.identifier, r.data_json
""",
(list(ids),),
)
return {identifier: merged for identifier, merged in cur.fetchall()}
[docs]
def tag_ids(self, ids: list[str], tag: str, value: str) -> None:
"""Attach a tag to a list of identifiers, overwriting any existing value for that tag.
Inserts ``(identifier, tag, value)`` into ``identifier_groups`` for each
identifier. If the ``(identifier, tag)`` pair already exists the value is
updated in place.
Args:
ids: Identifiers to tag. Each must already exist in ``records``.
tag: Tag key.
value: Tag value to assign.
Raises:
psycopg.errors.ForeignKeyViolation: If any identifier is not present
in ``records``.
"""
rows = [(str(i), tag, str(value)) for i in ids]
conn = self._get_conn()
with conn.transaction(), conn.cursor() as cur:
cur.executemany(
"""
INSERT INTO identifier_groups (identifier, tag_key, tag_value)
VALUES (%s, %s, %s)
ON CONFLICT (identifier, tag_key) DO UPDATE
SET tag_value = EXCLUDED.tag_value
""",
rows,
)
[docs]
def purge(self, ids: list[str]) -> None:
"""Delete DB rows for the given identifiers from records and identifier_groups.
Requires ``allow_purge=True`` at instantiation time. Does not touch
artifact storage; use :meth:`PostgresArchiver.purge` for that.
Args:
ids: Identifiers to delete.
Raises:
PermissionError: If ``allow_purge`` was not set to True.
"""
if not self._allow_purge:
raise PermissionError("purge() is disabled; set allow_purge=True to enable it")
id_list = list(ids)
conn = self._get_conn()
with conn.transaction(), conn.cursor() as cur:
cur.execute("DELETE FROM identifier_groups WHERE identifier = ANY(%s)", (id_list,))
cur.execute("DELETE FROM records WHERE identifier = ANY(%s)", (id_list,))
[docs]
def _find_expired(self, days: int) -> list[str]:
"""Return identifiers older than ``days`` days, excluding saved ones.
Records tagged with ``_save`` set to a truthy value (``'1'``,
``'true'``, ``'True'``) are excluded from the result.
Args:
days: Retention window in days.
Returns:
List of expired identifier strings.
"""
with self._get_conn().cursor() as cur:
cur.execute(
"""
SELECT identifier FROM records
WHERE updated_at < now() - interval '1 day' * %s
AND identifier NOT IN (
SELECT identifier FROM identifier_groups
WHERE tag_key = '_save'
AND tag_value IN ('1', 'true', 'True')
)
""",
(days,),
)
return [row[0] for row in cur.fetchall()]
[docs]
class PostgresRecordChangeListener:
"""Yield batches of changed record identifiers from PostgreSQL notifications.
Args:
db_url: Psycopg connection string.
"""
def __init__(self, db_url: str):
self._db_url = db_url
[docs]
def iter_changes(self, stop_event: threading.Event) -> Iterator[list[str]]:
"""Yield deduplicated notification batches until stopped.
Args:
stop_event: Event that stops the notification loop.
Yields:
Lists of changed record identifiers.
"""
conn = psycopg.connect(self._db_url, autocommit=True)
try:
conn.execute("LISTEN record_changed")
while not stop_event.is_set():
batch = []
for notify in conn.notifies(timeout=0.05):
batch.append(notify.payload)
if stop_event.is_set():
break
if batch:
yield list(dict.fromkeys(batch))
finally:
conn.close()
_USE_DEFAULT_S3 = object()
[docs]
class PostgresRecordStore(PostgresDatabase):
"""HEROS-independent PostgreSQL record store with artifact handling.
Owns the full artifact pipeline: numpy array detection, artifact storage
writes and reads, per-bucket resolution, and artifact deletion on purge.
:class:`PostgresDatabase` handles only SQL; this class adds the artifact
layer. It does not subscribe to HEROS events or PostgreSQL notifications.
Each identifier maps to one row in ``records``. Successive payloads are
shallow-merged at write time (``||``); later events win on key collision.
Numpy arrays are replaced by artifact_ref dicts in ``data_json`` and stored
in the configured artifact storage. Each ref embeds the bucket it was
written to, enabling transparent multi-bucket resolution and deletion.
Args:
db_url: psycopg connection string.
identifier_key: Key used to look up the record identifier.
allow_purge: Enable the :meth:`purge` method.
artifact_storage_kwargs: Passed as ``S3ArtifactStorage(**kwargs)``. Pass
``None`` to use ``InMemoryArtifactStorage`` (no persistence, useful
for dev/test). Omit to use local RustFS defaults.
array_key_template: Jinja2 template rendered against per-event metadata to
form the dict key for bare numpy-array payloads. ``source_name`` is
always available in the template context alongside all metadata fields.
Defaults to ``"{{ source_name }}"`` which preserves existing behaviour.
use_single_bucket: Store new artifacts in the configured static bucket
instead of UTC daily buckets. Defaults to False.
retention_days: Default retention window in days used by :meth:`cleanup`.
``0`` (default) means keep forever.
artifact_storage: Explicit storage instance. Mutually exclusive with
``artifact_storage_kwargs``.
"""
def __init__(
self,
db_url: str,
identifier_key: str = "identifier",
allow_purge: bool = False,
artifact_storage_kwargs: dict | None = _USE_DEFAULT_S3,
array_key_template: str = "{{ source_name }}",
use_single_bucket: bool = False,
retention_days: int = 0,
artifact_storage: ArtifactStorage | None = None,
):
if artifact_storage is not None and artifact_storage_kwargs is not _USE_DEFAULT_S3:
raise ValueError("artifact_storage and artifact_storage_kwargs are mutually exclusive")
if artifact_storage is not None:
storage = artifact_storage
elif artifact_storage_kwargs is _USE_DEFAULT_S3:
storage = S3ArtifactStorage()
elif artifact_storage_kwargs is None:
storage = InMemoryArtifactStorage()
else:
storage = S3ArtifactStorage(**artifact_storage_kwargs)
self._artifact_storage = storage
self._array_key_template = Template(array_key_template, undefined=StrictUndefined)
self._use_single_bucket = use_single_bucket
PostgresDatabase.__init__(self, db_url, identifier_key, allow_purge, retention_days)
[docs]
def store(self, source_name: str, payload: Any, metadata: dict) -> None:
"""Store one payload after extracting any numpy-array artifacts.
A bare numpy array payload is wrapped as ``{template_key: array}`` first.
``metadata`` is shallow-merged into the payload (payload wins on key
collision). Arrays are replaced by artifact_ref dicts with the bucket
embedded; artifacts are uploaded to storage before the DB write.
Args:
source_name: Name of the event source.
payload: The actual data (dict or numpy array).
metadata: Incoming metadata merged with ``default_metadata``.
"""
if isinstance(payload, np.ndarray):
try:
key = self._array_key_template.render({"source_name": source_name, **metadata})
except UndefinedError as exc:
log.warning(f"array_key_template failed for {source_name}: {exc}, dropping record")
return
payload = {key: payload}
if not isinstance(payload, dict):
log.warning(f"unsupported payload type {type(payload).__name__} from {source_name}, dropping record")
return
merged = metadata | payload
if merged.get(self._identifier_key) is None:
log.error(
f"identifier key '{self._identifier_key}' not found in payload from {source_name}, dropping record"
)
return
bucket = self._artifact_storage.default_bucket
if not self._use_single_bucket:
bucket = _daily_artifact_bucket(bucket)
clean_payload, artifacts = _substitute_artifacts(merged, bucket=bucket)
for artifact_id, _name, array in artifacts:
self._artifact_storage.put(str(artifact_id), array, bucket=bucket)
PostgresDatabase._store(self, clean_payload)
[docs]
def _store(self, source_name: str, payload: Any, metadata: dict) -> None:
"""Store one payload through the synchronous record-store API.
Args:
source_name: Name of the event source.
payload: The data to store.
metadata: Metadata merged into the payload.
"""
self.store(source_name, payload, metadata)
[docs]
def get(
self,
ids: list[str],
resolve_artifacts: bool = True,
size_limit_bytes: int | None = None,
) -> dict[str, dict]:
"""Fetch merged records and optionally resolve artifact_refs to numpy arrays.
Tags stored via :meth:`tag_ids` are merged into the returned dict
alongside ``data_json`` fields. Tags overwrite ``data_json`` values on
key collision (post-hoc annotations take priority).
Artifact_refs are resolved per-bucket: each ref's stored ``bucket``
field determines which bucket the array is fetched from. Refs without a
``bucket`` field (legacy records) fall back to the storage
``default_bucket``.
Args:
ids: Identifiers to fetch.
resolve_artifacts: Replace artifact_ref dicts with numpy arrays.
size_limit_bytes: When set, only artifacts whose stored ``size_bytes``
is at or below this threshold are fetched. Larger artifacts remain
as artifact_ref dicts. Refs without ``size_bytes`` are always
resolved. ``None`` resolves all regardless of size.
Returns:
Dict mapping identifier to the merged record dict.
Identifiers with no record are absent from the result.
"""
result = PostgresDatabase.get(self, ids)
if not resolve_artifacts:
return result
default_bucket = self._artifact_storage.default_bucket
all_refs: list[dict] = []
for record in result.values():
all_refs.extend(_collect_artifact_refs(record))
eligible = [ref for ref in all_refs if size_limit_bytes is None or ref.get("size_bytes", 0) <= size_limit_bytes]
by_bucket: dict[str | None, list[str]] = {}
for ref in eligible:
b = ref.get("bucket") or default_bucket
by_bucket.setdefault(b, []).append(ref["artifact_id"])
lookup: dict[str, np.ndarray] = {}
for b, aids in by_bucket.items():
unique_aids = list(dict.fromkeys(aids))
arrays = self._artifact_storage.get_many(unique_aids, bucket=b)
lookup.update(zip(unique_aids, arrays))
if lookup:
for identifier, record in result.items():
result[identifier] = _resolve_artifacts_batch(record, lookup)
return result
[docs]
def get_artifact(self, artifact_id: str, bucket: str | None = None) -> np.ndarray:
"""Retrieve one decoded artifact from the configured artifact storage.
Args:
artifact_id: Identifier of the artifact to retrieve.
bucket: Bucket override; uses the storage default when None.
Returns:
The decoded numpy array.
"""
return self._artifact_storage.get_artifact(artifact_id, bucket=bucket)
[docs]
def get_artifacts(self, locations: Sequence[tuple[str, str | None]]) -> Mapping[str, np.ndarray]:
"""Retrieve a bounded set of artifacts from the configured storage.
Args:
locations: Artifact ID and optional bucket pairs.
Returns:
Mapping of artifact IDs to decoded numpy arrays.
"""
return self._artifact_storage.get_artifacts(locations)
[docs]
def purge(self, ids: list[str], remove_artifacts: bool = True) -> None:
"""Delete all data for the given identifiers, including stored artifacts.
Artifacts are removed per-bucket using each ref's stored ``bucket``
field. Refs without a ``bucket`` field fall back to ``default_bucket``.
DB rows are deleted after artifact removal.
Requires ``allow_purge=True`` at instantiation time.
Args:
ids: Identifiers to purge.
remove_artifacts: Also remove referenced artifacts from storage.
Raises:
PermissionError: If ``allow_purge`` was not set to True.
"""
if not self._allow_purge:
msg = "purge() is disabled; set allow_purge=True to enable it"
log.error(msg)
raise PermissionError(msg)
id_list = list(ids)
if remove_artifacts:
default_bucket = self._artifact_storage.default_bucket
records = PostgresDatabase.get(self, id_list)
by_bucket: dict[str | None, list[str]] = {}
for data in records.values():
for ref in _collect_artifact_refs(data):
b = ref.get("bucket") or default_bucket
by_bucket.setdefault(b, []).append(ref["artifact_id"])
for b, aids in by_bucket.items():
self._artifact_storage.remove_many(aids, bucket=b)
PostgresDatabase.purge(self, id_list)
[docs]
def cleanup(
self,
retention_days: int | None = None,
remove_artifacts: bool = True,
dry_run: bool = True,
) -> list[str]:
"""Delete records older than the retention window, skipping saved ones.
Requires ``allow_purge=True``. Records tagged with ``_save`` set to a
truthy value (``'1'``, ``'true'``, ``'True'``) are never deleted.
Returns the list of expired identifiers (deleted or would-be-deleted).
Args:
retention_days: Purge records whose ``updated_at`` is older than this
many days. ``None`` uses the instance default set at construction.
``0`` is a no-op (keep forever).
remove_artifacts: Also remove referenced artifacts from storage.
dry_run: When True (default), only return the expired identifiers
without deleting anything.
Returns:
List of expired identifiers.
Raises:
PermissionError: If ``allow_purge`` was not set to True.
"""
if not self._allow_purge:
msg = "cleanup() is disabled; set allow_purge=True to enable it"
log.error(msg)
raise PermissionError(msg)
days = retention_days if retention_days is not None else self._retention_days
if days == 0:
return []
expired = self._find_expired(days)
if not dry_run and expired:
self.purge(expired, remove_artifacts=remove_artifacts)
return expired
[docs]
class PostgresArchiver(HERODataArchiver, PostgresRecordStore):
"""HEROS event adapter for :class:`PostgresRecordStore`.
Subscribes to incoming data events, queues and retries writes, and publishes
``record_changed`` notifications. Use :class:`PostgresRecordStore` for local
current-state queries without HEROS or notification connections.
"""
def __init__(
self,
db_url: str,
identifier_key: str = "identifier",
allow_purge: bool = False,
artifact_storage_kwargs: dict | None = _USE_DEFAULT_S3,
array_key_template: str = "{{ source_name }}",
use_single_bucket: bool = False,
retention_days: int = 0,
artifact_storage: ArtifactStorage | None = None,
*args,
**kwargs,
):
PostgresRecordStore.__init__(
self,
db_url=db_url,
identifier_key=identifier_key,
allow_purge=allow_purge,
artifact_storage_kwargs=artifact_storage_kwargs,
array_key_template=array_key_template,
use_single_bucket=use_single_bucket,
retention_days=retention_days,
artifact_storage=artifact_storage,
)
HERODataArchiver.__init__(self, *args, **kwargs)
self._change_listener = PostgresRecordChangeListener(self._db_url)
self._listener_thread = threading.Thread(target=self._listen_loop, daemon=True)
self._listener_thread.start()
@event
def record_changed(self, identifiers: list[str]) -> list[str]:
"""Publish changed record identifiers to HEROS subscribers.
Args:
identifiers: Deduplicated identifiers from PostgreSQL notifications.
Returns:
The published identifiers.
"""
return identifiers
[docs]
def _store(self, source_name: str, payload: Any, metadata: dict) -> None:
"""Store a queued HEROS payload through the record store.
Args:
source_name: Name of the event source.
payload: The data to store.
metadata: Metadata merged into the payload.
"""
self.store(source_name, payload, metadata)
[docs]
def _listen_loop(self) -> None:
"""Publish PostgreSQL change notifications through the HEROS event."""
for identifiers in self._change_listener.iter_changes(self._stop_event):
self.record_changed(identifiers)
[docs]
def _process_queue(self) -> None:
"""Drain the queue, then close this thread's DB connection on exit."""
super()._process_queue()
self.close()
[docs]
def _teardown(self) -> None:
"""Stop worker and listener threads; worker closes its own DB connection."""
self._stop()
self._listener_thread.join()