aiData/Tools/T5_ClearHistory.py

import asyncio
import os
import sys
import logging
from typing import List, Dict

# 确保项目根目录在 sys.path 中，并且优先于外部同名模块
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if project_root not in sys.path:
    sys.path.insert(0, project_root)

from Util import Win32Patch
from DbKit.Db import Db
from Util.RedisKit import RedisKit
from sqlalchemy.sql import text

try:
    from Config.Config import DB_URL
except ModuleNotFoundError:
    import importlib.util
    _root_dir_cfg = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    _config_path_cfg = os.path.join(_root_dir_cfg, "Config", "Config.py")
    _spec_cfg = importlib.util.spec_from_file_location("project_config_fallback", _config_path_cfg)
    _cfg_mod = importlib.util.module_from_spec(_spec_cfg)
    assert _spec_cfg.loader is not None
    _spec_cfg.loader.exec_module(_cfg_mod)
    DB_URL = _cfg_mod.DB_URL

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("T5_ClearHistory")

# 供应商配置映射
VENDORS = {
    "1": {"name": "特来电", "redis_pattern": "crawled:tld:*"},
    "2": {"name": "新电途", "redis_pattern": "crawled:xdt:*"},
    "3": {"name": "艾特吉易充", "redis_pattern": "crawled:aite:*"},
    "4": {"name": "驿来特", "redis_pattern": "crawled:ylt:*"},
}

async def clear_vendor_data(db: Db, redis_kit: RedisKit, vendor_info: Dict):
    """
    清理特定供应商的数据库和 Redis 数据
    """
    operator = vendor_info["name"]
    pattern = vendor_info["redis_pattern"]
    
    logger.info(f"--- 正在清理供应商: {operator} ---")
    
    # 1. 数据库清理
    try:
        async with db.AsyncSessionLocal() as session:
            async with session.begin():
                # 先删除从表记录（通过 station_hash 关联）
                # 1. t_station_status_scd
                sql_status = """
                DELETE FROM t_station_status_scd 
                WHERE station_hash IN (
                    SELECT station_hash FROM t_station_profile_scd WHERE operator = :operator
                )
                """
                result_status = await session.execute(text(sql_status), {"operator": operator})
                logger.info(f"[{operator}] t_station_status_scd 已删除 {result_status.rowcount} 行。")

                # 2. t_station_price_schedule_scd
                sql_price = """
                DELETE FROM t_station_price_schedule_scd 
                WHERE station_hash IN (
                    SELECT station_hash FROM t_station_profile_scd WHERE operator = :operator
                )
                """
                result_price = await session.execute(text(sql_price), {"operator": operator})
                logger.info(f"[{operator}] t_station_price_schedule_scd 已删除 {result_price.rowcount} 行。")

                # 3. 最后删除主表 t_station_profile_scd
                sql_profile = "DELETE FROM t_station_profile_scd WHERE operator = :operator"
                result_profile = await session.execute(text(sql_profile), {"operator": operator})
                logger.info(f"[{operator}] t_station_profile_scd 已删除 {result_profile.rowcount} 行。")
                
            logger.info(f"[{operator}] 数据库记录清理完成。")
    except Exception as e:
        logger.error(f"[{operator}] 数据库清理失败: {e}")

    # 2. Redis 清理
    try:
        keys = await redis_kit.keys(pattern)
        if keys:
            logger.info(f"[{operator}] 匹配到 {len(keys)} 个 Redis 键，正在删除...")
            await redis_kit.delete(*keys)
            logger.info(f"[{operator}] Redis 缓存清理完成。")
        else:
            logger.info(f"[{operator}] 未匹配到模式为 '{pattern}' 的 Redis 键。")
    except Exception as e:
        logger.error(f"[{operator}] Redis 清理失败: {e}")

async def main():
    print("\n" + "="*30)
    print("      供应商数据清理工具")
    print("="*30)
    for key, info in VENDORS.items():
        print(f"{key}. {info['name']}")
    print("0. 全部清理")
    print("q. 退出")
    print("-" * 30)
    
    choice = input("请输入选择的编号: ").strip().lower()
    
    if choice == 'q':
        return

    selected_vendors = []
    if choice == '0':
        selected_vendors = list(VENDORS.values())
    elif choice in VENDORS:
        selected_vendors = [VENDORS[choice]]
    else:
        print("无效的选择，请重新运行脚本。")
        return

    # 初始化资源
    db = Db(db_url=DB_URL)
    await db.init_db()
    redis_kit = RedisKit()
    
    try:
        for vendor in selected_vendors:
            await clear_vendor_data(db, redis_kit, vendor)
    finally:
        logger.info("正在关闭数据库连接...")
        await db.close()

if __name__ == "__main__":
    Win32Patch.patch()
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass
    except Exception as e:
        logger.error(f"程序运行异常: {e}")