91 lines
3.1 KiB
Python
91 lines
3.1 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""工作手机总控平台改版的 Mongo 集合与索引初始化。
|
||
|
||
示例：
|
||
python3 sdk/app/scripts/init_console_redesign_mongo.py --dry-run
|
||
python3 sdk/app/scripts/init_console_redesign_mongo.py
|
||
python3 sdk/app/scripts/init_console_redesign_mongo.py --rollback
|
||
|
||
默认只使用 sdk/app/config.py 中的 Mongo 配置。脚本不会删除集合或业务数据；
|
||
``--rollback`` 也只删除本脚本创建的 ``console11_`` 索引。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import sys
|
||
from typing import Any, Callable
|
||
|
||
APP_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
if APP_DIR not in sys.path:
|
||
sys.path.insert(0, APP_DIR)
|
||
|
||
from pymongo import MongoClient
|
||
|
||
from config import settings
|
||
from services.control_plane_indexes import (
|
||
CONTROL_PLANE_INDEXES,
|
||
CONTROL_PLANE_COLLECTIONS,
|
||
apply_control_plane_indexes,
|
||
index_plan,
|
||
rollback_control_plane_indexes,
|
||
)
|
||
|
||
|
||
def build_parser() -> argparse.ArgumentParser:
|
||
parser = argparse.ArgumentParser(description="初始化总控平台改版 Mongo 集合和索引")
|
||
parser.add_argument("--uri", default=settings.MONGO_URI, help="MongoDB URI，默认读取配置")
|
||
parser.add_argument("--database", default=settings.MONGO_DB, help="数据库名，默认读取配置")
|
||
parser.add_argument("--dry-run", action="store_true", help="只列出集合和索引，不连接数据库、不写入")
|
||
parser.add_argument("--rollback", action="store_true", help="只移除本任务 console11_ 索引，不删除数据")
|
||
parser.add_argument("--server-selection-timeout-ms", type=int, default=3000, help="Mongo 连接探测超时")
|
||
return parser
|
||
|
||
|
||
def _print_plan(*, rollback: bool = False) -> None:
|
||
plan = index_plan()
|
||
print(json.dumps({"collections": len(plan), "indexes": len(CONTROL_PLANE_INDEXES)}, ensure_ascii=False))
|
||
for collection, specs in plan.items():
|
||
print(f"集合 {collection}: {len(specs)} 个索引")
|
||
for spec in specs:
|
||
operation = "将移除" if rollback else "将创建"
|
||
print(f" - {operation} {spec['name']} {spec['keys']}")
|
||
|
||
|
||
def run(
|
||
argv: list[str] | None = None,
|
||
*,
|
||
client_factory: Callable[..., Any] = MongoClient,
|
||
) -> dict[str, Any]:
|
||
args = build_parser().parse_args(argv)
|
||
if args.dry_run:
|
||
_print_plan(rollback=args.rollback)
|
||
return {
|
||
"dry_run": True,
|
||
"rollback": args.rollback,
|
||
"collection_count": len(CONTROL_PLANE_COLLECTIONS),
|
||
"index_count": len(CONTROL_PLANE_INDEXES),
|
||
}
|
||
|
||
client = client_factory(args.uri, serverSelectionTimeoutMS=args.server_selection_timeout_ms)
|
||
try:
|
||
db = client[args.database]
|
||
if args.rollback:
|
||
summary = rollback_control_plane_indexes(db)
|
||
else:
|
||
summary = apply_control_plane_indexes(db)
|
||
print(json.dumps(summary, ensure_ascii=False, sort_keys=True))
|
||
return summary
|
||
finally:
|
||
close = getattr(client, "close", None)
|
||
if close:
|
||
close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
run()
|
||
|