629e1ba4ea
engine.py held four unrelated things: what the tables are, how to open a connection, how to read a database's state, and the Store that coordinates writes. At 1240 lines the Store's own concerns — locks, migrations, vacuum, tags — were hard to find among them. Table records, Arrow schemas, index_specs, ensure_indexes, REQUIRED_TABLES and query_to_pydantic move to store/schema.py, which imports nothing from haiku.rag: it describes the tables and never opens or mutates one. gather_database_info, get_database_stats, DatabaseInfo and its result models move to store/info.py. Nothing in Store calls them — they are read paths for the CLI, doctor, inspector and ingester API — so info depends on engine and not the reverse. engine.py keeps the Store, ConnectionMode, connect_lancedb, the tag helpers and the restore-order and retention constants. No re-exports: importers point at the new modules. test_app_info_uses_connect_lancedb_for_remote patched haiku.rag.store.engine.connect_lancedb; gather_database_info now binds that name in info.py, so the patch targets where the call is looked up.
105 lines
4.1 KiB
Python
105 lines
4.1 KiB
Python
import pytest
|
|
|
|
from haiku.rag.store.engine import Store
|
|
from haiku.rag.store.info import get_database_stats
|
|
|
|
|
|
class TestGetDatabaseStats:
|
|
@pytest.mark.asyncio
|
|
async def test_empty_database_stats(self, temp_db_path):
|
|
"""get_database_stats() on a fresh database reports zero rows and no vector index."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
stats = await get_database_stats(store.db)
|
|
|
|
for name in ("documents", "chunks", "document_items", "settings"):
|
|
assert stats[name]["exists"] is True
|
|
assert stats[name]["num_rows"] >= 0
|
|
assert stats[name]["total_bytes"] >= 0
|
|
assert stats[name]["num_versions"] >= 1
|
|
|
|
assert stats["documents"]["num_rows"] == 0
|
|
assert stats["chunks"]["num_rows"] == 0
|
|
assert stats["chunks"]["has_vector_index"] is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_missing_tables_report_absent(self, temp_db_path):
|
|
"""Tables that don't exist on the connection are reported as absent."""
|
|
import lancedb
|
|
from lancedb.pydantic import LanceModel
|
|
from pydantic import Field
|
|
|
|
class SettingsRecord(LanceModel):
|
|
id: str = Field(default="settings")
|
|
settings: str = Field(default="{}")
|
|
|
|
db = await lancedb.connect_async(temp_db_path)
|
|
await db.create_table("settings", schema=SettingsRecord)
|
|
|
|
stats = await get_database_stats(db)
|
|
|
|
assert stats["settings"]["exists"] is True
|
|
assert stats["documents"] == {"exists": False}
|
|
assert stats["chunks"] == {"exists": False}
|
|
assert stats["document_items"] == {"exists": False}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stats_after_adding_document(self, temp_db_path):
|
|
"""get_database_stats() reflects document and chunk counts after inserts."""
|
|
from haiku.rag.store.models import Chunk, Document
|
|
from haiku.rag.store.repositories.chunk import ChunkRepository
|
|
from haiku.rag.store.repositories.document import DocumentRepository
|
|
|
|
async with Store(temp_db_path, create=True) as store:
|
|
doc_repo = DocumentRepository(store)
|
|
chunk_repo = ChunkRepository(store)
|
|
|
|
doc = await doc_repo.create(Document(content="hello world"))
|
|
assert doc.id is not None
|
|
|
|
await chunk_repo.create(
|
|
Chunk(
|
|
content="hello world",
|
|
document_id=doc.id,
|
|
embedding=[0.0] * store.embedder._vector_dim,
|
|
)
|
|
)
|
|
|
|
stats = await get_database_stats(store.db)
|
|
assert stats["documents"]["num_rows"] == 1
|
|
assert stats["chunks"]["num_rows"] == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stats_with_vector_index(self, temp_db_path):
|
|
"""get_database_stats() reports vector index details once an index exists."""
|
|
from datetime import timedelta
|
|
|
|
from lancedb.index import IvfPq
|
|
|
|
async with Store(temp_db_path, create=True) as store:
|
|
dim = store.embedder._vector_dim
|
|
|
|
# Need >=256 rows for IVF_PQ training.
|
|
rows = [
|
|
{
|
|
"id": f"chunk-{i}",
|
|
"document_id": "doc-1",
|
|
"content": f"content {i}",
|
|
"content_fts": "",
|
|
"metadata": "{}",
|
|
"order": i,
|
|
"vector": [float(i % 7) + 0.01 * j for j in range(dim)],
|
|
}
|
|
for i in range(256)
|
|
]
|
|
await store.chunks_table.add(rows)
|
|
await store.chunks_table.create_index(
|
|
"vector", config=IvfPq(distance_type="cosine"), replace=True
|
|
)
|
|
await store.chunks_table.wait_for_index(
|
|
["vector_idx"], timeout=timedelta(minutes=1)
|
|
)
|
|
|
|
stats = await get_database_stats(store.db)
|
|
assert stats["chunks"]["has_vector_index"] is True
|
|
assert stats["chunks"]["num_indexed_rows"] >= 0
|
|
assert "num_unindexed_rows" in stats["chunks"]
|