ai_work_start/utlit/update_db.py

import sqlite3
import os
from typing import List, Dict, Optional

# 假设 DB_PATH 和 DB_TABLE 在脚本顶部已定义
# DB_PATH = "./kimi_tokens.db"
# DB_TABLE = "tokens"

def upsert_token(db_path: str, table: str, phone: Optional[str], platform_name: str, token_value: str, status: str = "active") -> bool:
    """
    插入或更新单条 token。
    - 如果表有 UNIQUE(phone, platform_name) 约束（推荐），会替换已有记录。
    - 如果没有 UNIQUE 约束，函数会先尝试按 (phone, platform_name) 查找更新，否则插入新行。
    返回 True 表示成功。
    """
    if not os.path.exists(db_path):
        # 如果 db 文件不存在，create table 需事先调用 ensure_tokens_table()
        # 这里仍尝试连接，会在后续创建表时被捕获
        pass

    conn = sqlite3.connect(db_path, timeout=5)
    cur = conn.cursor()
    try:
        # 先尝试按 phone+platform_name 更新（如果 phone 为 None，则只按 platform_name 更新——谨慎使用）
        if phone is not None:
            cur.execute(f"SELECT id FROM {table} WHERE phone = ? AND platform_name = ?", (phone, platform_name))
            row = cur.fetchone()
            if row:
                cur.execute(f"UPDATE {table} SET token = ?, status = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?", (token_value, status, row[0]))
                conn.commit()
                return True
        else:
            # phone is None：按 platform_name 更新最新一条 active/任意记录（谨慎）
            cur.execute(f"SELECT id FROM {table} WHERE platform_name = ? ORDER BY created_at DESC LIMIT 1", (platform_name,))
            row = cur.fetchone()
            if row:
                cur.execute(f"UPDATE {table} SET token = ?, status = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?", (token_value, status, row[0]))
                conn.commit()
                return True

        # 若没找到可更新的记录，插入新行
        cur.execute(f"INSERT INTO {table} (phone, platform_name, token, status) VALUES (?, ?, ?, ?)", (phone, platform_name, token_value, status))
        conn.commit()
        return True
    except sqlite3.IntegrityError:
        # 如果你的表定义了 UNIQUE(phone,platform_name) 并触发了冲突，可以用 replace
        try:
            cur.execute(f"INSERT OR REPLACE INTO {table} (phone, platform_name, token, status) VALUES (?, ?, ?, ?)", (phone, platform_name, token_value, status))
            conn.commit()
            return True
        except Exception as e:
            raise
    finally:
        cur.close()
        conn.close()


def bulk_upsert_tokens(db_path: str, table: str, rows: List[Dict[str, str]]) -> int:
    """
    批量插入/更新 tokens。rows 是 dict 列表，每项包含 keys: phone, platform_name, token, status (可选)
    返回成功处理的行数。
    示例 rows:
      [{"phone":"138...", "platform_name":"kimi", "token":"xxx"}, {...}, ...]
    """
    if not rows:
        return 0
    conn = sqlite3.connect(db_path, timeout=10)
    cur = conn.cursor()
    processed = 0
    try:
        # 使用事务加速批量写入
        cur.execute("BEGIN")
        for r in rows:
            phone = r.get("phone")
            platform_name = r["platform_name"]
            token_value = r["token"]
            status = r.get("status", "active")
            # 尝试按 phone+platform 更新
            if phone is not None:
                cur.execute(f"SELECT id FROM {table} WHERE phone = ? AND platform_name = ?", (phone, platform_name))
                row = cur.fetchone()
                if row:
                    cur.execute(f"UPDATE {table} SET token = ?, status = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?", (token_value, status, row[0]))
                    processed += 1
                    continue
            else:
                cur.execute(f"SELECT id FROM {table} WHERE platform_name = ? ORDER BY created_at DESC LIMIT 1", (platform_name,))
                row = cur.fetchone()
                if row:
                    cur.execute(f"UPDATE {table} SET token = ?, status = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?", (token_value, status, row[0]))
                    processed += 1
                    continue
            # 否则插入
            cur.execute(f"INSERT INTO {table} (phone, platform_name, token, status) VALUES (?, ?, ?, ?)", (phone, platform_name, token_value, status))
            processed += 1
        conn.commit()
        return processed
    except Exception as e:
        conn.rollback()
        raise
    finally:
        cur.close()
        conn.close()

def query_tokens(platform_name: str = None, status: str = None):
    TABLE  =  "tokens"
    conn = sqlite3.connect("./kimi_tokens.db")
    cur = conn.cursor()
    sql = f"SELECT id, phone, platform_name, token, status, created_at, updated_at FROM {TABLE} WHERE 1=1"
    params = []
    if platform_name:
        sql += " AND platform_name=?"
        params.append(platform_name)
    if status:
        sql += " AND status=?"
        params.append(status)
    sql += " ORDER BY created_at ASC"
    cur.execute(sql, params)
    rows = cur.fetchall()
    cur.close()
    conn.close()
    for row in rows:
        print(row)

# 示例
# query_tokens(platform_name="kimi", status="active")
# -------------------------
# 示例用法
# -------------------------
if __name__ == "__main__":
    # 假设 ensure_tokens_table() 已经被调用过
    DB = "./kimi_tokens.db"
    TABLE = "tokens"

    # 单条 upsert
    ok = upsert_token(DB, TABLE, phone="16223053815", platform_name="mita", token_value="JSESSIONID=356E6FEEBF81933FAFE924ED6A7C324C; aliyungf_tc=70df8ace7303f22dccbacee689c1abc32b0a7629a28579fdfefad08c7d04f536; tid=cc159aeb-4f9b-463a-9080-2ebc2d88f0bc; __eventn_id_UMO2dYNwFz=065s1a391e; traceid=2dc94bfd49cc4992; sid=c764bb074d3b496fbeb4ad5eaa8835f0; uid=69255736dddbb6ba7df1905c", status="active")
    print("upsert single:", ok)
    query_tokens(platform_name="mita", status="active")
    # # 批量 upsert
    # rows = [
    #     {"phone": "13800138000", "platform_name": "kimi", "token": "token_a"},
    #     {"phone": "13800138001", "platform_name": "kimi", "token": "token_b"},
    #     {"phone": None, "platform_name": "other", "token": "token_c"},
    # ]
    # n = bulk_upsert_tokens(DB, TABLE, rows)
    # print("bulk processed:", n)