databank/scripts/seed_leagues_mongo.py

"""Seed initial league data into a local MongoDB instance.

Usage (PowerShell):
  python -m pip install "pymongo>=4.7"
  $env:DATABANK_DB_URI = "mongodb://localhost:27017"
  $env:DATABANK_DB_NAME = "databank"
  python scripts/seed_leagues_mongo.py
"""

from __future__ import annotations

import os
from typing import Any, Dict, Iterable

from pymongo import MongoClient


def get_connection() -> tuple[MongoClient, str]:
    """Create a MongoDB client and return ``(client, db_name)``.

    Reads env vars with sensible defaults.
    """
    uri_default = "mongodb://localhost:27017"
    name_default = "databank"

    uri = os.getenv("DATABANK_DB_URI", uri_default)
    name = os.getenv("DATABANK_DB_NAME", name_default)

    client = MongoClient(uri)
    return client, name


def ensure_unique_index_leagues(db) -> None:
    """Ensure a unique index on ``leagues.league_id``."""
    leagues = db["leagues"]
    leagues.create_index("league_id", unique=True)


def upsert_leagues(db, docs: Iterable[Dict[str, Any]]) -> dict:
    """Upsert provided league docs by ``league_id``; return counters."""
    leagues = db["leagues"]
    inserted = 0
    updated = 0
    for doc in docs:
        key = {"league_id": doc["league_id"]}
        # Replace entire document on upsert to keep fields consistent
        result = leagues.replace_one(key, doc, upsert=True)
        if result.matched_count:
            updated += int(result.modified_count == 1)
        else:
            inserted += 1
    return {"inserted": inserted, "updated": updated}


def main() -> None:
    """Entry point to seed initial league documents into MongoDB."""
    client, db_name = get_connection()
    try:
        db = client[db_name]
        ensure_unique_index_leagues(db)
        seed = [
            {"league_id": 2079, "league_name": "英超", "max_round": 38},
            {"league_id": 2080, "league_name": "英冠", "max_round": 46},
            {"league_id": 2081, "league_name": "英甲", "max_round": 46},
            {"league_id": 2082, "league_name": "英乙", "max_round": 46},
        ]
        stats = upsert_leagues(db, seed)
        print(f"Seed done: inserted={stats['inserted']}, updated={stats['updated']}")
    finally:
        client.close()


if __name__ == "__main__":
    main()