bot-bottle/bot_bottle/contrib/gitea/forge_state.py

"""Forge state persistence for the orchestrator (PRD prd-new: fold orchestrator).

`ForgeState` is a dataclass that mirrors the orchestrator's `RunRecord`
field-for-field, held here so the store implementation is in bot-bottle
where the Gitea contrib lives.

`SqliteForgeStateStore` backs it with a single SQLite table. The DB path
is optional; passing `None` uses `:memory:` (useful for tests and status
commands that don't need persistence).
"""

from __future__ import annotations

import json
import sqlite3
from dataclasses import dataclass, field
from pathlib import Path


@dataclass
class ForgeState:
    """Persisted state for one forge-targeted issue's bottle lifecycle."""

    owner: str
    repo: str
    issue_number: int
    slug: str
    agent_name: str
    bottle_names: list[str] = field(default_factory=list)
    backend_name: str = ""
    agent_git_user: str = ""
    pr_number: int | None = None
    status: str = ""
    last_checkin_at: str = ""


_DDL = """
CREATE TABLE IF NOT EXISTS forge_state (
    owner        TEXT NOT NULL,
    repo         TEXT NOT NULL,
    issue_number INTEGER NOT NULL,
    slug         TEXT NOT NULL,
    agent_name   TEXT NOT NULL,
    bottle_names TEXT NOT NULL DEFAULT '[]',
    backend_name TEXT NOT NULL DEFAULT '',
    agent_git_user TEXT NOT NULL DEFAULT '',
    pr_number    INTEGER,
    status       TEXT NOT NULL DEFAULT '',
    last_checkin_at TEXT NOT NULL DEFAULT '',
    PRIMARY KEY (owner, repo, issue_number)
)
"""


class SqliteForgeStateStore:
    """SQLite-backed `ForgeState` store.

    Thread-safety: a single connection is used; callers that share a
    store across threads must serialise access externally.
    """

    def __init__(self, db_path: Path | None) -> None:
        path = str(db_path) if db_path is not None else ":memory:"
        self._conn = sqlite3.connect(path, check_same_thread=False)
        self._conn.row_factory = sqlite3.Row
        self._conn.execute(_DDL)
        self._conn.commit()

    def upsert(self, state: ForgeState) -> None:
        self._conn.execute(
            """
            INSERT INTO forge_state
                (owner, repo, issue_number, slug, agent_name,
                 bottle_names, backend_name, agent_git_user,
                 pr_number, status, last_checkin_at)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ON CONFLICT(owner, repo, issue_number) DO UPDATE SET
                slug            = excluded.slug,
                agent_name      = excluded.agent_name,
                bottle_names    = excluded.bottle_names,
                backend_name    = excluded.backend_name,
                agent_git_user  = excluded.agent_git_user,
                pr_number       = excluded.pr_number,
                status          = excluded.status,
                last_checkin_at = excluded.last_checkin_at
            """,
            (
                state.owner,
                state.repo,
                state.issue_number,
                state.slug,
                state.agent_name,
                json.dumps(state.bottle_names),
                state.backend_name,
                state.agent_git_user,
                state.pr_number,
                state.status,
                state.last_checkin_at,
            ),
        )
        self._conn.commit()

    def get(self, owner: str, repo: str, issue_number: int) -> ForgeState | None:
        row = self._conn.execute(
            "SELECT * FROM forge_state WHERE owner=? AND repo=? AND issue_number=?",
            (owner, repo, issue_number),
        ).fetchone()
        return _row_to_state(row) if row is not None else None

    def delete(self, owner: str, repo: str, issue_number: int) -> None:
        self._conn.execute(
            "DELETE FROM forge_state WHERE owner=? AND repo=? AND issue_number=?",
            (owner, repo, issue_number),
        )
        self._conn.commit()

    def all(self) -> list[ForgeState]:
        rows = self._conn.execute(
            "SELECT * FROM forge_state ORDER BY owner, repo, issue_number"
        ).fetchall()
        return [_row_to_state(r) for r in rows]


def _row_to_state(row: sqlite3.Row) -> ForgeState:
    return ForgeState(
        owner=row["owner"],
        repo=row["repo"],
        issue_number=row["issue_number"],
        slug=row["slug"],
        agent_name=row["agent_name"],
        bottle_names=json.loads(row["bottle_names"]),
        backend_name=row["backend_name"],
        agent_git_user=row["agent_git_user"],
        pr_number=row["pr_number"],
        status=row["status"],
        last_checkin_at=row["last_checkin_at"],
    )