#!/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()