coverage

haiku_rag_slim/haiku/rag/store/upgrades/v0_58_0.py

@@ -74,7 +74,9 @@ async def _apply_split_document_meta(store: Store) -> None:
         # lancedb's .stats() stub claims TableStatistics but returns a plain dict.
         stats: dict = await store.documents_table.stats()  # type: ignore[assignment]  # ty: ignore[invalid-assignment]
         live_bytes = int(stats.get("total_bytes", 0))
-    except Exception:
+    except (
+        Exception
+    ):  # pragma: no cover - defensive; stats() failure shouldn't block the split
         live_bytes = 0
     free_bytes = shutil.disk_usage(store.db_path).free
     if live_bytes and free_bytes < live_bytes:

tests/store/test_document_meta_split.py

@@ -106,3 +106,56 @@ async def test_migrate_creates_and_populates_document_meta(temp_db_path):
         repo = DocumentRepository(store)
         doc = await repo.get_by_id("d")
         assert doc is not None and doc.uri == "u"
+
+
+@pytest.mark.asyncio
+async def test_repository_empty_and_missing_id_paths(temp_db_path):
+    """Cover the repository's early-return branches for empty input and
+    missing ids, plus get_content."""
+    async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
+        repo = DocumentRepository(store)
+
+        assert await repo.create([]) == []
+        assert await repo.get_content("nope") is None
+        assert await repo.get_docling_data("nope") is None
+        assert await repo.get_pages_data("nope") is None
+        assert await repo.delete("nope") is False
+
+        doc = await repo.create(Document(content="hello body", uri="u1"))
+        assert doc.id is not None
+        assert await repo.get_content(doc.id) == "hello body"
+
+
+@pytest.mark.asyncio
+async def test_batch_create_rolls_back_meta_on_failure(temp_db_path, monkeypatch):
+    """The list create path also rolls back its document_meta rows if the
+    documents write fails."""
+    async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
+        repo = DocumentRepository(store)
+
+        async def boom(*_a, **_k):
+            raise RuntimeError("documents add failed")
+
+        monkeypatch.setattr(store.documents_table, "add", boom)
+        with pytest.raises(RuntimeError, match="documents add failed"):
+            await repo.create(
+                [Document(content="x", uri="u1"), Document(content="y", uri="u2")]
+            )
+        monkeypatch.undo()
+
+        assert await store.documents_table.count_rows() == 0
+        assert await store.document_meta_table.count_rows() == 0
+
+
+@pytest.mark.asyncio
+async def test_get_by_uri_with_orphan_meta_returns_none(temp_db_path):
+    """Defensive: a document_meta row whose documents row is missing (an
+    invariant violation) resolves to None, not a half-hydrated document."""
+    from haiku.rag.store.engine import DocumentMetaRecord
+
+    async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
+        repo = DocumentRepository(store)
+        await store.document_meta_table.add(
+            [DocumentMetaRecord(document_id="ghost", uri="u-ghost", metadata="{}")]
+        )
+        assert await repo.get_by_uri("u-ghost") is None

tests/store/test_v0_58_0_migration.py

@@ -96,3 +96,101 @@ class TestV0_58_0Migration:
             meta_rows = await store.document_meta_table.query().to_list()
             assert len(meta_rows) == 1
             assert meta_rows[0]["document_id"] == "doc-1"
+
+
+@pytest.mark.asyncio
+class TestV0_58_0MigrationEdgeCases:
+    async def test_resume_skips_already_migrated_rows(self, temp_db_path):
+        """A half-finished prior run leaves some document_meta rows; re-running
+        migrates only the rest and never duplicates."""
+        from haiku.rag.store.engine import DocumentMetaRecord
+
+        async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
+            await seed_legacy_documents(
+                store,
+                [
+                    LegacyDocumentRecord(id="a", content="x", uri="u-a", metadata="{}"),
+                    LegacyDocumentRecord(id="b", content="y", uri="u-b", metadata="{}"),
+                ],
+            )
+            # Pretend a prior run already moved doc "a".
+            await store.document_meta_table.add(
+                [DocumentMetaRecord(document_id="a", uri="u-a", metadata="{}")]
+            )
+            await store.set_haiku_version("0.57.0")
+
+        async with Store(temp_db_path, skip_migration_check=True) as store:
+            await store.migrate()
+            rows = await store.document_meta_table.query().to_list()
+            by_id = {r["document_id"]: r for r in rows}
+            assert set(by_id) == {"a", "b"}  # exactly one row each, no duplicates
+            assert len(rows) == 2
+
+    async def test_skips_vacuum_when_disk_is_tight(self, temp_db_path, monkeypatch):
+        """When free disk can't cover one compacted copy, the split still
+        completes but the reclaim vacuum is skipped."""
+        from types import SimpleNamespace
+
+        from haiku.rag.store.upgrades import v0_58_0
+
+        async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
+            await seed_legacy_documents(
+                store,
+                [LegacyDocumentRecord(id="a", content="x", uri="u", metadata="{}")],
+            )
+            await store.set_haiku_version("0.57.0")
+
+        # Pretend almost no free disk so the reclaim vacuum is skipped.
+        monkeypatch.setattr(
+            v0_58_0.shutil,
+            "disk_usage",
+            lambda _p: SimpleNamespace(total=1, used=1, free=1),
+        )
+
+        vacuum_calls: list[int] = []
+
+        async with Store(temp_db_path, skip_migration_check=True) as store:
+            # Force a known nonzero live size so the free<live check is
+            # deterministic (real stats().total_bytes can be 0 for a tiny table).
+            async def fake_stats():
+                return {"total_bytes": 10_000_000}
+
+            monkeypatch.setattr(store.documents_table, "stats", fake_stats)
+
+            orig_vacuum = store.vacuum
+
+            async def tracking_vacuum(*args, **kwargs):
+                vacuum_calls.append(1)
+                return await orig_vacuum(*args, **kwargs)
+
+            monkeypatch.setattr(store, "vacuum", tracking_vacuum)
+
+            await store.migrate()
+            # Split still happened despite the skipped vacuum.
+            names = {f.name for f in await store.documents_table.schema()}
+            assert _LEGACY_COLUMNS.isdisjoint(names)
+            repo = DocumentRepository(store)
+            migrated = await repo.get_by_id("a")
+            assert migrated is not None and migrated.uri == "u"
+
+        assert vacuum_calls == []  # reclaim vacuum was skipped
+
+
+@pytest.mark.asyncio
+async def test_delete_all_clears_document_meta(temp_db_path):
+    """delete_all drops and recreates both documents and document_meta."""
+    from haiku.rag.store.models.document import Document
+
+    async with Store(temp_db_path, create=True) as store:
+        repo = DocumentRepository(store)
+        await repo.create(Document(content="x", uri="u1"))
+        await repo.create(Document(content="y", uri="u2"))
+        assert await store.document_meta_table.count_rows() == 2
+
+        await repo.delete_all()
+
+        assert await store.documents_table.count_rows() == 0
+        assert await store.document_meta_table.count_rows() == 0
+        # Tables are usable again after recreation.
+        await repo.create(Document(content="z", uri="u3"))
+        assert await store.document_meta_table.count_rows() == 1