From 3fe953eccba6d2dc8653f7af202b6cbea2d99ea1 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Wed, 8 Apr 2026 12:53:09 +0300
Subject: [PATCH] Minor fixes

---
 haiku_rag_slim/haiku/rag/store/compression.py |  2 +-
 .../haiku/rag/store/models/document.py        |  6 +-
 .../haiku/rag/store/upgrades/v0_38_0.py       |  8 +--
 tests/test_document.py                        | 57 +++++++++++++++++++
 4 files changed, 64 insertions(+), 9 deletions(-)

diff --git a/haiku_rag_slim/haiku/rag/store/compression.py b/haiku_rag_slim/haiku/rag/store/compression.py
index 7e2b56ef..80abf712 100644
--- a/haiku_rag_slim/haiku/rag/store/compression.py
+++ b/haiku_rag_slim/haiku/rag/store/compression.py
@@ -1,6 +1,6 @@
 import json
 
-try:
+try:  # pragma: no cover
     from compression.zstd import (  # ty: ignore[unresolved-import]
         compress as _zstd_compress,  # type: ignore[import-not-found]
     )
diff --git a/haiku_rag_slim/haiku/rag/store/models/document.py b/haiku_rag_slim/haiku/rag/store/models/document.py
index 024205cd..a7e406f5 100644
--- a/haiku_rag_slim/haiku/rag/store/models/document.py
+++ b/haiku_rag_slim/haiku/rag/store/models/document.py
@@ -15,13 +15,11 @@ _docling_document_cache: LRUCache[str, "DoclingDocument"] = LRUCache(maxsize=100
 
 
 def _validate_without_pages(compressed_data: bytes) -> "DoclingDocument":
-    """Decompress and validate DoclingDocument, stripping page images."""
+    """Decompress and validate DoclingDocument."""
     from docling_core.types.doc.document import DoclingDocument
 
     json_str = decompress_json(compressed_data)
-    data = json.loads(json_str)
-    data.pop("pages", None)
-    return DoclingDocument.model_validate(data)
+    return DoclingDocument.model_validate_json(json_str)
 
 
 def _get_cached_docling_document(
diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_38_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_38_0.py
index 7b255348..0c28ed77 100644
--- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_38_0.py
+++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_38_0.py
@@ -106,7 +106,7 @@ def _apply_split_pages_zstd(store: Store) -> None:  # pragma: no cover
             .to_arrow()
             .to_pylist()
         ]
-    except Exception:
+    except (pa.ArrowInvalid, pa.ArrowNotImplementedError, OSError):
         ids = []
 
     if not ids:
@@ -132,7 +132,7 @@ def _apply_split_pages_zstd(store: Store) -> None:  # pragma: no cover
                     range(0, len(staging_ids), BATCH_SIZE), 1
                 ):
                     batch_ids = staging_ids[i : i + BATCH_SIZE]
-                    id_list = ", ".join(f"'{id}'" for id in batch_ids)
+                    id_list = ", ".join(f"'{doc_id}'" for doc_id in batch_ids)
                     batch = (
                         staging_table.search()
                         .where(f"id IN ({id_list})")
@@ -174,7 +174,7 @@ def _apply_split_pages_zstd(store: Store) -> None:  # pragma: no cover
 
     for batch_num, i in enumerate(range(0, len(ids), BATCH_SIZE), 1):
         batch_ids = ids[i : i + BATCH_SIZE]
-        id_list = ", ".join(f"'{id}'" for id in batch_ids)
+        id_list = ", ".join(f"'{doc_id}'" for doc_id in batch_ids)
 
         batch = (
             store.documents_table.search()
@@ -212,7 +212,7 @@ def _apply_split_pages_zstd(store: Store) -> None:  # pragma: no cover
 
     for batch_num, i in enumerate(range(0, len(staging_ids), BATCH_SIZE), 1):
         batch_ids = staging_ids[i : i + BATCH_SIZE]
-        id_list = ", ".join(f"'{id}'" for id in batch_ids)
+        id_list = ", ".join(f"'{doc_id}'" for doc_id in batch_ids)
 
         batch = (
             staging_table.search().where(f"id IN ({id_list})").to_arrow().to_pylist()
diff --git a/tests/test_document.py b/tests/test_document.py
index 1ce394b8..635c31d8 100644
--- a/tests/test_document.py
+++ b/tests/test_document.py
@@ -367,6 +367,63 @@ async def test_get_docling_data_loads_only_docling_columns(
     store.close()
 
 
+@pytest.mark.asyncio
+async def test_get_pages_data_loads_only_pages_column(qa_corpus: Dataset, temp_db_path):
+    """get_pages_data returns only page image data for a document."""
+    import json
+
+    from haiku.rag.store.compression import compress_json
+
+    pages_blob = compress_json(
+        json.dumps({"1": {"size": {"width": 612, "height": 792}, "page_no": 1}})
+    )
+
+    store = Store(temp_db_path, create=True)
+    doc_repo = DocumentRepository(store)
+
+    doc = Document(
+        content=qa_corpus[0]["document_extracted"],
+        uri="https://example.com/doc.txt",
+        docling_pages=pages_blob,
+    )
+    created = await doc_repo.create(doc)
+    assert created.id is not None
+
+    result = await doc_repo.get_pages_data(created.id)
+    assert result is not None
+    assert result.id == created.id
+    assert result.content == ""
+    assert result.docling_pages == pages_blob
+
+    # Non-existent ID returns None
+    assert await doc_repo.get_pages_data("nonexistent-id") is None
+
+    store.close()
+
+
+@pytest.mark.asyncio
+async def test_get_pages_data_none_for_markdown_document(
+    qa_corpus: Dataset, temp_db_path
+):
+    """Markdown documents have no page images — get_pages_data returns None pages."""
+    store = Store(temp_db_path, create=True)
+    doc_repo = DocumentRepository(store)
+
+    doc = Document(
+        content=qa_corpus[0]["document_extracted"],
+        uri="https://example.com/doc.md",
+    )
+    created = await doc_repo.create(doc)
+    assert created.id is not None
+
+    result = await doc_repo.get_pages_data(created.id)
+    assert result is not None
+    assert result.id == created.id
+    assert result.docling_pages is None
+
+    store.close()
+
+
 @pytest.mark.asyncio
 async def test_document_get_by_uri_with_special_characters(
     qa_corpus: Dataset, temp_db_path