more_dots/scripts/update_sql_prompts.py

#!/usr/bin/env python3
"""
SQL 提示词 Redis 热更新脚本

用法:
    python scripts/update_sql_prompts.py                     # 同步所有本地文件到 Redis
    python scripts/update_sql_prompts.py --tables apbo_eta_ful apbo_eta_milestone  # 同步指定表
    python scripts/update_sql_prompts.py --list              # 列出 Redis 中的所有表
    python scripts/update_sql_prompts.py --delete apbo_eta_ful  # 删除指定表
    python scripts/update_sql_prompts.py --from-file path/to/file.json --table apbo_eta_ful  # 从指定文件更新
"""

import argparse
import json
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from config import Config
from services.storage.cache import RedisCache


def _first_line(value) -> str:
    if value is None:
        return ""
    text = str(value)
    return text.splitlines()[0].strip() if text else ""


def _to_bool(value, default: bool = False) -> bool:
    text = _first_line(value).lower()
    if not text:
        return default
    return text in ("1", "true", "yes", "y", "on")


def _to_int(value, default: int = 0) -> int:
    text = _first_line(value)
    if not text:
        return default
    try:
        return int(text)
    except Exception:
        return default


def get_redis_cache() -> RedisCache:
    redis_cfg = Config.get_section("redis")
    enabled = _to_bool(redis_cfg.get("enabled", "false"))
    if not enabled:
        raise RuntimeError("Redis 未启用，请检查配置 redis.enabled")

    url = _first_line(redis_cfg.get("url"))
    db = _to_int(redis_cfg.get("db", redis_cfg.get("database", 0)), default=0)
    if not url:
        host = _first_line(redis_cfg.get("host"))
        port = _first_line(redis_cfg.get("port", "6379")) or "6379"
        password = _first_line(redis_cfg.get("password", ""))
        username = _first_line(redis_cfg.get("username", ""))
        database = _first_line(redis_cfg.get("database", str(db))) or str(db)
        if host:
            from urllib.parse import quote_plus
            if username and password:
                auth = f"{quote_plus(username)}:{quote_plus(password)}@"
            elif password:
                auth = f":{quote_plus(password)}@"
            else:
                auth = ""
            url = f"redis://{auth}{host}:{port}/{database}"

    if not url:
        raise RuntimeError("Redis 配置不完整，请检查 redis.url 或 redis.host")

    return RedisCache(url=url, db=db)


def get_ttl() -> int:
    redis_cfg = Config.get_section("redis")
    return _to_int(redis_cfg.get("sql_prompt_ttl", 0), default=0)


def sync_from_local_files(cache: RedisCache, tables: list = None, ttl: int = None):
    prompts_dir = PROJECT_ROOT / "config" / "sql_gen_prompts"

    if tables:
        files = [prompts_dir / f"{t}.json" for t in tables]
    else:
        files = list(prompts_dir.glob("*.json"))

    results = {}
    for file_path in files:
        if not file_path.exists():
            print(f"[跳过] 文件不存在: {file_path}")
            continue

        table_name = file_path.stem
        try:
            with open(file_path, "r", encoding="utf-8") as f:
                prompt = json.load(f)

            key = f"sql_prompt:{table_name}"
            cache.set(key, json.dumps(prompt, ensure_ascii=False), ttl)
            results[table_name] = "success"
            print(f"[成功] {table_name}")
        except Exception as e:
            results[table_name] = f"failed: {e}"
            print(f"[失败] {table_name}: {e}")

    return results


def update_from_file(cache: RedisCache, file_path: str, table_name: str, ttl: int = None):
    path = Path(file_path)
    if not path.exists():
        print(f"[错误] 文件不存在: {file_path}")
        return False

    try:
        with open(path, "r", encoding="utf-8") as f:
            prompt = json.load(f)

        key = f"sql_prompt:{table_name}"
        cache.set(key, json.dumps(prompt, ensure_ascii=False), ttl)
        print(f"[成功] 已更新 {table_name}")
        return True
    except Exception as e:
        print(f"[失败] {table_name}: {e}")
        return False


def list_tables(cache: RedisCache):
    keys = cache.keys("sql_prompt:*")
    tables = []
    for key in keys:
        parts = key.split(":", 1)
        if len(parts) == 2 and parts[1] != "table_list":
            tables.append(parts[1])

    if tables:
        print("Redis 中的 SQL 提示词表:")
        for t in sorted(tables):
            print(f"  - {t}")
    else:
        print("Redis 中没有 SQL 提示词")
    return tables


def delete_table(cache: RedisCache, table_name: str):
    key = f"sql_prompt:{table_name}"
    cache.delete(key)
    print(f"[成功] 已删除 {table_name}")


def main():
    parser = argparse.ArgumentParser(description="SQL 提示词 Redis 热更新工具")
    parser.add_argument("--tables", nargs="*", help="指定要同步的表名列表")
    parser.add_argument("--list", action="store_true", help="列出 Redis 中的所有表")
    parser.add_argument("--delete", type=str, help="删除指定表")
    parser.add_argument("--from-file", type=str, help="从指定文件更新")
    parser.add_argument("--table", type=str, help="目标表名（与 --from-file 配合使用）")

    args = parser.parse_args()

    try:
        cache = get_redis_cache()
        ttl = get_ttl()
    except Exception as e:
        print(f"[错误] {e}")
        sys.exit(1)

    if args.list:
        list_tables(cache)
    elif args.delete:
        delete_table(cache, args.delete)
    elif args.from_file:
        if not args.table:
            print("[错误] 使用 --from-file 时必须指定 --table")
            sys.exit(1)
        update_from_file(cache, args.from_file, args.table, ttl)
    else:
        sync_from_local_files(cache, args.tables, ttl)


if __name__ == "__main__":
    main()