extract get_database_stats and share it across info, inspector, and backend

app/backend/main.py

@@ -149,18 +149,20 @@ async def db_info(_: Request) -> JSONResponse:
             }
         )
 
+    from haiku.rag.store.engine import get_database_stats
+
     client = get_client()
-    stats = client.store.get_stats()
+    stats = get_database_stats(client.store.db)
 
     return JSONResponse(
         {
             "exists": True,
             "path": str(db_path),
-            "documents": stats.get("documents", {}).get("num_rows", 0),
+            "documents": stats["documents"].get("num_rows", 0),
-            "chunks": stats.get("chunks", {}).get("num_rows", 0),
+            "chunks": stats["chunks"].get("num_rows", 0),
-            "documents_bytes": stats.get("documents", {}).get("total_bytes", 0),
+            "documents_bytes": stats["documents"].get("total_bytes", 0),
-            "chunks_bytes": stats.get("chunks", {}).get("total_bytes", 0),
+            "chunks_bytes": stats["chunks"].get("total_bytes", 0),
-            "has_vector_index": stats.get("chunks", {}).get("has_vector_index", False),
+            "has_vector_index": stats["chunks"].get("has_vector_index", False),
         }
     )
 

haiku_rag_slim/haiku/rag/app.py

@@ -68,7 +68,7 @@ class HaikuRAGApp:  # pragma: no cover
     async def info(self):
         """Display read-only information about the database without modifying it."""
 
-        from haiku.rag.store.engine import connect_lancedb
+        from haiku.rag.store.engine import connect_lancedb, get_database_stats
         from haiku.rag.store.upgrades import get_pending_upgrades
 
         if self.before is not None:
@@ -89,9 +89,9 @@ class HaikuRAGApp:  # pragma: no cover
         # Connect directly. Don't go through Store so a database that is
         # missing tables (e.g. pre-migration) still reports what it can.
         db = connect_lancedb(self.config, self.db_path)
-        existing_tables = set(db.list_tables().tables)
+        stats = get_database_stats(db)
 
-        if not existing_tables:
+        if not any(entry["exists"] for entry in stats.values()):
             self.console.print(
                 "[red]Database is empty. Use 'haiku-rag init' to initialize.[/red]"
             )
@@ -103,7 +103,7 @@ class HaikuRAGApp:  # pragma: no cover
         embed_provider = "unknown"
         embed_model = "unknown"
         vector_dim = None
-        if "settings" in existing_tables:
+        if stats["settings"]["exists"]:
             settings_tbl = db.open_table("settings")
             rows = (
                 settings_tbl.search()
@@ -131,16 +131,14 @@ class HaikuRAGApp:  # pragma: no cover
             f"{embed_provider}/{embed_model} (dim: {dim_part})"
         )
 
-        # Per-table row counts and sizes. Missing required tables are shown
+        # Per-table row counts and sizes. Missing required tables are
-        # as "absent" rather than raising.
+        # reported as "absent" rather than raising.
         for name in ("documents", "chunks", "document_items"):
-            if name in existing_tables:
+            entry = stats[name]
-                stats = db.open_table(name).stats()
+            if entry["exists"]:
-                num_rows = stats.get("num_rows", 0)
-                total_bytes = stats.get("total_bytes", 0)
                 self.console.print(
-                    f"  [repr.attrib_name]{name}[/repr.attrib_name]: {num_rows} "
+                    f"  [repr.attrib_name]{name}[/repr.attrib_name]: {entry['num_rows']} "
-                    f"({format_bytes(total_bytes)})"
+                    f"({format_bytes(entry['total_bytes'])})"
                 )
             else:
                 self.console.print(
@@ -148,20 +146,11 @@ class HaikuRAGApp:  # pragma: no cover
                 )
 
         # Vector index information
-        if "chunks" in existing_tables:
+        if stats["chunks"]["exists"]:
-            chunks_tbl = db.open_table("chunks")
+            num_chunks = stats["chunks"]["num_rows"]
-            num_chunks = chunks_tbl.stats().get("num_rows", 0)
+            if stats["chunks"].get("has_vector_index"):
-            indices = chunks_tbl.list_indices()
+                num_indexed_rows = stats["chunks"].get("num_indexed_rows", 0)
-            has_vector_index = any("vector" in str(idx).lower() for idx in indices)
+                num_unindexed_rows = stats["chunks"].get("num_unindexed_rows", 0)
-
-            if has_vector_index:
-                index_stats = chunks_tbl.index_stats("vector_idx")
-                num_indexed_rows = (
-                    index_stats.num_indexed_rows if index_stats is not None else 0
-                )
-                num_unindexed_rows = (
-                    index_stats.num_unindexed_rows if index_stats is not None else 0
-                )
                 self.console.print(
                     "  [repr.attrib_name]vector index[/repr.attrib_name]: ✓ exists"
                 )
@@ -189,15 +178,15 @@ class HaikuRAGApp:  # pragma: no cover
                         f"(need {256 - num_chunks} more chunks)"
                     )
 
-        if "documents" in existing_tables:
+        if stats["documents"]["exists"]:
-            doc_versions = len(list(db.open_table("documents").list_versions()))
             self.console.print(
-                f"  [repr.attrib_name]versions (documents)[/repr.attrib_name]: {doc_versions}"
+                f"  [repr.attrib_name]versions (documents)[/repr.attrib_name]: "
+                f"{stats['documents']['num_versions']}"
             )
-        if "chunks" in existing_tables:
+        if stats["chunks"]["exists"]:
-            chunk_versions = len(list(db.open_table("chunks").list_versions()))
             self.console.print(
-                f"  [repr.attrib_name]versions (chunks)[/repr.attrib_name]: {chunk_versions}"
+                f"  [repr.attrib_name]versions (chunks)[/repr.attrib_name]: "
+                f"{stats['chunks']['num_versions']}"
             )
 
         # Migration status
@@ -207,8 +196,8 @@ class HaikuRAGApp:  # pragma: no cover
         self.console.rule()
         if pending:
             self.console.print(
-                f"[bold yellow]Pending migrations: {len(pending)}[/bold yellow] "
+                f"[bold yellow]{len(pending)} migration(s) pending.[/bold yellow] "
-                "(run [cyan]haiku-rag migrate[/cyan] to upgrade)"
+                "Run [cyan]haiku-rag migrate[/cyan] to upgrade."
             )
             for step in pending:
                 desc = step.description or ""

haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py

@@ -64,7 +64,11 @@ class InfoModal(ModalScreen):
 
     async def on_mount(self) -> None:
         """Load and display database info."""
-        from haiku.rag.store.engine import ConnectionMode, connect_lancedb
+        from haiku.rag.store.engine import (
+            ConnectionMode,
+            connect_lancedb,
+            get_database_stats,
+        )
 
         lines: list[str] = []
 
@@ -81,7 +85,7 @@ class InfoModal(ModalScreen):
         config = self.client.store._config
         try:
             db = connect_lancedb(config, self.db_path)
-            table_names = set(db.list_tables().tables)
+            stats = get_database_stats(db)
         except Exception as e:
             lines.append(f"[red]Failed to open database: {e}[/red]")
             self._content_widget.update("\n".join(lines))
@@ -90,16 +94,13 @@ class InfoModal(ModalScreen):
         # Get versions
         versions = get_package_versions()
 
-        # Get stats from store
-        table_stats = self.client.store.get_stats()
-
         # Read settings
         stored_version = "unknown"
         embed_provider: str | None = None
         embed_model: str | None = None
         vector_dim: int | None = None
 
-        if "settings" in table_names:
+        if stats["settings"]["exists"]:
             settings_tbl = db.open_table("settings")
             arrow = settings_tbl.search().where("id = 'settings'").limit(1).to_arrow()
             rows = arrow.to_pylist() if arrow is not None else []
@@ -113,27 +114,18 @@ class InfoModal(ModalScreen):
                 embed_model = embed_model_obj.get("name")
                 vector_dim = embed_model_obj.get("vector_dim")
 
-        num_docs = table_stats["documents"].get("num_rows", 0)
+        num_docs = stats["documents"].get("num_rows", 0)
-        doc_bytes = table_stats["documents"].get("total_bytes", 0)
+        doc_bytes = stats["documents"].get("total_bytes", 0)
 
-        num_chunks = table_stats["chunks"].get("num_rows", 0)
+        num_chunks = stats["chunks"].get("num_rows", 0)
-        chunk_bytes = table_stats["chunks"].get("total_bytes", 0)
+        chunk_bytes = stats["chunks"].get("total_bytes", 0)
 
-        has_vector_index = table_stats["chunks"].get("has_vector_index", False)
+        has_vector_index = stats["chunks"].get("has_vector_index", False)
-        num_indexed_rows = table_stats["chunks"].get("num_indexed_rows", 0)
+        num_indexed_rows = stats["chunks"].get("num_indexed_rows", 0)
-        num_unindexed_rows = table_stats["chunks"].get("num_unindexed_rows", 0)
+        num_unindexed_rows = stats["chunks"].get("num_unindexed_rows", 0)
 
-        # Table versions
+        doc_versions = stats["documents"].get("num_versions", 0)
-        doc_versions = (
+        chunk_versions = stats["chunks"].get("num_versions", 0)
-            len(list(db.open_table("documents").list_versions()))
-            if "documents" in table_names
-            else 0
-        )
-        chunk_versions = (
-            len(list(db.open_table("chunks").list_versions()))
-            if "chunks" in table_names
-            else 0
-        )
 
         # Build output
         lines.append(

haiku_rag_slim/haiku/rag/store/engine.py

@@ -121,6 +121,50 @@ class SettingsRecord(LanceModel):
     settings: str = Field(default="{}")
 
 
+REQUIRED_TABLES: tuple[str, ...] = ("documents", "chunks", "document_items", "settings")
+
+
+def get_database_stats(db: lancedb.DBConnection) -> dict:
+    """Collect stats for every haiku.rag table on the connection.
+
+    Missing tables return ``{"exists": False}``. Present tables include
+    ``num_rows``, ``total_bytes``, and ``num_versions``. The ``chunks``
+    entry additionally reports vector index status and, when an index
+    exists, ``num_indexed_rows`` and ``num_unindexed_rows``.
+    """
+    existing = set(db.list_tables().tables)
+    stats: dict = {}
+    tables: dict = {}
+
+    for name in REQUIRED_TABLES:
+        if name not in existing:
+            stats[name] = {"exists": False}
+            continue
+        tbl = db.open_table(name)
+        tables[name] = tbl
+        # lancedb's .stats() stub claims TableStatistics but returns a plain dict at runtime.
+        tbl_stats: dict = tbl.stats()  # ty: ignore[invalid-assignment]
+        stats[name] = {
+            "exists": True,
+            "num_rows": tbl_stats.get("num_rows", 0),
+            "total_bytes": tbl_stats.get("total_bytes", 0),
+            "num_versions": len(list(tbl.list_versions())),
+        }
+
+    if stats["chunks"]["exists"]:
+        chunks_tbl = tables["chunks"]
+        indices = chunks_tbl.list_indices()
+        has_vector_index = any("vector" in str(idx).lower() for idx in indices)
+        stats["chunks"]["has_vector_index"] = has_vector_index
+        if has_vector_index:
+            index_stats = chunks_tbl.index_stats("vector_idx")
+            if index_stats is not None:
+                stats["chunks"]["num_indexed_rows"] = index_stats.num_indexed_rows
+                stats["chunks"]["num_unindexed_rows"] = index_stats.num_unindexed_rows
+
+    return stats
+
+
 class Store:
     def __init__(
         self,
@@ -277,51 +321,6 @@ class Store:
     def _connection_mode(self) -> ConnectionMode:
         return ConnectionMode.from_config(self._config)
 
-    def get_stats(self) -> dict:
-        """Get comprehensive table statistics.
-
-        Returns:
-            Dictionary with statistics for documents and chunks tables including:
-            - Row counts
-            - Storage sizes
-            - Vector index status and statistics
-        """
-        stats_dict: dict = {
-            "documents": {"exists": False},
-            "chunks": {"exists": False},
-        }
-
-        # Documents table stats
-        doc_stats: dict = self.documents_table.stats()
-        stats_dict["documents"] = {
-            "exists": True,
-            "num_rows": doc_stats.get("num_rows", 0),
-            "total_bytes": doc_stats.get("total_bytes", 0),
-        }
-
-        # Chunks table stats
-        chunk_stats: dict = self.chunks_table.stats()
-        stats_dict["chunks"] = {
-            "exists": True,
-            "num_rows": chunk_stats.get("num_rows", 0),
-            "total_bytes": chunk_stats.get("total_bytes", 0),
-        }
-
-        # Vector index stats
-        indices = self.chunks_table.list_indices()
-        has_vector_index = any("vector" in str(idx).lower() for idx in indices)
-        stats_dict["chunks"]["has_vector_index"] = has_vector_index
-
-        if has_vector_index:
-            index_stats = self.chunks_table.index_stats("vector_idx")
-            if index_stats is not None:
-                stats_dict["chunks"]["num_indexed_rows"] = index_stats.num_indexed_rows
-                stats_dict["chunks"]["num_unindexed_rows"] = (
-                    index_stats.num_unindexed_rows
-                )
-
-        return stats_dict
-
     def _ensure_vector_index(self) -> None:
         """Create or rebuild vector index on chunks table.
 
@@ -368,8 +367,7 @@ class Store:
     def _init_tables(self):
         """Initialize database tables (create if they don't exist)."""
         existing_tables = self.db.list_tables().tables
-        required_tables = {"documents", "chunks", "document_items", "settings"}
+        missing_tables = set(REQUIRED_TABLES) - set(existing_tables)
-        missing_tables = required_tables - set(existing_tables)
 
         if missing_tables and self._read_only:
             raise ReadOnlyError(

tests/store/test_engine.py

@@ -0,0 +1,103 @@
+import pytest
+
+from haiku.rag.store import Store
+from haiku.rag.store.engine import get_database_stats
+
+
+class TestGetDatabaseStats:
+    def test_empty_database_stats(self, temp_db_path):
+        """get_database_stats() on a fresh database reports zero rows and no vector index."""
+        store = Store(temp_db_path, create=True)
+
+        stats = 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
+
+        store.close()
+
+    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 = lancedb.connect(temp_db_path)
+        db.create_table("settings", schema=SettingsRecord)
+
+        stats = 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
+
+        store = Store(temp_db_path, create=True)
+        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 = get_database_stats(store.db)
+        assert stats["documents"]["num_rows"] == 1
+        assert stats["chunks"]["num_rows"] == 1
+        store.close()
+
+    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
+
+        store = Store(temp_db_path, create=True)
+        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)
+        ]
+        store.chunks_table.add(rows)
+        store.chunks_table.create_index(
+            metric="cosine", index_type="IVF_PQ", replace=True
+        )
+        store.chunks_table.wait_for_index(["vector_idx"], timeout=timedelta(minutes=1))
+
+        stats = 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"]
+        store.close()

tests/test_info.py

@@ -247,7 +247,7 @@ async def test_app_info_with_missing_document_items_table(temp_db_path, capsys):
     assert "document_items: absent" in out
 
     # Migration status should be surfaced
-    assert "Pending migrations:" in out
+    assert "migration(s) pending" in out
     assert "haiku-rag migrate" in out
 
 
@@ -308,7 +308,7 @@ async def test_app_info_reports_up_to_date(temp_db_path, capsys):
 
     out = capsys.readouterr().out
     assert "Database is up to date." in out
-    assert "Pending migrations" not in out
+    assert "migration(s) pending" not in out
 
 
 @pytest.mark.asyncio

tests/test_s3_integration.py

@@ -52,9 +52,11 @@ def _make_config() -> AppConfig:
 
 
 def test_store_connect_and_create(tmp_path):
+    from haiku.rag.store.engine import get_database_stats
+
     config = _make_config()
     store = Store(tmp_path / "unused", config=config, create=True)
-    stats = store.get_stats()
+    stats = get_database_stats(store.db)
     assert stats["documents"]["exists"]
     assert stats["chunks"]["exists"]
     store.close()
@@ -69,7 +71,7 @@ async def test_store_vacuum(tmp_path):
 
 
 def test_store_add_document(tmp_path):
-    from haiku.rag.store.engine import DocumentRecord
+    from haiku.rag.store.engine import DocumentRecord, get_database_stats
 
     config = _make_config()
     store = Store(tmp_path / "unused", config=config, create=True)
@@ -77,7 +79,7 @@ def test_store_add_document(tmp_path):
     doc = DocumentRecord(content="The quick brown fox jumps over the lazy dog.")
     store.documents_table.add([doc])
 
-    stats = store.get_stats()
+    stats = get_database_stats(store.db)
     assert stats["documents"]["num_rows"] == 1
     store.close()