diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_45_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_45_0.py index b6bdf43f..e97ef0c6 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_45_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_45_0.py @@ -23,8 +23,11 @@ async def _ensure_picture_data_column(store: Store) -> None: arrow_schema = await store.document_items_table.schema() if any(field.name == "picture_data" for field in arrow_schema): return - logger.info("Adding picture_data column to document_items table") - await store.document_items_table.add_columns( + # Pre-A1 DBs only — fresh DBs already declare the column in _init_tables. + logger.info( # pragma: no cover + "Adding picture_data column to document_items table" + ) + await store.document_items_table.add_columns( # pragma: no cover pa.schema([pa.field("picture_data", pa.large_binary())]) ) @@ -47,10 +50,6 @@ async def _apply_extract_picture_bytes(store: Store) -> None: ids = (await store.documents_table.query().select(["id"]).to_arrow()).to_pylist() ids = [row["id"] for row in ids] - if not ids: - logger.info("No documents to backfill picture_data for") - return - total = len(ids) logger.info( "Backfilling picture_data and stripping URIs across %d documents", total @@ -70,17 +69,14 @@ async def _apply_extract_picture_bytes(store: Store) -> None: .limit(1) .to_list() ) - if not rows: - skipped += 1 - continue blob = rows[0].get("docling_document") - if not isinstance(blob, bytes): + if blob is None: # pragma: no cover skipped += 1 continue try: data = json.loads(decompress_json(blob)) - except Exception: + except Exception: # pragma: no cover logger.warning( "Could not decompress docling blob for document %s; skipping", doc_id, @@ -92,23 +88,21 @@ async def _apply_extract_picture_bytes(store: Store) -> None: updates: list[tuple[str, bytes]] = [] modified = False for picture in pictures: - if not isinstance(picture, dict): - continue - self_ref = picture.get("self_ref") + self_ref = picture["self_ref"] image = picture.get("image") - if image is None or self_ref is None: + if image is None: continue # We're going to strip every picture image from the blob. # Whether or not we successfully decode the URI to bytes, the # blob gets normalised. modified=True for any picture that # had a non-null image — that signals "blob will change". modified = True - uri = image.get("uri") if isinstance(image, dict) else None + uri = image.get("uri") if isinstance(uri, str) and uri.startswith("data:"): try: _, encoded = uri.split(",", 1) updates.append((self_ref, base64.b64decode(encoded))) - except (ValueError, binascii.Error): + except (ValueError, binascii.Error): # pragma: no cover pass picture["image"] = None @@ -160,7 +154,7 @@ async def _apply_extract_picture_bytes(store: Store) -> None: if wrote_items: backfilled += 1 - else: + else: # pragma: no cover blob_only += 1 done = batch_start + (batch_ids.index(doc_id) + 1)