crypto_monitor/lib/account_ledger/account_ledger_db.py

"""账户流水 SQLite 缓存."""
from __future__ import annotations

import time
from typing import Any, Optional

from lib.account_ledger.account_ledger_normalize import PAGE_SIZE, VALID_ACCOUNTS


def ensure_account_ledger_tables(conn) -> None:
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS account_ledger_entries (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            account TEXT NOT NULL,
            ccy TEXT NOT NULL,
            amount REAL NOT NULL,
            balance_after REAL,
            kind TEXT,
            raw_type TEXT,
            symbol TEXT,
            ref_id TEXT NOT NULL,
            ts_ms INTEGER NOT NULL,
            note TEXT,
            synced_at REAL,
            UNIQUE(account, ref_id, ccy, ts_ms)
        )
        """
    )
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_account_ledger_acc_ts "
        "ON account_ledger_entries(account, ts_ms DESC)"
    )
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS account_ledger_meta (
            key TEXT PRIMARY KEY,
            value TEXT
        )
        """
    )
    conn.commit()


def meta_get(conn, key: str, default: str = "") -> str:
    row = conn.execute(
        "SELECT value FROM account_ledger_meta WHERE key=?", (key,)
    ).fetchone()
    if not row:
        return default
    try:
        return str(row[0] if not hasattr(row, "keys") else row["value"])
    except Exception:
        return default


def meta_set(conn, key: str, value: str) -> None:
    conn.execute(
        "INSERT INTO account_ledger_meta(key, value) VALUES(?, ?) "
        "ON CONFLICT(key) DO UPDATE SET value=excluded.value",
        (key, str(value)),
    )


def upsert_entries(conn, rows: list[dict[str, Any]]) -> int:
    if not rows:
        return 0
    now = time.time()
    n = 0
    for r in rows:
        try:
            conn.execute(
                """
                INSERT INTO account_ledger_entries(
                    account, ccy, amount, balance_after, kind, raw_type,
                    symbol, ref_id, ts_ms, note, synced_at
                ) VALUES (?,?,?,?,?,?,?,?,?,?,?)
                ON CONFLICT(account, ref_id, ccy, ts_ms) DO UPDATE SET
                    amount=excluded.amount,
                    balance_after=excluded.balance_after,
                    kind=excluded.kind,
                    raw_type=excluded.raw_type,
                    symbol=excluded.symbol,
                    note=excluded.note,
                    synced_at=excluded.synced_at
                """,
                (
                    r["account"],
                    r["ccy"],
                    float(r["amount"]),
                    r.get("balance_after"),
                    r.get("kind") or "other",
                    r.get("raw_type") or "",
                    r.get("symbol") or "",
                    r["ref_id"],
                    int(r["ts_ms"]),
                    r.get("note") or "",
                    now,
                ),
            )
            n += 1
        except Exception:
            continue
    conn.commit()
    return n


def query_entries(
    conn,
    *,
    account: str,
    start_ms: int,
    end_ms: int,
    page: int = 1,
    page_size: int = PAGE_SIZE,
    currencies: Optional[list[str]] = None,
) -> dict[str, Any]:
    acc = (account or "").strip().lower()
    if acc not in VALID_ACCOUNTS:
        return {"items": [], "total": 0, "page": 1, "page_size": page_size, "pages": 0}
    page = max(1, int(page or 1))
    page_size = max(1, min(50, int(page_size or PAGE_SIZE)))
    start_ms = int(start_ms)
    end_ms = int(end_ms)
    params: list[Any] = [acc, start_ms, end_ms]
    ccy_sql = ""
    if currencies:
        ccy_list = [c.strip().upper() for c in currencies if c and str(c).strip()]
        if ccy_list:
            placeholders = ",".join("?" for _ in ccy_list)
            ccy_sql = f" AND ccy IN ({placeholders})"
            params.extend(ccy_list)
    total = conn.execute(
        f"SELECT COUNT(*) FROM account_ledger_entries "
        f"WHERE account=? AND ts_ms>=? AND ts_ms<=?{ccy_sql}",
        params,
    ).fetchone()[0]
    total = int(total or 0)
    pages = (total + page_size - 1) // page_size if total else 0
    if pages and page > pages:
        page = pages
    offset = (page - 1) * page_size
    rows = conn.execute(
        f"""
        SELECT account, ccy, amount, balance_after, kind, raw_type, symbol,
               ref_id, ts_ms, note
        FROM account_ledger_entries
        WHERE account=? AND ts_ms>=? AND ts_ms<=?{ccy_sql}
        ORDER BY ts_ms DESC, id DESC
        LIMIT ? OFFSET ?
        """,
        params + [page_size, offset],
    ).fetchall()
    items = []
    for r in rows:
        if hasattr(r, "keys"):
            d = {k: r[k] for k in r.keys()}
        else:
            d = {
                "account": r[0],
                "ccy": r[1],
                "amount": r[2],
                "balance_after": r[3],
                "kind": r[4],
                "raw_type": r[5],
                "symbol": r[6],
                "ref_id": r[7],
                "ts_ms": r[8],
                "note": r[9],
            }
        from lib.account_ledger.account_ledger_normalize import kind_label_zh

        d["kind_label"] = kind_label_zh(d.get("kind") or "")
        items.append(d)
    return {
        "items": items,
        "total": total,
        "page": page,
        "page_size": page_size,
        "pages": pages,
    }


def prune_older_than(conn, min_ts_ms: int) -> None:
    conn.execute("DELETE FROM account_ledger_entries WHERE ts_ms < ?", (int(min_ts_ms),))
    conn.commit()