"""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//.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 .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())