balanceboard/migrate_content_to_db.py

#!/usr/bin/env python
"""Backfill on-disk content JSON into the Postgres ``posts``/``comments`` tables.

Phase 3 prep (Agent D): reads ``data/posts/*.json`` and ``data/comments/*.json``
and upserts them into the ``Post`` / ``Comment`` models. Live reads/writes still
go through ``PostService`` (disk JSON) — this script only populates the DB so a
later cutover has data to read. It is idempotent: existing rows are skipped by
``uuid`` (run it again after collecting new content to backfill only the new
files).

Usage::

    python migrate_content_to_db.py [--data-dir data] [--batch-size 500] [--dry-run]

Requires the full Flask/Postgres stack (``DATABASE_URL`` or the POSTGRES_*
env vars) — run inside the docker compose environment, not locally.
"""

import argparse
import json
import logging
import sys
from pathlib import Path

from app import create_app
from database import db
from models import Comment, Post

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("migrate_content_to_db")


def _load_json(path: Path):
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except (json.JSONDecodeError, OSError) as e:
        logger.warning("Skipping unreadable file %s: %s", path, e)
        return None


def _existing_uuids(model, uuids):
    """Return the subset of ``uuids`` already present in the table."""
    if not uuids:
        return set()
    found = set()
    # Chunk to avoid huge IN clauses.
    for i in range(0, len(uuids), 500):
        chunk = uuids[i:i + 500]
        rows = db.session.query(model.uuid).filter(model.uuid.in_(chunk)).all()
        found.update(r[0] for r in rows)
    return found


def backfill_posts(posts_dir: Path, batch_size: int, dry_run: bool) -> int:
    files = sorted(posts_dir.glob("*.json")) if posts_dir.exists() else []
    if not files:
        logger.info("No post files found in %s", posts_dir)
        return 0

    records = []
    for pf in files:
        data = _load_json(pf)
        if not data or not data.get("uuid"):
            continue
        records.append(data)

    seen = _existing_uuids(Post, [r["uuid"] for r in records])
    inserted = 0
    batch = []
    for r in records:
        if r["uuid"] in seen:
            continue
        batch.append(Post(
            uuid=r["uuid"],
            external_id=r.get("id"),
            platform=r.get("platform", "") or "",
            source=r.get("source", "") or "",
            title=(r.get("title") or "")[:500],
            author=r.get("author"),
            url=r.get("url"),
            content=r.get("content"),
            score=int(r.get("score", 0) or 0),
            timestamp=int(r.get("timestamp", 0) or 0),
            tags=r.get("tags"),
            moderation_uuid=r.get("moderation_uuid"),
        ))
        inserted += 1
        if len(batch) >= batch_size:
            _flush(batch, dry_run)
            batch = []
    _flush(batch, dry_run)
    logger.info("Posts: backfilled %d new (%d already present)", inserted, len(seen))
    return inserted


def backfill_comments(comments_dir: Path, batch_size: int, dry_run: bool) -> int:
    files = sorted(comments_dir.glob("*.json")) if comments_dir.exists() else []
    if not files:
        logger.info("No comment files found in %s", comments_dir)
        return 0

    records = []
    for cf in files:
        data = _load_json(cf)
        if not data or not data.get("uuid"):
            continue
        records.append(data)

    seen = _existing_uuids(Comment, [r["uuid"] for r in records])
    inserted = 0
    batch = []
    for r in records:
        if r["uuid"] in seen:
            continue
        batch.append(Comment(
            uuid=r["uuid"],
            post_uuid=r.get("post_uuid") or "",
            platform=r.get("platform"),
            parent_comment_uuid=r.get("parent_comment_uuid"),
            comment_id=r.get("comment_id"),
            author=r.get("author"),
            content=r.get("content"),
            score=int(r.get("score", 0) or 0),
            timestamp=int(r.get("timestamp", 0) or 0),
            depth=int(r.get("depth", 0) or 0),
            moderation_uuid=r.get("moderation_uuid"),
        ))
        inserted += 1
        if len(batch) >= batch_size:
            _flush(batch, dry_run)
            batch = []
    _flush(batch, dry_run)
    logger.info("Comments: backfilled %d new (%d already present)", inserted, len(seen))
    return inserted


def _flush(batch, dry_run):
    if not batch:
        return
    if dry_run:
        logger.info("[dry-run] would insert %d rows", len(batch))
        return
    db.session.bulk_save_objects(batch)
    db.session.commit()


def main(argv=None):
    parser = argparse.ArgumentParser(description="Backfill content JSON into Postgres.")
    parser.add_argument("--data-dir", default="data", help="Root data directory.")
    parser.add_argument("--batch-size", type=int, default=500, help="Insert batch size.")
    parser.add_argument("--dry-run", action="store_true", help="Log counts without writing.")
    args = parser.parse_args(argv)

    app = create_app()
    with app.app_context():
        posts_dir = Path(args.data_dir) / "posts"
        comments_dir = Path(args.data_dir) / "comments"
        p = backfill_posts(posts_dir, args.batch_size, args.dry_run)
        c = backfill_comments(comments_dir, args.batch_size, args.dry_run)
        logger.info("Done. posts=%d comments=%d", p, c)
    return 0


if __name__ == "__main__":
    sys.exit(main())