rakuten-api/app/gateway/db.py

"""网关 SQLite 访问层：三张表 + 原生 SQL，无业务逻辑

业务规则（状态机、并发度、幂等）放在 `task_queue.py`，本模块只负责持久化与查询，
返回 dataclass 行。所有时间戳以 ISO8601 UTC 字符串存储（以 `Z` 结尾），便于
跨进程对账；时间戳运算在 task_queue 层完成。

写操作由 task_queue 的 asyncio.Lock 串行化（详见该模块），本层不重复加锁，
因此**调用方必须确保写操作在外层锁的保护下进行**。
"""
from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import aiosqlite

logger = logging.getLogger(__name__)

SCHEMA = """
CREATE TABLE IF NOT EXISTS tasks (
    task_id          TEXT PRIMARY KEY,
    site             TEXT NOT NULL,
    intent_json      TEXT NOT NULL,
    status           TEXT NOT NULL,
    lease_owner      TEXT,
    lease_expires_at TEXT,
    lease_count      INTEGER NOT NULL DEFAULT 0,
    created_at       TEXT NOT NULL,
    updated_at       TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_tasks_pending ON tasks(status, created_at);

CREATE TABLE IF NOT EXISTS task_reports (
    task_id       TEXT NOT NULL,
    state         TEXT NOT NULL,
    payable_yen   INTEGER,
    pay_deadline  TEXT,
    site_order_id TEXT,
    evidence_ref  TEXT,
    detail        TEXT,
    reported_at   TEXT NOT NULL,
    PRIMARY KEY (task_id, state)
);

CREATE TABLE IF NOT EXISTS workers (
    worker_id    TEXT PRIMARY KEY,
    last_seen_at TEXT NOT NULL
);
"""


@dataclass(slots=True)
class TaskRow:
    """tasks 表一行的强类型视图"""

    task_id: str
    site: str
    intent_json: str
    status: str
    lease_owner: str | None
    lease_expires_at: str | None
    lease_count: int
    created_at: str
    updated_at: str

    @property
    def intent(self) -> dict[str, Any]:
        return json.loads(self.intent_json)


@dataclass(slots=True)
class ReportRow:
    """task_reports 表一行的强类型视图"""

    task_id: str
    state: str
    payable_yen: int | None
    pay_deadline: str | None
    site_order_id: str | None
    evidence_ref: str | None
    detail: str
    reported_at: str


@dataclass(slots=True)
class WorkerRow:
    """workers 表一行的强类型视图"""

    worker_id: str
    last_seen_at: str


def _row_to_task(row: aiosqlite.Row) -> TaskRow:
    return TaskRow(
        task_id=row["task_id"],
        site=row["site"],
        intent_json=row["intent_json"],
        status=row["status"],
        lease_owner=row["lease_owner"],
        lease_expires_at=row["lease_expires_at"],
        lease_count=row["lease_count"],
        created_at=row["created_at"],
        updated_at=row["updated_at"],
    )


def _row_to_report(row: aiosqlite.Row) -> ReportRow:
    return ReportRow(
        task_id=row["task_id"],
        state=row["state"],
        payable_yen=row["payable_yen"],
        pay_deadline=row["pay_deadline"],
        site_order_id=row["site_order_id"],
        evidence_ref=row["evidence_ref"],
        detail=row["detail"],
        reported_at=row["reported_at"],
    )


def _row_to_worker(row: aiosqlite.Row) -> WorkerRow:
    return WorkerRow(worker_id=row["worker_id"], last_seen_at=row["last_seen_at"])


class GatewayDB:
    """网关 SQLite 访问对象

    单连接 + aiosqlite 的内部线程。生命周期由 `GatewayContainer` 管理：
    `start()` 在 lifespan 启动时调用，`close()` 在关闭时调用。
    """

    def __init__(self, db_path: Path):
        self._db_path = db_path
        self._conn: aiosqlite.Connection | None = None

    async def start(self) -> None:
        """打开连接并初始化 schema（IF NOT EXISTS，可重复执行）"""
        self._db_path.parent.mkdir(parents=True, exist_ok=True)
        self._conn = await aiosqlite.connect(str(self._db_path))
        self._conn.row_factory = aiosqlite.Row
        await self._conn.executescript(SCHEMA)
        await self._conn.commit()
        logger.info("网关 DB 已就绪：%s", self._db_path)

    async def close(self) -> None:
        if self._conn is not None:
            await self._conn.close()
            self._conn = None

    @property
    def conn(self) -> aiosqlite.Connection:
        if self._conn is None:
            raise RuntimeError("GatewayDB 未启动：先调用 start()")
        return self._conn

    # ---- tasks ----

    async def get_task(self, task_id: str) -> TaskRow | None:
        async with self.conn.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)) as cur:
            row = await cur.fetchone()
            return _row_to_task(row) if row else None

    async def insert_task(self, row: TaskRow) -> bool:
        """插入新任务。返回 True=新建，False=task_id 已存在（幂等命中）"""
        try:
            await self.conn.execute(
                "INSERT INTO tasks (task_id, site, intent_json, status, "
                "lease_owner, lease_expires_at, lease_count, created_at, updated_at) "
                "VALUES (?, ?, ?, ?, NULL, NULL, 0, ?, ?)",
                (
                    row.task_id,
                    row.site,
                    row.intent_json,
                    row.status,
                    row.created_at,
                    row.updated_at,
                ),
            )
            await self.conn.commit()
            return True
        except aiosqlite.IntegrityError:
            return False

    async def update_task(
        self,
        task_id: str,
        *,
        status: str | None = None,
        lease_owner: str | None = None,
        lease_expires_at: str | None = None,
        lease_count: int | None = None,
        updated_at: str | None = None,
        clear_lease: bool = False,
    ) -> None:
        """更新任务字段。clear_lease=True 时把 lease_owner/expires_at 置 NULL"""
        sets: list[str] = []
        params: list[Any] = []
        if status is not None:
            sets.append("status = ?")
            params.append(status)
        if lease_owner is not None:
            sets.append("lease_owner = ?")
            params.append(lease_owner)
        if lease_expires_at is not None:
            sets.append("lease_expires_at = ?")
            params.append(lease_expires_at)
        if lease_count is not None:
            sets.append("lease_count = ?")
            params.append(lease_count)
        if updated_at is not None:
            sets.append("updated_at = ?")
            params.append(updated_at)
        if clear_lease:
            sets.append("lease_owner = NULL")
            sets.append("lease_expires_at = NULL")
        if not sets:
            return
        params.append(task_id)
        await self.conn.execute(
            f"UPDATE tasks SET {', '.join(sets)} WHERE task_id = ?",
            params,
        )
        await self.conn.commit()

    async def list_tasks(
        self,
        *,
        status: str | None = None,
        site: str | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> tuple[list[TaskRow], int]:
        """分页列出任务，按 created_at 升序"""
        where = []
        params: list[Any] = []
        if status:
            where.append("status = ?")
            params.append(status)
        if site:
            where.append("site = ?")
            params.append(site)
        clause = f"WHERE {' AND '.join(where)}" if where else ""

        async with self.conn.execute(
            f"SELECT COUNT(*) FROM tasks {clause}",
            params,
        ) as cur:
            total = (await cur.fetchone())[0]

        sql = f"SELECT * FROM tasks {clause} ORDER BY created_at ASC LIMIT ? OFFSET ?"
        async with self.conn.execute(sql, [*params, limit, offset]) as cur:
            rows = await cur.fetchall()
        return [_row_to_task(r) for r in rows], total

    async def list_tasks_in_statuses(self, statuses: tuple[str, ...]) -> list[TaskRow]:
        """取处于给定状态集合的全部任务（不分页，用于健康检查的 active 列表）"""
        if not statuses:
            return []
        placeholders = ", ".join("?" for _ in statuses)
        sql = (
            f"SELECT * FROM tasks WHERE status IN ({placeholders}) "
            "ORDER BY created_at ASC"
        )
        async with self.conn.execute(sql, statuses) as cur:
            rows = await cur.fetchall()
        return [_row_to_task(r) for r in rows]

    async def pick_queued_task(self, site: str | None) -> TaskRow | None:
        """取最早的 queued 任务，可选按站点过滤"""
        if site:
            sql = "SELECT * FROM tasks WHERE status = ? AND site = ? ORDER BY created_at ASC LIMIT 1"
            params: tuple[Any, ...] = ("queued", site)
        else:
            sql = "SELECT * FROM tasks WHERE status = ? ORDER BY created_at ASC LIMIT 1"
            params = ("queued",)
        async with self.conn.execute(sql, params) as cur:
            row = await cur.fetchone()
            return _row_to_task(row) if row else None

    # ---- task_reports ----

    async def upsert_report(self, row: ReportRow) -> bool:
        """写入一条 report。同一 (task_id, state) 已存在则覆盖（幂等）。

        返回 True=本次新写入，False=覆盖了既有行。

        先 SELECT 再 INSERT OR REPLACE，避免依赖 SQLite rowcount——后者在
        INSERT OR REPLACE 上即使命中既有行也报 1，无法区分。
        """
        async with self.conn.execute(
            "SELECT 1 FROM task_reports WHERE task_id = ? AND state = ?",
            (row.task_id, row.state),
        ) as cur:
            existed = await cur.fetchone() is not None

        await self.conn.execute(
            "INSERT OR REPLACE INTO task_reports "
            "(task_id, state, payable_yen, pay_deadline, site_order_id, "
            "evidence_ref, detail, reported_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            (
                row.task_id,
                row.state,
                row.payable_yen,
                row.pay_deadline,
                row.site_order_id,
                row.evidence_ref,
                row.detail,
                row.reported_at,
            ),
        )
        await self.conn.commit()
        return not existed

    async def get_reports(self, task_id: str) -> list[ReportRow]:
        async with self.conn.execute(
            "SELECT * FROM task_reports WHERE task_id = ? ORDER BY reported_at ASC",
            (task_id,),
        ) as cur:
            rows = await cur.fetchall()
        return [_row_to_report(r) for r in rows]

    async def get_latest_state(self, task_id: str) -> str | None:
        """取该任务最近一次 report 的 state，无 report 返回 None"""
        async with self.conn.execute(
            "SELECT state FROM task_reports WHERE task_id = ? "
            "ORDER BY reported_at DESC LIMIT 1",
            (task_id,),
        ) as cur:
            row = await cur.fetchone()
            return row["state"] if row else None

    # ---- workers ----

    async def upsert_worker(self, worker_id: str, last_seen_at: str) -> None:
        await self.conn.execute(
            "INSERT INTO workers (worker_id, last_seen_at) VALUES (?, ?) "
            "ON CONFLICT(worker_id) DO UPDATE SET last_seen_at = excluded.last_seen_at",
            (worker_id, last_seen_at),
        )
        await self.conn.commit()

    async def list_workers(self) -> list[WorkerRow]:
        async with self.conn.execute("SELECT * FROM workers ORDER BY last_seen_at DESC") as cur:
            rows = await cur.fetchall()
        return [_row_to_worker(r) for r in rows]

    # ---- 计数 ----

    async def count_by_status(self, status: str) -> int:
        async with self.conn.execute(
            "SELECT COUNT(*) FROM tasks WHERE status = ?",
            (status,),
        ) as cur:
            return (await cur.fetchone())[0]