From f5b31f809fa8e2b74ce640a1ac8acec163571d9c Mon Sep 17 00:00:00 2001 From: Pier-Jean Malandrino Date: Thu, 2 Apr 2026 12:05:12 +0200 Subject: [PATCH] Add document_json and chunks_json persistence layer Schema migration for new columns, repository read/write support, and dedicated update_chunks method for re-chunking operations. --- document-parser/persistence/analysis_repo.py | 44 ++++++++++++++------ document-parser/persistence/database.py | 20 +++++++++ 2 files changed, 51 insertions(+), 13 deletions(-) diff --git a/document-parser/persistence/analysis_repo.py b/document-parser/persistence/analysis_repo.py index f8b3df7..0652e5b 100644 --- a/document-parser/persistence/analysis_repo.py +++ b/document-parser/persistence/analysis_repo.py @@ -16,6 +16,7 @@ def _parse_dt(value: str | None) -> datetime | None: def _row_to_job(row) -> AnalysisJob: + keys = row.keys() return AnalysisJob( id=row["id"], document_id=row["document_id"], @@ -23,11 +24,13 @@ def _row_to_job(row) -> AnalysisJob: content_markdown=row["content_markdown"], content_html=row["content_html"], pages_json=row["pages_json"], + document_json=row["document_json"] if "document_json" in keys else None, + chunks_json=row["chunks_json"] if "chunks_json" in keys else None, error_message=row["error_message"], started_at=_parse_dt(row["started_at"]), completed_at=_parse_dt(row["completed_at"]), created_at=_parse_dt(row["created_at"]) or datetime.now(), - document_filename=row["filename"] if "filename" in row.keys() else None, # noqa: SIM118 + document_filename=row["filename"] if "filename" in keys else None, ) @@ -60,9 +63,7 @@ async def find_all(*, limit: int = 200, offset: int = 0) -> list[AnalysisJob]: async def find_by_id(job_id: str) -> AnalysisJob | None: async with get_connection() as db: - cursor = await db.execute( - f"{_SELECT_WITH_DOC} WHERE aj.id = ?", (job_id,) - ) + cursor = await db.execute(f"{_SELECT_WITH_DOC} WHERE aj.id = ?", (job_id,)) row = await cursor.fetchone() return _row_to_job(row) if row else None @@ -72,17 +73,36 @@ async def update_status(job: AnalysisJob) -> None: await db.execute( """UPDATE analysis_jobs SET status = ?, content_markdown = ?, content_html = ?, - pages_json = ?, error_message = ?, started_at = ?, completed_at = ? + pages_json = ?, document_json = ?, chunks_json = ?, + error_message = ?, started_at = ?, completed_at = ? WHERE id = ?""", - (job.status.value, job.content_markdown, job.content_html, - job.pages_json, job.error_message, - str(job.started_at) if job.started_at else None, - str(job.completed_at) if job.completed_at else None, - job.id), + ( + job.status.value, + job.content_markdown, + job.content_html, + job.pages_json, + job.document_json, + job.chunks_json, + job.error_message, + str(job.started_at) if job.started_at else None, + str(job.completed_at) if job.completed_at else None, + job.id, + ), ) await db.commit() +async def update_chunks(job_id: str, chunks_json: str) -> bool: + """Update only the chunks_json column for a completed analysis.""" + async with get_connection() as db: + cursor = await db.execute( + "UPDATE analysis_jobs SET chunks_json = ? WHERE id = ?", + (chunks_json, job_id), + ) + await db.commit() + return cursor.rowcount > 0 + + async def delete(job_id: str) -> bool: async with get_connection() as db: cursor = await db.execute("DELETE FROM analysis_jobs WHERE id = ?", (job_id,)) @@ -93,8 +113,6 @@ async def delete(job_id: str) -> bool: async def delete_by_document(document_id: str) -> int: """Delete all analysis jobs for a given document. Returns count deleted.""" async with get_connection() as db: - cursor = await db.execute( - "DELETE FROM analysis_jobs WHERE document_id = ?", (document_id,) - ) + cursor = await db.execute("DELETE FROM analysis_jobs WHERE document_id = ?", (document_id,)) await db.commit() return cursor.rowcount diff --git a/document-parser/persistence/database.py b/document-parser/persistence/database.py index 6aaa37b..93c7d94 100644 --- a/document-parser/persistence/database.py +++ b/document-parser/persistence/database.py @@ -30,6 +30,8 @@ CREATE TABLE IF NOT EXISTS analysis_jobs ( content_markdown TEXT, content_html TEXT, pages_json TEXT, + document_json TEXT, + chunks_json TEXT, error_message TEXT, started_at TEXT, completed_at TEXT, @@ -42,11 +44,29 @@ CREATE INDEX IF NOT EXISTS idx_documents_created_at ON documents(created_at); """ +_MIGRATIONS = [ + ("document_json", "ALTER TABLE analysis_jobs ADD COLUMN document_json TEXT"), + ("chunks_json", "ALTER TABLE analysis_jobs ADD COLUMN chunks_json TEXT"), +] + + +async def _run_migrations(db: aiosqlite.Connection) -> None: + """Add columns that may be missing in older databases.""" + cursor = await db.execute("PRAGMA table_info(analysis_jobs)") + existing = {row[1] for row in await cursor.fetchall()} + for col_name, ddl in _MIGRATIONS: + if col_name not in existing: + await db.execute(ddl) + logger.info("Migration: added column %s to analysis_jobs", col_name) + await db.commit() + + async def init_db() -> None: """Create database file and tables if they don't exist.""" os.makedirs(os.path.dirname(DB_PATH) or ".", exist_ok=True) async with aiosqlite.connect(DB_PATH) as db: await db.executescript(_SCHEMA) + await _run_migrations(db) await db.commit() logger.info("Database initialized at %s", DB_PATH)