From 99d2013d519c5867f5ddbc30bee8293761513c8f Mon Sep 17 00:00:00 2001 From: krypticmouse Date: Thu, 26 Mar 2026 19:41:35 +0000 Subject: [PATCH] feat: add content-addressed AttachmentStore with SHA-256 dedup Implements AttachmentStore that writes blobs to {base_dir}/{sha[:2]}/{sha} with a SQLite metadata index tracking filename, MIME type, size, and the accumulated list of source_doc_ids for each content-identical blob. Includes 7 unit tests covering SHA-256 return, file path layout, idempotent dedup, multi-source tracking, metadata retrieval, content round-trip, and missing-blob None return. Co-Authored-By: Claude Sonnet 4.6 --- src/openjarvis/connectors/attachment_store.py | 175 ++++++++++++++++++ tests/connectors/test_attachment_store.py | 102 ++++++++++ 2 files changed, 277 insertions(+) create mode 100644 src/openjarvis/connectors/attachment_store.py create mode 100644 tests/connectors/test_attachment_store.py diff --git a/src/openjarvis/connectors/attachment_store.py b/src/openjarvis/connectors/attachment_store.py new file mode 100644 index 00000000..7b653562 --- /dev/null +++ b/src/openjarvis/connectors/attachment_store.py @@ -0,0 +1,175 @@ +"""AttachmentStore — content-addressed blob storage for Deep Research attachments. + +Stores binary attachments at ``{base_dir}/{sha256[:2]}/{sha256}`` and tracks +metadata in a SQLite database at ``{base_dir}/attachments.db``. + +Deduplication is automatic: re-storing the same bytes accumulates source_doc_ids +but does not write a second copy of the file. + +Pure Python ``sqlite3`` (no Rust extension required). +""" + +from __future__ import annotations + +import hashlib +import json +import sqlite3 +import time +from pathlib import Path +from typing import Dict, List, Optional + +# --------------------------------------------------------------------------- +# DDL +# --------------------------------------------------------------------------- + +_CREATE_TABLE = """ +CREATE TABLE IF NOT EXISTS attachments ( + sha256 TEXT PRIMARY KEY, + filename TEXT NOT NULL DEFAULT '', + mime_type TEXT NOT NULL DEFAULT '', + size_bytes INTEGER NOT NULL DEFAULT 0, + source_doc_ids TEXT NOT NULL DEFAULT '[]', + created_at REAL NOT NULL +); +""" + +# --------------------------------------------------------------------------- +# AttachmentStore +# --------------------------------------------------------------------------- + + +class AttachmentStore: + """Content-addressed blob store with SQLite metadata index. + + Files are written to ``{base_dir}/{sha256[:2]}/{sha256}`` so that the + directory fan-out stays bounded even for millions of blobs. The SQLite + metadata table tracks filename, MIME type, size, and the list of + source document IDs that reference each blob. + """ + + def __init__(self, base_dir: str = "") -> None: + if not base_dir: + from openjarvis.core.config import DEFAULT_CONFIG_DIR + + base_dir = str(DEFAULT_CONFIG_DIR / "blobs") + + self._base_dir = Path(base_dir) + self._base_dir.mkdir(parents=True, exist_ok=True) + + db_path = self._base_dir / "attachments.db" + self._conn = sqlite3.connect(str(db_path), check_same_thread=False) + self._conn.row_factory = sqlite3.Row + self._setup() + + # ------------------------------------------------------------------ + # Internal setup + # ------------------------------------------------------------------ + + def _setup(self) -> None: + self._conn.execute("PRAGMA journal_mode=WAL;") + self._conn.execute(f"{_CREATE_TABLE}") + self._conn.commit() + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def store( + self, + content: bytes, + *, + filename: str, + mime_type: str = "", + source_doc_id: str = "", + ) -> str: + """Store *content* and return its SHA-256 hex digest. + + The call is idempotent with respect to the blob file — if the same + bytes are stored again only the metadata (source_doc_ids) is updated. + """ + sha = hashlib.sha256(content).hexdigest() + + # Write blob file (idempotent) + blob_dir = self._base_dir / sha[:2] + blob_dir.mkdir(parents=True, exist_ok=True) + blob_path = blob_dir / sha + if not blob_path.exists(): + blob_path.write_bytes(content) + + # Upsert metadata row + existing = self._conn.execute( + "SELECT source_doc_ids FROM attachments WHERE sha256 = ?", (sha,) + ).fetchone() + + if existing is None: + source_ids: List[str] = [] + if source_doc_id: + source_ids.append(source_doc_id) + self._conn.execute( + """ + INSERT INTO attachments + (sha256, filename, mime_type, size_bytes, + source_doc_ids, created_at) + VALUES (?, ?, ?, ?, ?, ?) + """, + ( + sha, + filename, + mime_type, + len(content), + json.dumps(source_ids), + time.time(), + ), + ) + else: + source_ids = json.loads(existing["source_doc_ids"]) + if source_doc_id and source_doc_id not in source_ids: + source_ids.append(source_doc_id) + self._conn.execute( + "UPDATE attachments SET source_doc_ids = ? WHERE sha256 = ?", + (json.dumps(source_ids), sha), + ) + + self._conn.commit() + return sha + + def get_metadata(self, sha: str) -> Optional[Dict]: + """Return the metadata dict for *sha*, or ``None`` if not found. + + Returned keys: ``sha256``, ``filename``, ``mime_type``, + ``size_bytes``, ``source_doc_ids`` (list), ``created_at``. + """ + row = self._conn.execute( + "SELECT * FROM attachments WHERE sha256 = ?", (sha,) + ).fetchone() + if row is None: + return None + return { + "sha256": row["sha256"], + "filename": row["filename"], + "mime_type": row["mime_type"], + "size_bytes": row["size_bytes"], + "source_doc_ids": json.loads(row["source_doc_ids"]), + "created_at": row["created_at"], + } + + def get_content(self, sha: str) -> Optional[bytes]: + """Return the raw bytes for *sha*, or ``None`` if the blob is missing.""" + blob_path = self._base_dir / sha[:2] / sha + if not blob_path.exists(): + return None + return blob_path.read_bytes() + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def close(self) -> None: + """Close the underlying SQLite connection.""" + try: + self._conn.close() + except Exception: + pass + + +__all__ = ["AttachmentStore"] diff --git a/tests/connectors/test_attachment_store.py b/tests/connectors/test_attachment_store.py new file mode 100644 index 00000000..562e8ff9 --- /dev/null +++ b/tests/connectors/test_attachment_store.py @@ -0,0 +1,102 @@ +"""Tests for AttachmentStore — content-addressed blob storage.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from openjarvis.connectors.attachment_store import AttachmentStore + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def store(tmp_path: Path) -> AttachmentStore: + return AttachmentStore(base_dir=str(tmp_path / "blobs")) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +def test_store_returns_sha256(store: AttachmentStore) -> None: + """store() returns a 64-character lowercase hex string (SHA-256).""" + sha = store.store(b"hello world", filename="hello.txt") + assert isinstance(sha, str) + assert len(sha) == 64 + assert sha == sha.lower() + # Only hex characters + int(sha, 16) + + +def test_store_creates_blob_file(store: AttachmentStore, tmp_path: Path) -> None: + """Blob file is written at {base_dir}/{sha[:2]}/{sha} with correct content.""" + data = b"binary content for blob test" + sha = store.store(data, filename="blob.bin", mime_type="application/octet-stream") + + blob_path = tmp_path / "blobs" / sha[:2] / sha + assert blob_path.exists(), f"Expected blob at {blob_path}" + assert blob_path.read_bytes() == data + + +def test_dedup_same_content(store: AttachmentStore) -> None: + """Storing identical bytes twice returns the same hash.""" + data = b"duplicate content" + sha1 = store.store(data, filename="file1.txt") + sha2 = store.store(data, filename="file2.txt") + assert sha1 == sha2 + + +def test_dedup_tracks_multiple_sources(store: AttachmentStore) -> None: + """Three stores of the same bytes accumulate three distinct source_doc_ids.""" + data = b"shared attachment content" + store.store(data, filename="attach.pdf", source_doc_id="doc:001") + store.store(data, filename="attach.pdf", source_doc_id="doc:002") + store.store(data, filename="attach.pdf", source_doc_id="doc:003") + + sha = store.store(data, filename="attach.pdf") # 4th call, no new source + meta = store.get_metadata(sha) + assert meta is not None + source_ids = meta["source_doc_ids"] + assert "doc:001" in source_ids + assert "doc:002" in source_ids + assert "doc:003" in source_ids + assert len(source_ids) == 3 # no duplicates + + +def test_get_metadata(store: AttachmentStore) -> None: + """get_metadata() returns correct filename, mime_type, and size_bytes.""" + data = b"metadata test payload" + sha = store.store( + data, + filename="report.pdf", + mime_type="application/pdf", + source_doc_id="doc:meta:001", + ) + + meta = store.get_metadata(sha) + assert meta is not None + assert meta["sha256"] == sha + assert meta["filename"] == "report.pdf" + assert meta["mime_type"] == "application/pdf" + assert meta["size_bytes"] == len(data) + assert "created_at" in meta + assert isinstance(meta["source_doc_ids"], list) + + +def test_get_content(store: AttachmentStore) -> None: + """get_content() returns the exact bytes that were stored.""" + data = b"\x00\x01\x02\x03binary\xff\xfe" + sha = store.store(data, filename="binary.bin") + retrieved = store.get_content(sha) + assert retrieved == data + + +def test_get_content_nonexistent(store: AttachmentStore) -> None: + """get_content() returns None for an unknown SHA-256.""" + fake_sha = "a" * 64 + assert store.get_content(fake_sha) is None