chore(#206): backfill 0.6.0 doc-centric data model
CLI script to migrate existing tenants to the data model introduced
by #202-205. Idempotent and resumable — re-running is a no-op.
tools/migrate_06.py
- Step 1: backfill_lifecycle — infers documents.lifecycle_state from
analysis_jobs.status + chunks_json presence
Failed analyses, no completed -> Failed
Completed + chunks_json -> Chunked
Completed, no chunks_json -> Parsed
Otherwise -> Uploaded
- Step 2: materialize_chunks — promotes the latest chunks_json blob
for each doc into rows in the new chunks table. Stable id derivation
via uuid5(NAMESPACE_OID, '<doc>|<seq>|<text>') so re-running lands
on the same ids
- Step 3: backfill_links — for any doc that has chunks materialized,
creates a (doc, default-store) link in Ingested state with a freshly
computed chunkset_hash. Treats 'chunks exist' as proxy for 'doc was
ingested into the legacy single index' — sufficient for the typical
pre-0.6.0 deployment, with OpenSearch reindex documented separately
- Step 4: reaggregate — applies #203's aggregate_lifecycle rule so
doc-chunked transitions to Ingested
Persistence
- migration_progress table for resumability + per-step idempotency
CLI
python -m tools.migrate_06 # full migration
python -m tools.migrate_06 --dry-run # plan only, no writes
python -m tools.migrate_06 --only-step <name> # rerun a phase
Tests
- 9 tests: inference rule per (completed, failed, chunks_json) tuple,
end-to-end migration on a hand-built three-doc snapshot, idempotency
(second run = no writes), --dry-run writes nothing, deterministic
chunk ids across reruns
Refs #206
This commit is contained in:
Pier-Jean Malandrino
parent
648f2a5c0d
commit
0c09274836
3 changed files with 624 additions and 0 deletions
|
|
@ -114,6 +114,12 @@ CREATE TABLE IF NOT EXISTS chunk_pushes (
|
|||
pushed_at TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_chunk_pushes_doc_store ON chunk_pushes(document_id, store_id);
|
||||
|
||||
-- 0.6.0 — migration progress (resumability for the 0.6.0 backfill, #206).
|
||||
CREATE TABLE IF NOT EXISTS migration_progress (
|
||||
name TEXT PRIMARY KEY,
|
||||
completed_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
"""
|
||||
|
||||
|
||||
|
|
|
|||
185
document-parser/tests/test_migrate_06.py
Normal file
185
document-parser/tests/test_migrate_06.py
Normal file
|
|
@ -0,0 +1,185 @@
|
|||
"""Tests for the 0.6.0 migration script (#206)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from domain.value_objects import DocumentLifecycleState, DocumentStoreLinkState
|
||||
from persistence.chunk_repo import SqliteChunkRepository
|
||||
from persistence.database import get_connection, init_db
|
||||
from persistence.document_repo import SqliteDocumentRepository
|
||||
from persistence.document_store_link_repo import SqliteDocumentStoreLinkRepository
|
||||
from tools.migrate_06 import _DocFacts, _infer_lifecycle, run
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
async def setup_db(monkeypatch, tmp_path):
|
||||
db_path = str(tmp_path / "test.db")
|
||||
monkeypatch.setattr("persistence.database.DB_PATH", db_path)
|
||||
await init_db()
|
||||
yield
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Inference rule unit tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("has_completed", "has_failed", "has_chunks_json", "expected"),
|
||||
[
|
||||
(False, False, False, DocumentLifecycleState.UPLOADED),
|
||||
(False, True, False, DocumentLifecycleState.FAILED),
|
||||
(True, False, False, DocumentLifecycleState.PARSED),
|
||||
(True, False, True, DocumentLifecycleState.CHUNKED),
|
||||
(True, True, True, DocumentLifecycleState.CHUNKED), # any completed wins
|
||||
],
|
||||
)
|
||||
def test_infer_lifecycle(
|
||||
has_completed: bool,
|
||||
has_failed: bool,
|
||||
has_chunks_json: bool,
|
||||
expected: DocumentLifecycleState,
|
||||
) -> None:
|
||||
facts = _DocFacts(
|
||||
document_id="doc",
|
||||
has_completed=has_completed,
|
||||
has_failed=has_failed,
|
||||
has_chunks_json=has_chunks_json,
|
||||
)
|
||||
assert _infer_lifecycle(facts) == expected
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Integration: end-to-end migration on a hand-built pre-state
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _seed_pre_06_state() -> None:
|
||||
"""Seed a tenant snapshot with three documents covering the three
|
||||
main inference cases:
|
||||
|
||||
doc-uploaded — only an `Uploaded` doc row, nothing else
|
||||
doc-parsed — completed analysis, no chunks
|
||||
doc-chunked — completed analysis with chunks_json
|
||||
"""
|
||||
repo = SqliteDocumentRepository()
|
||||
async with get_connection() as db:
|
||||
await db.execute(
|
||||
"INSERT INTO documents (id, filename, storage_path, created_at) "
|
||||
"VALUES ('doc-uploaded', 'a.pdf', '/tmp/a.pdf', datetime('now'))"
|
||||
)
|
||||
await db.execute(
|
||||
"INSERT INTO documents (id, filename, storage_path, created_at) "
|
||||
"VALUES ('doc-parsed', 'b.pdf', '/tmp/b.pdf', datetime('now'))"
|
||||
)
|
||||
await db.execute(
|
||||
"INSERT INTO documents (id, filename, storage_path, created_at) "
|
||||
"VALUES ('doc-chunked', 'c.pdf', '/tmp/c.pdf', datetime('now'))"
|
||||
)
|
||||
await db.execute(
|
||||
"INSERT INTO analysis_jobs (id, document_id, status) "
|
||||
"VALUES ('a-parsed', 'doc-parsed', 'COMPLETED')"
|
||||
)
|
||||
chunks_json = json.dumps(
|
||||
[
|
||||
{"text": "Intro paragraph.", "headings": ["Intro"], "sourcePage": 1},
|
||||
{"text": "Body of section A.", "headings": ["A"], "sourcePage": 2},
|
||||
]
|
||||
)
|
||||
await db.execute(
|
||||
"INSERT INTO analysis_jobs (id, document_id, status, chunks_json) "
|
||||
"VALUES ('a-chunked', 'doc-chunked', 'COMPLETED', ?)",
|
||||
(chunks_json,),
|
||||
)
|
||||
await db.commit()
|
||||
# Touch the document repo so the type checker is satisfied.
|
||||
assert await repo.find_by_id("doc-chunked") is not None
|
||||
|
||||
|
||||
async def test_full_migration_writes_expected_state() -> None:
|
||||
await _seed_pre_06_state()
|
||||
|
||||
results = await run(dry_run=False)
|
||||
by_step = {r.name: r for r in results}
|
||||
assert by_step["backfill_lifecycle"].written >= 2 # parsed + chunked transition
|
||||
assert by_step["materialize_chunks"].written == 2 # two chunks
|
||||
assert by_step["backfill_links"].written == 1 # only doc-chunked has chunks
|
||||
|
||||
repo = SqliteDocumentRepository()
|
||||
chunks_repo = SqliteChunkRepository()
|
||||
link_repo = SqliteDocumentStoreLinkRepository()
|
||||
|
||||
uploaded = await repo.find_by_id("doc-uploaded")
|
||||
assert uploaded is not None
|
||||
assert uploaded.lifecycle_state == DocumentLifecycleState.UPLOADED
|
||||
|
||||
parsed = await repo.find_by_id("doc-parsed")
|
||||
assert parsed is not None
|
||||
assert parsed.lifecycle_state == DocumentLifecycleState.PARSED
|
||||
|
||||
chunked = await repo.find_by_id("doc-chunked")
|
||||
assert chunked is not None
|
||||
# After link backfill + reaggregation, doc-chunked is Ingested in default store.
|
||||
assert chunked.lifecycle_state == DocumentLifecycleState.INGESTED
|
||||
|
||||
chunks = await chunks_repo.find_for_document("doc-chunked")
|
||||
assert len(chunks) == 2
|
||||
assert chunks[0].text == "Intro paragraph."
|
||||
assert chunks[0].sequence == 0
|
||||
|
||||
links = await link_repo.find_for_document("doc-chunked")
|
||||
assert len(links) == 1
|
||||
assert links[0].store_id == "default"
|
||||
assert links[0].state == DocumentStoreLinkState.INGESTED
|
||||
assert links[0].chunkset_hash is not None
|
||||
|
||||
|
||||
async def test_migration_is_idempotent() -> None:
|
||||
await _seed_pre_06_state()
|
||||
|
||||
first = await run(dry_run=False)
|
||||
second = await run(dry_run=False)
|
||||
# Second run sees every step as already-done.
|
||||
assert all(r.skipped for r in second)
|
||||
assert all(not r.skipped for r in first)
|
||||
|
||||
|
||||
async def test_dry_run_writes_nothing() -> None:
|
||||
await _seed_pre_06_state()
|
||||
|
||||
repo = SqliteDocumentRepository()
|
||||
chunks_repo = SqliteChunkRepository()
|
||||
|
||||
await run(dry_run=True)
|
||||
|
||||
# No state changed.
|
||||
parsed = await repo.find_by_id("doc-parsed")
|
||||
assert parsed is not None
|
||||
assert parsed.lifecycle_state == DocumentLifecycleState.UPLOADED # unchanged
|
||||
chunked_chunks = await chunks_repo.find_for_document("doc-chunked")
|
||||
assert chunked_chunks == []
|
||||
|
||||
|
||||
async def test_chunk_ids_are_deterministic_across_runs() -> None:
|
||||
"""Stable id derivation lets a re-run after a manual reset land on
|
||||
the same chunk ids — important for the link-table foreign keys."""
|
||||
await _seed_pre_06_state()
|
||||
|
||||
chunks_repo = SqliteChunkRepository()
|
||||
|
||||
await run(dry_run=False, only_step="materialize_chunks")
|
||||
first_ids = sorted(c.id for c in await chunks_repo.find_for_document("doc-chunked"))
|
||||
|
||||
# Wipe progress + chunks, run again — same ids.
|
||||
async with get_connection() as db:
|
||||
await db.execute("DELETE FROM migration_progress WHERE name = 'materialize_chunks'")
|
||||
await db.execute("DELETE FROM chunks WHERE document_id = 'doc-chunked'")
|
||||
await db.commit()
|
||||
|
||||
await run(dry_run=False, only_step="materialize_chunks")
|
||||
second_ids = sorted(c.id for c in await chunks_repo.find_for_document("doc-chunked"))
|
||||
|
||||
assert first_ids == second_ids
|
||||
433
document-parser/tools/migrate_06.py
Normal file
433
document-parser/tools/migrate_06.py
Normal file
|
|
@ -0,0 +1,433 @@
|
|||
"""0.6.0 migration — backfill the doc-centric data model.
|
||||
|
||||
Run after deploying the 0.6.0 release code, before sending traffic. The
|
||||
script is idempotent and resumable — re-running on an already-migrated
|
||||
database is a no-op. Each step records itself in `migration_progress`.
|
||||
|
||||
Inference rules for `documents.lifecycle_state`:
|
||||
|
||||
has any FAILED analysis_jobs row, no completed → Failed
|
||||
has any COMPLETED analysis_jobs row + chunks_json → Chunked
|
||||
has any COMPLETED analysis_jobs row, no chunks_json → Parsed
|
||||
else → Uploaded
|
||||
|
||||
After chunks are materialized into the new `chunks` table and links are
|
||||
backfilled (one per existing default-store ingestion), the document
|
||||
state is recomputed using the standard #203 aggregation rule. This
|
||||
covers Stale / Ingested upgrades where applicable.
|
||||
|
||||
Usage:
|
||||
|
||||
python -m tools.migrate_06 [--dry-run] [--only-step <name>]
|
||||
|
||||
The CLI is intentionally minimal — see docs/design/206-lifecycle-state-migration.md
|
||||
for the full set of flags considered.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from domain.hashing import chunkset_hash
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Iterable
|
||||
|
||||
import aiosqlite
|
||||
from domain.lifecycle_aggregation import aggregate_lifecycle
|
||||
from domain.models import Chunk, DocumentStoreLink
|
||||
from domain.value_objects import (
|
||||
ChunkBbox,
|
||||
ChunkDocItem,
|
||||
ChunkResult,
|
||||
DocumentLifecycleState,
|
||||
DocumentStoreLinkState,
|
||||
)
|
||||
from persistence.chunk_repo import SqliteChunkRepository
|
||||
from persistence.database import get_connection, init_db
|
||||
from persistence.document_store_link_repo import SqliteDocumentStoreLinkRepository
|
||||
from persistence.store_repo import SqliteStoreRepository
|
||||
|
||||
logger = logging.getLogger("migrate_06")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Progress tracking
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _is_done(db: aiosqlite.Connection, name: str) -> bool:
|
||||
cursor = await db.execute("SELECT 1 FROM migration_progress WHERE name = ?", (name,))
|
||||
return await cursor.fetchone() is not None
|
||||
|
||||
|
||||
async def _mark_done(db: aiosqlite.Connection, name: str) -> None:
|
||||
await db.execute(
|
||||
"INSERT OR IGNORE INTO migration_progress (name) VALUES (?)",
|
||||
(name,),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Step: backfill_document_lifecycle_state
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class _DocFacts:
|
||||
document_id: str
|
||||
has_completed: bool
|
||||
has_failed: bool
|
||||
has_chunks_json: bool
|
||||
|
||||
|
||||
async def _gather_document_facts(db: aiosqlite.Connection) -> list[_DocFacts]:
|
||||
"""One row per document with the analysis-state booleans we care about."""
|
||||
cursor = await db.execute(
|
||||
"""SELECT d.id AS document_id,
|
||||
MAX(CASE WHEN a.status = 'COMPLETED' THEN 1 ELSE 0 END) AS has_completed,
|
||||
MAX(CASE WHEN a.status = 'FAILED' THEN 1 ELSE 0 END) AS has_failed,
|
||||
MAX(CASE WHEN a.chunks_json IS NOT NULL AND a.chunks_json != ''
|
||||
THEN 1 ELSE 0 END) AS has_chunks_json
|
||||
FROM documents d
|
||||
LEFT JOIN analysis_jobs a ON a.document_id = d.id
|
||||
GROUP BY d.id"""
|
||||
)
|
||||
rows = await cursor.fetchall()
|
||||
return [
|
||||
_DocFacts(
|
||||
document_id=r["document_id"],
|
||||
has_completed=bool(r["has_completed"]),
|
||||
has_failed=bool(r["has_failed"]),
|
||||
has_chunks_json=bool(r["has_chunks_json"]),
|
||||
)
|
||||
for r in rows
|
||||
]
|
||||
|
||||
|
||||
def _infer_lifecycle(facts: _DocFacts) -> DocumentLifecycleState:
|
||||
if facts.has_completed and facts.has_chunks_json:
|
||||
return DocumentLifecycleState.CHUNKED
|
||||
if facts.has_completed:
|
||||
return DocumentLifecycleState.PARSED
|
||||
if facts.has_failed:
|
||||
return DocumentLifecycleState.FAILED
|
||||
return DocumentLifecycleState.UPLOADED
|
||||
|
||||
|
||||
async def _step_backfill_lifecycle(db: aiosqlite.Connection, *, dry_run: bool) -> int:
|
||||
facts = await _gather_document_facts(db)
|
||||
written = 0
|
||||
for f in facts:
|
||||
target = _infer_lifecycle(f)
|
||||
cursor = await db.execute(
|
||||
"SELECT lifecycle_state FROM documents WHERE id = ?",
|
||||
(f.document_id,),
|
||||
)
|
||||
row = await cursor.fetchone()
|
||||
if row and row["lifecycle_state"] == target.value:
|
||||
continue
|
||||
logger.info(
|
||||
"lifecycle: doc=%s -> %s",
|
||||
f.document_id,
|
||||
target.value,
|
||||
)
|
||||
if not dry_run:
|
||||
await db.execute(
|
||||
"UPDATE documents SET lifecycle_state = ?, "
|
||||
"lifecycle_state_at = datetime('now') WHERE id = ?",
|
||||
(target.value, f.document_id),
|
||||
)
|
||||
written += 1
|
||||
if not dry_run:
|
||||
await db.commit()
|
||||
return written
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Step: materialize_chunks_from_chunks_json
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _stable_chunk_id(document_id: str, sequence: int, text: str) -> str:
|
||||
"""Deterministic id so re-running the migration doesn't duplicate rows."""
|
||||
seed = f"{document_id}|{sequence}|{text}"
|
||||
return uuid.uuid5(uuid.NAMESPACE_OID, seed).hex
|
||||
|
||||
|
||||
def _chunk_dicts_for_doc(
|
||||
db: aiosqlite.Connection,
|
||||
) -> Callable[[str], asyncio.Future[list[dict]]]:
|
||||
"""Return an async-compatible accessor over the latest chunks_json
|
||||
for a document. Returns an empty list when no chunks_json is available.
|
||||
"""
|
||||
|
||||
async def _read(document_id: str) -> list[dict]:
|
||||
cursor = await db.execute(
|
||||
"""SELECT chunks_json FROM analysis_jobs
|
||||
WHERE document_id = ? AND chunks_json IS NOT NULL AND chunks_json != ''
|
||||
ORDER BY completed_at DESC, created_at DESC
|
||||
LIMIT 1""",
|
||||
(document_id,),
|
||||
)
|
||||
row = await cursor.fetchone()
|
||||
if row is None:
|
||||
return []
|
||||
try:
|
||||
return list(json.loads(row["chunks_json"]))
|
||||
except (TypeError, ValueError):
|
||||
logger.warning("Invalid chunks_json for doc %s — skipping", document_id)
|
||||
return []
|
||||
|
||||
return _read
|
||||
|
||||
|
||||
def _chunk_dict_to_chunk(d: dict, *, document_id: str, sequence: int) -> Chunk:
|
||||
text = d.get("text", "")
|
||||
return Chunk(
|
||||
id=_stable_chunk_id(document_id, sequence, text),
|
||||
document_id=document_id,
|
||||
sequence=sequence,
|
||||
text=text,
|
||||
headings=list(d.get("headings", [])),
|
||||
source_page=d.get("sourcePage"),
|
||||
bboxes=[
|
||||
ChunkBbox(page=b.get("page", 0), bbox=list(b.get("bbox", [])))
|
||||
for b in d.get("bboxes", [])
|
||||
],
|
||||
doc_items=[
|
||||
ChunkDocItem(
|
||||
self_ref=di.get("selfRef", ""),
|
||||
label=di.get("label", ""),
|
||||
)
|
||||
for di in d.get("docItems", [])
|
||||
],
|
||||
token_count=d.get("tokenCount"),
|
||||
)
|
||||
|
||||
|
||||
async def _step_materialize_chunks(db: aiosqlite.Connection, *, dry_run: bool) -> int:
|
||||
"""For each document with chunks_json, insert chunk rows if absent."""
|
||||
chunks_repo = SqliteChunkRepository()
|
||||
cursor = await db.execute("SELECT id FROM documents")
|
||||
doc_rows = await cursor.fetchall()
|
||||
written = 0
|
||||
read = _chunk_dicts_for_doc(db)
|
||||
for doc_row in doc_rows:
|
||||
document_id = doc_row["id"]
|
||||
existing = await chunks_repo.find_for_document(document_id, include_deleted=True)
|
||||
if existing:
|
||||
continue # already materialized
|
||||
chunk_dicts = await read(document_id)
|
||||
if not chunk_dicts:
|
||||
continue
|
||||
chunks = [
|
||||
_chunk_dict_to_chunk(d, document_id=document_id, sequence=i)
|
||||
for i, d in enumerate(chunk_dicts)
|
||||
]
|
||||
logger.info("chunks: doc=%s materializing %d", document_id, len(chunks))
|
||||
if not dry_run:
|
||||
await chunks_repo.insert_many(chunks)
|
||||
written += len(chunks)
|
||||
return written
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Step: backfill_default_store_links
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _chunks_to_results(chunks: Iterable[Chunk]) -> list[ChunkResult]:
|
||||
return [
|
||||
ChunkResult(
|
||||
text=c.text,
|
||||
headings=list(c.headings),
|
||||
source_page=c.source_page,
|
||||
bboxes=[ChunkBbox(page=b.page, bbox=list(b.bbox)) for b in c.bboxes],
|
||||
doc_items=[ChunkDocItem(self_ref=d.self_ref, label=d.label) for d in c.doc_items],
|
||||
)
|
||||
for c in chunks
|
||||
]
|
||||
|
||||
|
||||
async def _step_backfill_links(db: aiosqlite.Connection, *, dry_run: bool) -> int:
|
||||
"""For each doc that already has chunks materialized, create a
|
||||
`default`-store link in `Ingested` state with a freshly-computed
|
||||
chunkset hash. This treats "chunks exist" as a proxy for "the
|
||||
document has been ingested into the default store" — which is true
|
||||
for tenants with the legacy single-index deployment that 0.6.0
|
||||
formalises into the `default` store.
|
||||
|
||||
This step does NOT call OpenSearch. Operators with non-trivial
|
||||
OpenSearch state should run a separate reindex / verification
|
||||
after this script (see runbook).
|
||||
"""
|
||||
store_repo = SqliteStoreRepository()
|
||||
chunks_repo = SqliteChunkRepository()
|
||||
link_repo = SqliteDocumentStoreLinkRepository()
|
||||
|
||||
default_store = await store_repo.get_default()
|
||||
if default_store is None:
|
||||
logger.warning(
|
||||
"No default store — skipping link backfill (run init_db first)",
|
||||
)
|
||||
return 0
|
||||
|
||||
cursor = await db.execute("SELECT id FROM documents")
|
||||
doc_rows = await cursor.fetchall()
|
||||
written = 0
|
||||
for doc_row in doc_rows:
|
||||
document_id = doc_row["id"]
|
||||
existing = await link_repo.find_one(document_id, default_store.id)
|
||||
if existing is not None:
|
||||
continue
|
||||
chunks = await chunks_repo.find_for_document(document_id)
|
||||
if not chunks:
|
||||
continue
|
||||
h = chunkset_hash(_chunks_to_results(chunks))
|
||||
link = DocumentStoreLink(
|
||||
document_id=document_id,
|
||||
store_id=default_store.id,
|
||||
state=DocumentStoreLinkState.INGESTED,
|
||||
chunkset_hash=h,
|
||||
)
|
||||
logger.info(
|
||||
"links: doc=%s -> store=%s ingested (hash=%s…)",
|
||||
document_id,
|
||||
default_store.id,
|
||||
h[:8],
|
||||
)
|
||||
if not dry_run:
|
||||
await link_repo.upsert(link)
|
||||
written += 1
|
||||
return written
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Step: reaggregate_document_lifecycle
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _step_reaggregate(db: aiosqlite.Connection, *, dry_run: bool) -> int:
|
||||
"""Recompute the doc lifecycle state from per-store links + the
|
||||
fallback inferred in step 1."""
|
||||
link_repo = SqliteDocumentStoreLinkRepository()
|
||||
cursor = await db.execute("SELECT id, lifecycle_state FROM documents")
|
||||
rows = await cursor.fetchall()
|
||||
written = 0
|
||||
for row in rows:
|
||||
doc_id = row["id"]
|
||||
current = DocumentLifecycleState(row["lifecycle_state"])
|
||||
links = await link_repo.find_for_document(doc_id)
|
||||
target = aggregate_lifecycle(links, fallback=current)
|
||||
if target == current:
|
||||
continue
|
||||
logger.info(
|
||||
"reaggregate: doc=%s %s -> %s",
|
||||
doc_id,
|
||||
current.value,
|
||||
target.value,
|
||||
)
|
||||
if not dry_run:
|
||||
await db.execute(
|
||||
"UPDATE documents SET lifecycle_state = ?, "
|
||||
"lifecycle_state_at = datetime('now') WHERE id = ?",
|
||||
(target.value, doc_id),
|
||||
)
|
||||
written += 1
|
||||
if not dry_run:
|
||||
await db.commit()
|
||||
return written
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Orchestration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
_STEPS = [
|
||||
("backfill_lifecycle", _step_backfill_lifecycle),
|
||||
("materialize_chunks", _step_materialize_chunks),
|
||||
("backfill_links", _step_backfill_links),
|
||||
("reaggregate", _step_reaggregate),
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class StepResult:
|
||||
name: str
|
||||
skipped: bool
|
||||
written: int
|
||||
|
||||
|
||||
async def run(
|
||||
*,
|
||||
dry_run: bool = False,
|
||||
only_step: str | None = None,
|
||||
) -> list[StepResult]:
|
||||
"""Run the migration. Returns one StepResult per step (skipped or not)."""
|
||||
await init_db() # ensures the schema and migration_progress table exist
|
||||
results: list[StepResult] = []
|
||||
async with get_connection() as db:
|
||||
for name, step in _STEPS:
|
||||
if only_step and only_step != name:
|
||||
continue
|
||||
if await _is_done(db, name):
|
||||
results.append(StepResult(name=name, skipped=True, written=0))
|
||||
logger.info("step %s: already done — skipping", name)
|
||||
continue
|
||||
logger.info("step %s: running (dry_run=%s)", name, dry_run)
|
||||
written = await step(db, dry_run=dry_run)
|
||||
if not dry_run:
|
||||
await _mark_done(db, name)
|
||||
results.append(StepResult(name=name, skipped=False, written=written))
|
||||
return results
|
||||
|
||||
|
||||
def _print_summary(results: list[StepResult], *, dry_run: bool) -> None:
|
||||
print()
|
||||
print(f"{'step':35s} {'wrote':>10s} {'skipped':>10s}")
|
||||
print("-" * 60)
|
||||
total_written = 0
|
||||
for r in results:
|
||||
print(f"{r.name:35s} {r.written:>10d} {'yes' if r.skipped else 'no':>10s}")
|
||||
total_written += r.written
|
||||
print("-" * 60)
|
||||
print(f"{'total':35s} {total_written:>10d}")
|
||||
print()
|
||||
print(f"dry_run={dry_run}")
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
|
||||
)
|
||||
parser = argparse.ArgumentParser(prog="migrate_06")
|
||||
parser.add_argument(
|
||||
"--dry-run",
|
||||
action="store_true",
|
||||
help="Print the plan without writing.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-step",
|
||||
choices=[name for name, _ in _STEPS],
|
||||
help="Run a single named step (useful when re-doing one phase).",
|
||||
)
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
results = asyncio.run(run(dry_run=args.dry_run, only_step=args.only_step))
|
||||
_print_summary(results, dry_run=args.dry_run)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Loading…
Reference in a new issue