monitor-shared-utils/scripts/bootstrap_db.py

"""Bootstrap the shared Mongo database: collections, indexes, and config seeds.

Rebuilds what the monitors and the discord consumer expect to find: the
configs/control/alerts collections in both modes, the alerts indexes the
consumer relies on for claim ordering and dedupe, one empty product
collection per monitor and mode, the per-monitor config documents from
scripts/db_seed/<mode>/<monitor>.json, and the control docs from
scripts/db_seed/control.json (the monitors registry; log_channel and
bot_owner hold Discord ids only the bot's commands know, and migrations
must stay absent so the consumer re-runs them on a fresh database).

Idempotent: existing collections and indexes are left alone and seed
documents use $setOnInsert, so re-running never clobbers operator edits.

Usage:
    python scripts/bootstrap_db.py [--dry-run] [--url URL] [--database NAME]
"""

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

from pymongo import ASCENDING, DESCENDING, MongoClient
from pymongo.collection import Collection
from pymongo.database import Database

from monitor_utils.constants import MONGO_DATABASE_NAME, MONGO_URL

LOGGER = logging.getLogger("bootstrap_db")

SEED_ROOT = Path(__file__).resolve().parent / "db_seed"
# collection suffix -> seed directory
MODES = {"": "production", "_testing": "testing"}
# monitors that call init_monitor and so own a product collection per mode
MONITORS = [
    "apple",
    "bestbuy",
    "corsair",
    "hp",
    "lg",
    "northerntool",
    "samsung",
    "samsung_shop",
    "woot",
]
SEED_KEYS = {"categories", "channels"}
# categories shape each monitor's task function accepts: run_monitor_cycle
# passes query_category_id= only for a mapping, and only bestbuy (required
# query) and apple (node -> route) take that argument; the rest take a list
CATEGORY_SHAPES = {
    "apple": dict,
    "bestbuy": dict,
    "corsair": list,
    "hp": list,
    "lg": list,
    "northerntool": list,
    "samsung": list,
    "samsung_shop": list,
    "woot": list,
}
CONTROL_SEED = SEED_ROOT / "control.json"

# shape setup_config (runner.py) reads: _id is the monitor name, categories is
# a list of names or a name -> query mapping, channels maps a bucket to ONE
# channel id as a digit string — the consumer's find_invalid_channel_ids
# rejects lists (the pre-07-16 spreading shape). Extra fields stay allowed.
CHANNEL_ID = {"bsonType": "string", "pattern": "^[0-9]+$"}
CONFIG_VALIDATOR = {
    "$jsonSchema": {
        "bsonType": "object",
        "required": ["_id", "categories", "channels"],
        "properties": {
            "_id": {"bsonType": "string"},
            "categories": {
                "anyOf": [
                    {"bsonType": "array", "items": {"bsonType": "string"}},
                    {
                        "bsonType": "object",
                        "additionalProperties": {"bsonType": "string"},
                    },
                ]
            },
            "channels": {
                "bsonType": "object",
                "additionalProperties": CHANNEL_ID,
            },
        },
    }
}
# identical to the consumer's cog_load create_index calls; a differing spec
# under the same key pattern would make its startup fail
ALERT_INDEXES = [
    ([("claimed", ASCENDING), ("priority", DESCENDING), ("createdAt", ASCENDING)], {}),
    (
        [("product_id", ASCENDING)],
        {"unique": True, "partialFilterExpression": {"claimed": False}},
    ),
]


def ensure_collection(
    database: Database, name: str, dry_run: bool = False, **options
) -> bool:
    """Create a collection when absent; return True when it was created."""

    if name in database.list_collection_names():
        LOGGER.info("collection exists", extra={"collection": name})
        return False

    LOGGER.info("creating collection", extra={"collection": name, "dry_run": dry_run})
    if not dry_run:
        database.create_collection(name, **options)

    return True


def ensure_config_validator(
    database: Database, name: str, dry_run: bool = False
) -> bool:
    """Apply CONFIG_VALIDATOR to an existing collection; True when it changed."""

    info = next(iter(database.list_collections(filter={"name": name})), None)
    if info is None or info["options"].get("validator") == CONFIG_VALIDATOR:
        return False

    LOGGER.info("updating validator", extra={"collection": name, "dry_run": dry_run})
    if not dry_run:
        database.command("collMod", name, validator=CONFIG_VALIDATOR)

    return True


def ensure_alert_indexes(collection: Collection, dry_run: bool = False) -> list[str]:
    """Create the consumer's alert indexes; return the index names."""

    names = []
    for keys, options in ALERT_INDEXES:
        if dry_run:
            names.append("+".join(f"{field}_{direction}" for field, direction in keys))
            continue

        names.append(collection.create_index(keys, **options))

    LOGGER.info(
        "alert indexes ensured",
        extra={"collection": collection.name, "indexes": names, "dry_run": dry_run},
    )
    return names


def load_seed_configs(seed_dir: Path) -> dict[str, dict]:
    """Load <monitor>.json seed docs keyed by monitor name.

    Raises:
        ValueError: When a seed file lacks categories or channels.
    """

    docs = {}
    for path in sorted(seed_dir.glob("*.json")):
        with path.open() as seed_file:
            doc = json.load(seed_file)

        missing = SEED_KEYS - set(doc)
        if missing:
            raise ValueError(f"{path}: missing {sorted(missing)}")

        bad_channels = {
            bucket: value
            for bucket, value in doc["channels"].items()
            if not (isinstance(value, str) and value.isdigit())
        }
        if bad_channels:
            raise ValueError(
                f"{path}: channels must be single id strings {bad_channels}"
            )

        expected = CATEGORY_SHAPES.get(path.stem)
        if expected is not None and not isinstance(doc["categories"], expected):
            raise ValueError(
                f"{path}: categories must be {expected.__name__} for {path.stem}"
            )

        docs[path.stem] = {key: doc[key] for key in SEED_KEYS}

    return docs


def seed_configs(
    collection: Collection, docs: dict[str, dict], dry_run: bool = False
) -> list[str]:
    """Insert absent config docs via $setOnInsert; return the ids inserted."""

    inserted = []
    for monitor_name, doc in docs.items():
        if dry_run:
            if collection.find_one({"_id": monitor_name}, {"_id": 1}) is None:
                inserted.append(monitor_name)
            continue

        result = collection.update_one(
            {"_id": monitor_name}, {"$setOnInsert": doc}, upsert=True
        )
        if result.upserted_id is not None:
            inserted.append(monitor_name)

    LOGGER.info(
        "configs seeded",
        extra={
            "collection": collection.name,
            "inserted": inserted,
            "skipped": sorted(set(docs) - set(inserted)),
            "dry_run": dry_run,
        },
    )
    return inserted


def load_control_seed(path: Path) -> dict[str, dict]:
    """Load control docs keyed by _id from one JSON file."""

    with path.open() as seed_file:
        docs = json.load(seed_file)

    for doc_id, doc in docs.items():
        if not isinstance(doc, dict):
            raise ValueError(f"{path}: {doc_id} must be an object")

    return docs


def seed_control(
    collection: Collection, docs: dict[str, dict], dry_run: bool = False
) -> list[str]:
    """Insert absent control docs via $setOnInsert; return the ids inserted."""

    inserted = []
    for doc_id, doc in docs.items():
        if dry_run:
            if collection.find_one({"_id": doc_id}, {"_id": 1}) is None:
                inserted.append(doc_id)
            continue

        result = collection.update_one(
            {"_id": doc_id}, {"$setOnInsert": doc}, upsert=True
        )
        if result.upserted_id is not None:
            inserted.append(doc_id)

    LOGGER.info(
        "control seeded",
        extra={
            "collection": collection.name,
            "inserted": inserted,
            "skipped": sorted(set(docs) - set(inserted)),
            "dry_run": dry_run,
        },
    )
    return inserted


def bootstrap(
    database: Database, seed_root: Path = SEED_ROOT, dry_run: bool = False
) -> dict:
    """Run every step for both modes and return a summary of what changed."""

    summary = {"created": [], "validators": [], "indexes": {}, "seeded": {}}
    for suffix, seed_mode in MODES.items():
        configs_name = f"configs{suffix}"
        if ensure_collection(
            database, configs_name, dry_run, validator=CONFIG_VALIDATOR
        ):
            summary["created"].append(configs_name)
        elif ensure_config_validator(database, configs_name, dry_run):
            summary["validators"].append(configs_name)

        for base in ("control", "alerts", *MONITORS):
            name = f"{base}{suffix}"
            if ensure_collection(database, name, dry_run):
                summary["created"].append(name)

        alerts_name = f"alerts{suffix}"
        summary["indexes"][alerts_name] = ensure_alert_indexes(
            database[alerts_name], dry_run
        )
        docs = load_seed_configs(seed_root / seed_mode)
        summary["seeded"][configs_name] = seed_configs(
            database[configs_name], docs, dry_run
        )
        control_name = f"control{suffix}"
        summary["seeded"][control_name] = seed_control(
            database[control_name],
            load_control_seed(seed_root / "control.json"),
            dry_run,
        )

    return summary


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    parser.add_argument("--url", default=MONGO_URL, help="Mongo connection url")
    parser.add_argument("--database", default=MONGO_DATABASE_NAME)
    parser.add_argument(
        "--dry-run", action="store_true", help="report actions without writing"
    )
    args = parser.parse_args(argv)
    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")

    client = MongoClient(args.url, serverSelectionTimeoutMS=30_000)
    LOGGER.info("connected", extra={"server": client.server_info()["version"]})
    summary = bootstrap(client[args.database], dry_run=args.dry_run)
    print(json.dumps(summary, indent=2))
    return 0


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