spb/company-atlas

1"""Content-addressable object store for raw observations, normalized text and block structures (spec §26, §96, §97).23    objects/<sha256[0:2]>/<sha256[2:4]>/<sha256>.zst45Objects are written once (dedupe by hash) and never mutated; many observations may point at the same key. zstd level 9 gives6~8× on HTML. `key` = sha256 hex of the *uncompressed* bytes; the relative path is derived, so the store can move as a whole.7"""8from __future__ import annotations910import hashlib11import os12from pathlib import Path1314import zstandard as zstd1516from companyatlas.config import settings1718_cctx = zstd.ZstdCompressor(level=9)19_dctx = zstd.ZstdDecompressor()202122def object_path(key: str, root: Path | None = None) -> Path:23    root = root or settings.objects_dir24    return root / key[:2] / key[2:4] / f"{key}.zst"252627def sha256_hex(data: bytes) -> str:28    return hashlib.sha256(data).hexdigest()293031def put_bytes(data: bytes, *, key: str | None = None) -> tuple[str, int, bool]:32    """Store bytes; return (key, stored_size_bytes, created). Idempotent."""33    key = key or sha256_hex(data)34    path = object_path(key)35    if path.exists():36        return key, path.stat().st_size, False37    path.parent.mkdir(parents=True, exist_ok=True)38    tmp = path.with_suffix(f".zst.tmp{os.getpid()}")39    with open(tmp, "wb") as fh:40        fh.write(_cctx.compress(data))41    tmp.replace(path)42    return key, path.stat().st_size, True434445def put_text(text: str) -> tuple[str, int, bool]:46    return put_bytes(text.encode("utf-8"))474849def get_bytes(key: str) -> bytes:50    with open(object_path(key), "rb") as fh:51        return _dctx.decompress(fh.read(), max_output_size=64 * 1024 * 1024)525354def get_text(key: str) -> str:55    return get_bytes(key).decode("utf-8", errors="replace")565758def exists(key: str) -> bool:59    return object_path(key).exists()606162def store_stats(root: Path | None = None) -> dict[str, int]:63    root = root or settings.objects_dir64    total = count = 065    if root.exists():66        for dirpath, _dirs, files in os.walk(root):67            for f in files:68                if f.endswith(".zst"):69                    try:70                        total += os.stat(os.path.join(dirpath, f)).st_size71                        count += 172                    except OSError:73                        pass74    return {"objects": count, "bytes": total}757677__all__ = ["exists", "get_bytes", "get_text", "object_path", "put_bytes", "put_text", "sha256_hex", "store_stats"]78