arborist/tests/test_wallet_bucket.py

"""Bucket-direct queries via HttpRangeVFS.

Spins up an in-process HTTP server that serves files with `Range` support,
opens an arborist shard `.db` through `HttpRangeVFS`, runs the same FTS5
queries you'd run against a local sqlite3 connection. Asserts results
match a direct sqlite3 open of the same file.

Then fetches a chunk body from a `blobs/<hash>/...` path served by the
same HTTP server and confirms its hash matches the stored leaf_hash.
"""
from __future__ import annotations

import os
import sqlite3
import threading
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Iterator

import pytest

apsw = pytest.importorskip("apsw")

from arborist.document import Document
from arborist.ingest import ingest_source
from arborist.merkle import hash_leaf
from arborist.source import Source
from arborist.store import connect
from arborist.wallet.bucket import BucketClient, BucketEndpoint


class _RangeHandler(SimpleHTTPRequestHandler):
    """SimpleHTTPRequestHandler that honors `Range:` requests.

    Stdlib version supports Range only on Python 3.7+; explicit support
    here makes the test portable to older interpreters and gives the
    server an obvious "yes I do partial reads" affirmation.
    """

    def log_message(self, fmt, *args):
        return  # quiet

    def do_GET(self):
        path = self.translate_path(self.path)
        if not os.path.isfile(path):
            self.send_error(404)
            return
        size = os.path.getsize(path)
        rng = self.headers.get("Range")
        if not rng:
            with open(path, "rb") as f:
                body = f.read()
            self.send_response(200)
            self.send_header("Content-Length", str(size))
            self.send_header("Content-Type", "application/octet-stream")
            self.send_header("Accept-Ranges", "bytes")
            self.end_headers()
            self.wfile.write(body)
            return
        # Parse `bytes=N-M` (single range; multi-range is overkill here).
        assert rng.startswith("bytes=")
        start_s, end_s = rng[len("bytes="):].split("-")
        start = int(start_s)
        end = int(end_s) if end_s else size - 1
        end = min(end, size - 1)
        length = end - start + 1
        with open(path, "rb") as f:
            f.seek(start)
            body = f.read(length)
        self.send_response(206)
        self.send_header("Content-Range", f"bytes {start}-{end}/{size}")
        self.send_header("Content-Length", str(length))
        self.send_header("Content-Type", "application/octet-stream")
        self.send_header("Accept-Ranges", "bytes")
        self.end_headers()
        self.wfile.write(body)


def _serve_dir(directory: Path) -> tuple[ThreadingHTTPServer, threading.Thread, str]:
    """Serve `directory` over HTTP on a random local port. Returns
    (httpd, thread, base_url)."""
    cwd_before = os.getcwd()
    os.chdir(directory)
    httpd = ThreadingHTTPServer(("127.0.0.1", 0), _RangeHandler)
    port = httpd.server_address[1]
    thread = threading.Thread(target=httpd.serve_forever, daemon=True)
    thread.start()
    # Restore cwd; SimpleHTTPRequestHandler captures it per-handler-instance,
    # not per-server, but our requests run in handler threads so they
    # need the cwd they were started in. Workaround: leave the chdir
    # for the test's lifetime; the cleanup restores.
    def restore():
        os.chdir(cwd_before)
    httpd._restore_cwd = restore  # type: ignore[attr-defined]
    return httpd, thread, f"http://127.0.0.1:{port}"


class _S(Source):
    source_type = "test"
    def __init__(self, ds): self.ds = ds
    def iter_documents(self) -> Iterator[Document]:
        yield from self.ds


@pytest.fixture
def bucket_layout(tmp_path: Path) -> tuple[Path, str, dict[str, bytes]]:
    """Build a small corpus, write the .db + per-chunk blobs into a
    `bucket/` directory ready to be served over HTTP.

    Returns (bucket_dir, shard_filename_under_bucket, blobs_map).
    """
    bucket = tmp_path / "bucket"
    bucket.mkdir()
    clones = bucket / "clones" / "snap-1"
    clones.mkdir(parents=True)
    blobs = bucket / "blobs"

    db_path = clones / "000.db"
    conn = connect(db_path)
    try:
        ingest_source(conn, _S([
            Document(
                uri=f"test://doc-{i}",
                source_type="test",
                title=f"Doc {i}",
                content=(
                    f"This is document number {i}. " * 10
                    + f"The unique phrase for entry {i} is 'lemma-{i}'. " * 10
                ),
            )
            for i in range(5)
        ]))
    finally:
        conn.close()

    # Externalize every chunk to blobs/<hash[:2]>/<hash[2:]>.
    blob_map: dict[str, bytes] = {}
    with sqlite3.connect(str(db_path)) as src:
        src.row_factory = sqlite3.Row
        for r in src.execute("SELECT leaf_hash, content FROM chunks WHERE content IS NOT NULL"):
            text = r["content"]
            if isinstance(text, bytes):
                # might be zstd-packed; the existing pack flow ships
                # plain bytes via unpack_chunk first. Mirror that.
                from arborist.compress import unpack_chunk
                body = (unpack_chunk(text) or "").encode("utf-8")
            else:
                body = text.encode("utf-8")
            if hash_leaf(body).hex() != r["leaf_hash"]:
                continue  # affinity edge case: skip mismatches
            blob_dir = blobs / r["leaf_hash"][:2]
            blob_dir.mkdir(parents=True, exist_ok=True)
            (blob_dir / r["leaf_hash"][2:]).write_bytes(body)
            blob_map[r["leaf_hash"]] = body

    return bucket, "clones/snap-1/000.db", blob_map


def test_bucket_open_and_fts_match_local_sqlite(bucket_layout):
    """Open the bucket-served .db via HttpRangeVFS; FTS5 results must
    equal a direct sqlite3 open of the same file."""
    bucket, shard_rel, _blobs = bucket_layout
    httpd, _t, base_url = _serve_dir(bucket)
    try:
        endpoint = BucketEndpoint(
            shard_url=f"{base_url}/{shard_rel}",
            blob_base=f"{base_url}/blobs",
        )
        client = BucketClient(endpoint, cache_bytes=4 * 1024 * 1024)
        try:
            remote = client.fts_search('"lemma-3"', limit=4)
        finally:
            client.close()

        # Reference: direct sqlite3 against the file.
        local = []
        with sqlite3.connect(str(bucket / shard_rel)) as conn:
            conn.row_factory = sqlite3.Row
            for r in conn.execute(
                "SELECT d.document_root, d.document_uri, d.title, "
                "       bm25(chunks_fts) AS score "
                "FROM chunks_fts JOIN chunks c ON c.rowid = chunks_fts.rowid "
                "JOIN documents d ON d.document_root = c.document_root "
                "WHERE chunks_fts MATCH ? ORDER BY score LIMIT ?",
                ('"lemma-3"', 4),
            ):
                local.append({"document_root": r[0], "document_uri": r[1],
                              "title": r[2], "score": r[3]})

        assert remote == local, (
            "bucket-direct FTS5 results must equal local sqlite3 results"
        )
        # 'lemma-3' is unique to doc 3 — must be the top hit.
        assert remote, "FTS should have matched at least one chunk"
        assert remote[0]["document_uri"] == "test://doc-3"
    finally:
        httpd.shutdown()
        httpd._restore_cwd()  # type: ignore[attr-defined]


def test_bucket_fetch_chunk_body_round_trip(bucket_layout):
    """Pick one leaf_hash from the corpus, fetch it via BucketClient,
    confirm body hashes back to the same leaf_hash."""
    bucket, shard_rel, blob_map = bucket_layout
    assert blob_map, "fixture should have at least one uploaded blob"
    leaf_hash, expected_body = next(iter(blob_map.items()))

    httpd, _t, base_url = _serve_dir(bucket)
    try:
        endpoint = BucketEndpoint(
            shard_url=f"{base_url}/{shard_rel}",
            blob_base=f"{base_url}/blobs",
        )
        client = BucketClient(endpoint)
        try:
            body = client.fetch_chunk_body(leaf_hash)
        finally:
            client.close()
        assert body == expected_body
        assert hash_leaf(body).hex() == leaf_hash
    finally:
        httpd.shutdown()
        httpd._restore_cwd()  # type: ignore[attr-defined]


def test_bucket_snapshot_root_matches_local(bucket_layout):
    """snapshot_root computed from the bucket-direct conn must match
    snapshot_root computed from a local connection."""
    from arborist.snapshot import compute_snapshot_root

    bucket, shard_rel, _blobs = bucket_layout
    with sqlite3.connect(str(bucket / shard_rel)) as conn:
        conn.row_factory = sqlite3.Row
        local_snap, _ = compute_snapshot_root(conn)

    httpd, _t, base_url = _serve_dir(bucket)
    try:
        endpoint = BucketEndpoint(
            shard_url=f"{base_url}/{shard_rel}",
            blob_base=f"{base_url}/blobs",
        )
        client = BucketClient(endpoint)
        try:
            remote_snap = client.snapshot_root()
        finally:
            client.close()
        assert remote_snap == local_snap
    finally:
        httpd.shutdown()
        httpd._restore_cwd()  # type: ignore[attr-defined]


def test_bucket_cache_amortizes_repeated_queries(bucket_layout):
    """First query pays segment traversal; second identical query
    should serve mostly from the LRU page cache."""
    bucket, shard_rel, _blobs = bucket_layout
    httpd, _t, base_url = _serve_dir(bucket)
    try:
        endpoint = BucketEndpoint(
            shard_url=f"{base_url}/{shard_rel}",
            blob_base=f"{base_url}/blobs",
        )
        client = BucketClient(endpoint)
        try:
            client.fts_search('"lemma-3"', limit=4)
            stats_after_first = client.stats()
            client.fts_search('"lemma-3"', limit=4)
            stats_after_second = client.stats()
        finally:
            client.close()
        # The cost signal is "second query is cheap." SQLite has its
        # OWN page cache (default ~2000 pages) which absorbs the second
        # query entirely on small corpora — so our LRU may never see a
        # hit. What matters is that HTTP traffic doesn't multiply.
        r0 = stats_after_first["http_requests"]
        r1 = stats_after_second["http_requests"]
        assert r1 <= r0, (
            f"second identical query should add 0 HTTP requests "
            f"(SQLite-cached), got {r0} → {r1}"
        )
    finally:
        httpd.shutdown()
        httpd._restore_cwd()  # type: ignore[attr-defined]