buff3.0/cleanup-api-keys/cleanup_api_keys.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
清理 api_keys 表中的过期数据
- 每天执行一次
- 删除两天前的所有数据（基于 created_at 字段）
- 记录清理结果到日志
"""
import os
import sys
import time
import configparser
from datetime import datetime, timezone, timedelta
from pathlib import Path

import pymysql

# 北京时间 UTC+8，与 backend-client 一致
BEIJING_TZ = timezone(timedelta(hours=8))


def beijing_now() -> datetime:
    """当前北京时间（naive）"""
    return datetime.now(BEIJING_TZ).replace(tzinfo=None)


def load_config():
    """从 config.ini 加载配置，返回 config 对象"""
    conf_path = Path(__file__).resolve().parent / "config.ini"
    if not conf_path.exists():
        print(f"错误: 未找到配置文件 {conf_path}，请复制 config.ini.example 为 config.ini 并填写", file=sys.stderr)
        sys.exit(1)
    config = configparser.ConfigParser()
    config.read(conf_path, encoding="utf-8")
    return config


def get_db_connection(config):
    """创建数据库连接"""
    return pymysql.connect(
        host=config.get("database", "host"),
        port=config.getint("database", "port"),
        user=config.get("database", "user"),
        password=config.get("database", "password"),
        database=config.get("database", "database"),
        charset=config.get("database", "charset"),
        cursorclass=pymysql.cursors.DictCursor,
    )


def run_once(config) -> int:
    """
    执行一次清理：删除两天前的 api_keys 数据
    返回: 删除的记录数量
    """
    now_dt = beijing_now()
    now_str = now_dt.strftime("%Y-%m-%d %H:%M:%S")

    # 计算两天前的时间戳
    two_days_ago = now_dt - timedelta(days=2)
    two_days_ago_str = two_days_ago.strftime("%Y-%m-%d %H:%M:%S")

    conn = get_db_connection(config)
    try:
        with conn.cursor() as cursor:
            # 先查询要删除的数据数量
            cursor.execute(
                "SELECT COUNT(*) as delete_count FROM api_keys WHERE created_at < %s",
                (two_days_ago_str,)
            )
            delete_count = cursor.fetchone()["delete_count"]

            if delete_count == 0:
                print(f"[{now_str}] 无两天前数据需要清理")
                return 0

            # 删除两天前的数据
            cursor.execute(
                "DELETE FROM api_keys WHERE created_at < %s",
                (two_days_ago_str,)
            )
            deleted_count = cursor.rowcount
            conn.commit()

            # 查询剩余记录数
            cursor.execute("SELECT COUNT(*) as remaining_count FROM api_keys")
            remaining_count = cursor.fetchone()["remaining_count"]

            print(f"[{now_str}] 清理完成: 删除 {deleted_count} 条两天前数据 (cutoff: {two_days_ago_str}), 剩余 {remaining_count} 条")
            return deleted_count
    finally:
        conn.close()


def main():
    config = load_config()
    interval_hours = config.getint("script", "interval_hours", fallback=24)

    print(f"=" * 50)
    print(f"api_keys 清理脚本已启动")
    print(f"清理周期: 每 {interval_hours} 小时执行一次")
    print(f"清理条件: 删除两天前的所有数据")
    print(f"启动时间: {beijing_now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"=" * 50)

    # 启动时先执行一次清理
    try:
        run_once(config)
    except Exception as e:
        print(f"[{beijing_now()}] 启动时清理失败: {e}", file=sys.stderr)

    print(f"\n开始定时清理，每 {interval_hours} 小时执行一次...\n")

    while True:
        time.sleep(interval_hours * 3600)
        try:
            run_once(config)
        except Exception as e:
            print(f"[{beijing_now()}] 清理异常: {e}", file=sys.stderr)


if __name__ == "__main__":
    main()