File size: 2,008 Bytes
f8b48da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | """Migrate keyword search from OpenSearch to Postgres tsvector.
Idempotent: safe to run on fresh DBs (creates tables) and existing DBs
(adds ``keyword_postings`` + trigger, backfills postings from ``chunks``).
Usage:
python scripts/migrate_tsvector.py
What it does:
1. ``Base.metadata.create_all`` (creates ``keyword_postings`` if missing).
2. Runs ``KEYWORD_DDL`` (tsv column, weighted trigger, GIN index).
3. Backfills postings for chunks that have none (existing corpus keeps
working without re-ingestion; embeddings/Qdrant untouched).
"""
from __future__ import annotations
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT))
from app.db.models import Base # noqa: E402
from app.db.session import SessionLocal, engine # noqa: E402
from app.retrieval.bm25 import KEYWORD_DDL # noqa: E402
from sqlalchemy import text # noqa: E402
def main() -> int:
Base.metadata.create_all(bind=engine)
print("Tables ensured.")
with SessionLocal() as session:
for stmt in KEYWORD_DDL:
session.execute(text(stmt))
session.commit()
print("tsvector column, trigger, and GIN index ensured.")
with SessionLocal() as session:
res = session.execute(
text(
"INSERT INTO keyword_postings"
" (chunk_id, document_id, text, source, page_number, section)"
" SELECT c.id, c.document_id, c.text, c.source,"
" c.page_number, c.section FROM chunks c"
" WHERE NOT EXISTS (SELECT 1 FROM keyword_postings k"
" WHERE k.chunk_id = c.id)"
)
)
session.commit()
print(f"Backfilled {res.rowcount} postings from chunks.")
total = session.execute(text("SELECT COUNT(*) FROM keyword_postings")).scalar()
print(f"Total keyword postings: {total}.")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|