chore/backend/scripts/migrate_to_mongodb.py

# python
"""
Migrate existing TinyDB JSON files into MongoDB.

Usage:
    cd backend
    python -m scripts.migrate_to_mongodb [--dry-run]

The script reads files from ``data/db/`` (or ``test_data/db/`` when
``DB_ENV=test``), maps each record's ``id`` field to MongoDB's ``_id`` field,
and inserts the records idempotently.  TinyDB files are backed up to
``<db_dir>/backups/<timestamp>/`` before the first migration run.
"""
import argparse
import json
import os
import shutil
import sys
from datetime import datetime
from pathlib import Path

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from config.paths import get_database_dir
from db.db import COLLECTION_INDEXES, ensure_mongodb_indexes
from db.mongo_client import get_mongo_client, get_mongo_db_name


# Map TinyDB JSON filenames to MongoDB collection names.
COLLECTION_FILE_MAP = {
    'children.json': 'children',
    'tasks.json': 'tasks',
    'routines.json': 'routines',
    'routine_items.json': 'routine_items',
    'routine_schedules.json': 'routine_schedules',
    'routine_extensions.json': 'routine_extensions',
    'rewards.json': 'rewards',
    'images.json': 'images',
    'pending_rewards.json': 'pending_rewards',
    'pending_confirmations.json': 'pending_confirmations',
    'users.json': 'users',
    'tracking_events.json': 'tracking_events',
    'child_overrides.json': 'child_overrides',
    'chore_schedules.json': 'chore_schedules',
    'task_extensions.json': 'task_extensions',
    'refresh_tokens.json': 'refresh_tokens',
    'push_subscriptions.json': 'push_subscriptions',
    'digest_action_tokens.json': 'digest_action_tokens',
}


def _load_tinydb_records(path: str) -> list[dict]:
    """Load all records from a TinyDB JSON file."""
    with open(path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    default_table = data.get('_default', {})
    return list(default_table.values())


def _doc_to_mongo(doc: dict) -> dict:
    """Map the model ``id`` field to MongoDB's ``_id`` field.

    The original ``id`` field is removed so documents do not store both
    ``_id`` and ``id`` with identical values.
    """
    mongo_doc = dict(doc)
    if 'id' in mongo_doc:
        mongo_doc['_id'] = mongo_doc.pop('id')
    return mongo_doc


def migrate(dry_run: bool = False) -> dict:
    """Migrate TinyDB files to MongoDB and return a per-collection summary."""
    db_dir = get_database_dir()
    if not os.path.isdir(db_dir):
        raise FileNotFoundError(f'Database directory does not exist: {db_dir}')

    client = get_mongo_client()
    db_name = get_mongo_db_name()
    db = client[db_name]

    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    backup_dir = os.path.join(db_dir, 'backups', timestamp)

    if not dry_run:
        os.makedirs(backup_dir, exist_ok=True)
        ensure_mongodb_indexes(client=client, db_name=db_name)

    summary: dict[str, dict] = {}

    for filename, collection_name in COLLECTION_FILE_MAP.items():
        path = os.path.join(db_dir, filename)
        if not os.path.exists(path):
            summary[collection_name] = {
                'source_file': filename,
                'total': 0,
                'migrated': 0,
                'skipped': 0,
                'status': 'missing',
            }
            continue

        records = _load_tinydb_records(path)

        if not dry_run:
            shutil.copy2(path, backup_dir)

        collection = db[collection_name]
        to_insert: list[dict] = []
        skipped = 0

        for record in records:
            doc_id = record.get('id')
            if not doc_id:
                skipped += 1
                continue

            if not dry_run:
                existing = collection.find_one({'_id': doc_id})
                if existing:
                    skipped += 1
                    continue

            to_insert.append(_doc_to_mongo(record))

        if not dry_run and to_insert:
            try:
                collection.insert_many(to_insert, ordered=False)
            except Exception as exc:  # pragma: no cover - defensive
                print(
                    f'  Warning: error inserting into {collection_name}: {exc}',
                    file=sys.stderr,
                )
                raise

        summary[collection_name] = {
            'source_file': filename,
            'total': len(records),
            'migrated': len(to_insert),
            'skipped': skipped,
            'status': 'migrated' if not dry_run else 'dry-run',
        }

    return summary


def _print_summary(summary: dict) -> None:
    """Print a human-readable migration summary."""
    print('\nMigration Summary')
    print('-' * 70)
    print(f'{"Collection":<30}{"Total":>8}{"Migrated":>10}{"Skipped":>10}{"Status":>10}')
    print('-' * 70)
    total_records = 0
    total_migrated = 0
    total_skipped = 0
    for collection_name, info in summary.items():
        print(
            f'{collection_name:<30}'
            f'{info["total"]:>8}'
            f'{info["migrated"]:>10}'
            f'{info["skipped"]:>10}'
            f'{info["status"]:>10}'
        )
        total_records += info['total']
        total_migrated += info['migrated']
        total_skipped += info['skipped']
    print('-' * 70)
    print(
        f'{"TOTAL":<30}'
        f'{total_records:>8}'
        f'{total_migrated:>10}'
        f'{total_skipped:>10}'
    )


def main():
    parser = argparse.ArgumentParser(
        description='Migrate TinyDB JSON files to MongoDB.'
    )
    parser.add_argument(
        '--dry-run',
        action='store_true',
        help='Analyze files and report counts without writing to MongoDB.',
    )
    args = parser.parse_args()

    if os.environ.get('USE_MONGODB', 'true').lower() != 'true':
        print('Set USE_MONGODB=true to run the migration.', file=sys.stderr)
        sys.exit(1)

    if not os.environ.get('MONGO_URI'):
        print('MONGO_URI is required when USE_MONGODB=true.', file=sys.stderr)
        sys.exit(1)

    if args.dry_run:
        print('Dry run: no data will be written to MongoDB.')

    summary = migrate(dry_run=args.dry_run)
    _print_summary(summary)


if __name__ == '__main__':
    main()