autool-dispatcher/migrate_to_new_db.py
2026-06-17 19:50:39 +08:00

135 lines
4.3 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
新数据库初始化/校验脚本。
运行时代码不再兼容旧库,也不创建 app_collect_summary / v_collection_latest
这类兼容视图。本脚本只负责用 schema_new.sql 初始化新库,并校验核心新表。
"""
import argparse
import os
import sqlite3
import sys
import time
from analytics import AnalyticsRepository
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
RUNTIME_MAIN = os.path.join(BASE_DIR, "runtime", "main")
DEFAULT_DB = os.path.join(RUNTIME_MAIN, "monitoring_new.sqlite3")
SCHEMA_SQL = os.path.join(BASE_DIR, "schema_new.sql")
REQUIRED_TABLES = {
"collection_task",
"app_catalog",
"apk_registry",
"worker_activity_summary",
"analytics_job",
"analytics_source_file",
"app_domain_traffic",
"app_traffic_component",
}
FORBIDDEN_OBJECTS = {
"app_collect_summary",
"v_collection_latest",
"v_collection_stats_by_tag",
}
def log(message: str) -> None:
print(f"[new-db] {message}", flush=True)
def _remove_sqlite_files(db_path: str) -> None:
for suffix in ("", "-wal", "-shm"):
path = f"{db_path}{suffix}"
if os.path.exists(path):
os.remove(path)
def initialize_new_db(db_path: str, *, replace: bool = False) -> None:
if not os.path.exists(SCHEMA_SQL):
raise FileNotFoundError(f"schema file not found: {SCHEMA_SQL}")
if os.path.exists(db_path):
if not replace:
log(f"数据库已存在,跳过初始化: {db_path}")
return
_remove_sqlite_files(db_path)
os.makedirs(os.path.dirname(db_path) or ".", exist_ok=True)
with open(SCHEMA_SQL, "r", encoding="utf-8") as handle:
schema_sql = handle.read()
with sqlite3.connect(db_path) as connection:
connection.executescript(schema_sql)
AnalyticsRepository(db_path=db_path)
log(f"✅ 新库 schema 初始化完成: {db_path}")
def validate_new_db(db_path: str) -> bool:
if not os.path.exists(db_path):
log(f"❌ 数据库不存在: {db_path}")
return False
ok = True
with sqlite3.connect(db_path) as connection:
connection.row_factory = sqlite3.Row
objects = {
str(row["name"])
for row in connection.execute(
"SELECT name FROM sqlite_master WHERE type IN ('table', 'view')"
).fetchall()
}
missing = sorted(REQUIRED_TABLES - objects)
forbidden = sorted(FORBIDDEN_OBJECTS & objects)
if missing:
ok = False
log("❌ 缺少新 schema 表: " + ", ".join(missing))
if forbidden:
ok = False
log("❌ 存在旧兼容对象: " + ", ".join(forbidden))
if ok:
log("✅ schema 对象校验通过")
for table_name in sorted(REQUIRED_TABLES & objects):
row = connection.execute(f"SELECT COUNT(*) AS total FROM {table_name}").fetchone()
log(f"{table_name}: {int(row['total'] or 0)}")
task_state = connection.execute(
"""
SELECT task_status, COUNT(*) AS total
FROM collection_task
GROUP BY task_status
ORDER BY total DESC
"""
).fetchall() if "collection_task" in objects else []
if task_state:
log(
"collection_task.task_status 分布: "
+ str({str(row["task_status"] or "NULL"): int(row["total"] or 0) for row in task_state})
)
return ok
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Initialize and validate the new analytics database schema.")
parser.add_argument("--db-path", default=DEFAULT_DB, help="Target sqlite path.")
parser.add_argument("--replace", action="store_true", help="Delete the target sqlite files before initializing.")
parser.add_argument("--validate-only", action="store_true", help="Only validate an existing database.")
return parser
def main() -> int:
args = build_parser().parse_args()
db_path = str(args.db_path or "").strip() or DEFAULT_DB
started_at = time.time()
if not args.validate_only:
initialize_new_db(db_path, replace=bool(args.replace))
ok = validate_new_db(db_path)
log(f"完成,用时 {time.time() - started_at:.1f}s")
return 0 if ok else 1
if __name__ == "__main__":
sys.exit(main())