cleveragents-core/tools/controller/db/models.py

"""SQLAlchemy 2.0 declarative models for the controller DB.

Schema follows plan v6/v9: unified ``workflows`` table with ``kind``
discriminator (pr/issue), workflow_attempts that doubles as the worker
queue (status + lock columns), controller_events for audit + Forgejo
replay, flake_history for CI flake learning, ci_observations for raw
CI state history.

Cross-dialect (SQLite + Postgres):
- ``JSON`` columns map to JSON on SQLite and JSONB on Postgres.
- ``Integer``/``BigInteger`` are dialect-portable; we use
  ``BigInteger`` for PKs that will grow (attempts) and ``Integer``
  for workflow ids (one-per-PR; bounded).
- Partial indexes use the SQLAlchemy ``Index(..., postgresql_where=)``
  + ``sqlite_where=`` pattern so both dialects honor them.
- ``CHECK`` constraints use SQLAlchemy ``CheckConstraint`` (both
  dialects).
- Per the v6 schema decision, the parent_workflow_id kind-validation
  is app-level (SQLite can't do subquery-CHECK; Postgres trigger
  defer to v6 plan body).
"""
from __future__ import annotations

from datetime import datetime
from typing import Any

from sqlalchemy import (
    BigInteger,
    Boolean,
    CheckConstraint,
    DateTime,
    Float,
    ForeignKey,
    Index,
    Integer,
    String,
    Text,
    UniqueConstraint,
    func,
    text,
)
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import (
    DeclarativeBase,
    Mapped,
    mapped_column,
    relationship,
)
from sqlalchemy.types import JSON


# Cross-dialect JSON: JSONB on Postgres, JSON on SQLite. JSONB-only
# features (GIN indexes, jsonb_path_ops) are not used; v1 stays
# portable.
JsonColumn = JSON().with_variant(JSONB(), "postgresql")

# Cross-dialect autoincrement PK: SQLite's INTEGER (4-byte rowid alias
# that auto-increments); Postgres's BIGSERIAL (8-byte). SQLAlchemy's
# BigInteger doesn't autoincrement on SQLite, so this variant gives us
# the right type on each dialect.
AutoincrementPk = Integer().with_variant(BigInteger(), "postgresql")


class Base(DeclarativeBase):
    """SQLAlchemy 2.0 declarative base for controller tables."""
    type_annotation_map = {dict[str, Any]: JsonColumn}


# ─── workflows ────────────────────────────────────────────────────────


class Workflow(Base):
    """One PR or issue under controller management.

    Plan v6 unified `workflows` table (replaces v5's separate
    pr_workflows / issue_workflows) with `kind` discriminator. An
    issue workflow can spawn a PR workflow via `parent_workflow_id`
    on the new row (CREATED_PR transition).
    """

    __tablename__ = "workflows"

    workflow_id: Mapped[int] = mapped_column(AutoincrementPk, primary_key=True, autoincrement=True)

    # Discriminator + identity
    kind: Mapped[str] = mapped_column(String(8), nullable=False)  # 'pr' | 'issue'
    owner: Mapped[str] = mapped_column(String(128), nullable=False)
    repo: Mapped[str] = mapped_column(String(128), nullable=False)
    entity_number: Mapped[int] = mapped_column(Integer, nullable=False)
    parent_workflow_id: Mapped[int | None] = mapped_column(
        AutoincrementPk,
        ForeignKey("workflows.workflow_id", ondelete="RESTRICT"),
        nullable=True,
    )

    # State machine
    current_state: Mapped[str] = mapped_column(String(32), nullable=False)
    current_tier: Mapped[int | None] = mapped_column(Integer, nullable=True)
    tier_last_succeeded: Mapped[int | None] = mapped_column(Integer, nullable=True)

    # Timestamps
    started_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    last_transition_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    entered_state_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )

    # Limits + workspace
    max_attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=6)
    workspace_dir: Mapped[str | None] = mapped_column(Text, nullable=True)

    # MERGING non-blocking retries
    merging_retry_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    merging_retry_next_attempt_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )

    # Phase 1k+: pre-pause state for label-gate PAUSE/RESUME. When the
    # operator removes CONTROLLER_OPT_IN_LABEL, the workflow's current
    # state moves to PAUSED and the prior state is captured here. On
    # label re-add, the master reads this column to determine the
    # resume target. NULL means "never paused" or "resumed".
    pre_pause_state: Mapped[str | None] = mapped_column(
        String(32), nullable=True
    )

    # Phase 1k++ (R3): the ``ci_flake_retries_remaining`` and
    # ``awaiting_ci_started_at`` columns added in Phase 1k were never
    # actually read or written by any producer — the round-3 review
    # caught them as dead schema weight. ``ci_poll.py`` uses
    # ``entered_state_at`` (already populated on every transition) as
    # the AWAITING_CI start timestamp; the flake-retries counter is a
    # future-phase item that should ship with its producer, not as
    # speculative schema.

    # Relationships
    attempts: Mapped[list[WorkflowAttempt]] = relationship(
        back_populates="workflow", cascade="all, delete-orphan"
    )

    __table_args__ = (
        CheckConstraint(
            "kind IN ('pr', 'issue')",
            name="ck_workflows_kind_valid",
        ),
        UniqueConstraint(
            "owner", "repo", "kind", "entity_number",
            name="uq_workflows_entity",
        ),
        # Filtered indexes — fast worklist queries skip terminal rows.
        Index(
            "ix_workflows_active",
            "current_state",
            postgresql_where=text(
                "current_state NOT IN ('MERGED', 'ABANDONED', 'STUCK', 'CREATED_PR')"
            ),
            sqlite_where=text(
                "current_state NOT IN ('MERGED', 'ABANDONED', 'STUCK', 'CREATED_PR')"
            ),
        ),
        Index(
            "ix_workflows_merging_retry",
            "merging_retry_next_attempt_at",
            postgresql_where=text("current_state = 'MERGING'"),
            sqlite_where=text("current_state = 'MERGING'"),
        ),
    )


# ─── workflow_attempts (IS the worker queue) ──────────────────────────


class WorkflowAttempt(Base):
    """One attempt by one role (estimator / implementer / reviewer /
    conflict_resolver / summarizer) against one workflow.

    Doubles as the worker queue: workers dequeue rows where
    ``status='pending'`` via ``dequeue.py`` (Postgres FOR UPDATE
    SKIP LOCKED, SQLite BEGIN IMMEDIATE). Lock columns (status,
    locked_by_instance, locked_at, lock_heartbeat_at, lock_ttl_seconds,
    pickup_count) implement the v9 multi-machine-safe lock protocol.
    """

    __tablename__ = "workflow_attempts"

    attempt_id: Mapped[int] = mapped_column(AutoincrementPk, primary_key=True, autoincrement=True)
    workflow_id: Mapped[int] = mapped_column(
        AutoincrementPk,
        ForeignKey("workflows.workflow_id", ondelete="CASCADE"),
        nullable=False,
    )
    attempt_number: Mapped[int] = mapped_column(Integer, nullable=False)
    role: Mapped[str] = mapped_column(String(32), nullable=False)
    tier: Mapped[int | None] = mapped_column(Integer, nullable=True)

    # v5 queue + lock columns
    status: Mapped[str] = mapped_column(String(16), nullable=False, default="pending")
    locked_by_instance: Mapped[str | None] = mapped_column(String(256), nullable=True)
    locked_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    lock_heartbeat_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True), nullable=True
    )
    lock_ttl_seconds: Mapped[int] = mapped_column(Integer, nullable=False, default=600)
    pickup_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    wallclock_seconds: Mapped[float | None] = mapped_column(Float, nullable=True)

    # SHA bookkeeping
    head_sha_before: Mapped[str | None] = mapped_column(String(64), nullable=True)
    head_sha_after: Mapped[str | None] = mapped_column(String(64), nullable=True)

    # Payloads (JSONB on Postgres, JSON on SQLite). Application caps
    # input_payload at 4MB (see payload_guard.py).
    input_payload: Mapped[dict[str, Any]] = mapped_column(JsonColumn, nullable=False)
    input_version: Mapped[str] = mapped_column(String(8), nullable=False)
    input_payload_truncated: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
    output_payload: Mapped[dict[str, Any] | None] = mapped_column(JsonColumn, nullable=True)
    output_version: Mapped[str | None] = mapped_column(String(8), nullable=True)
    outcome: Mapped[str | None] = mapped_column(String(64), nullable=True)

    # Worker metadata
    session_id: Mapped[str | None] = mapped_column(String(128), nullable=True)
    worker_subprocess_pid: Mapped[int | None] = mapped_column(Integer, nullable=True)
    strict_parse_retries: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

    workflow: Mapped[Workflow] = relationship(back_populates="attempts")

    __table_args__ = (
        CheckConstraint(
            "status IN ('pending', 'in_progress', 'complete', 'failed', 'reaped')",
            name="ck_attempts_status_valid",
        ),
        CheckConstraint("pickup_count >= 0", name="ck_attempts_pickup_count_nonneg"),
        # Worker-dequeue index: hot path of every worker tick.
        Index(
            "ix_attempts_pending_by_role",
            "role", "created_at",
            postgresql_where=text("status = 'pending'"),
            sqlite_where=text("status = 'pending'"),
        ),
        # Reaper index: master tick scans this every 60s.
        Index(
            "ix_attempts_in_progress_by_heartbeat",
            "lock_heartbeat_at",
            postgresql_where=text("status = 'in_progress'"),
            sqlite_where=text("status = 'in_progress'"),
        ),
        # Master driver: scans complete attempts to advance state machine.
        Index(
            "ix_attempts_complete_by_workflow",
            "workflow_id", "attempt_number",
            postgresql_where=text("status IN ('complete', 'failed')"),
            sqlite_where=text("status IN ('complete', 'failed')"),
        ),
    )


# ─── controller_events ────────────────────────────────────────────────


class ControllerEvent(Base):
    """Audit trail + Forgejo-write-replay log.

    Every state transition is a row. Forgejo writes (status comments,
    labels, merges) live here so a crash between DB commit + Forgejo
    write can be replayed on next master tick (v9 simplified protocol:
    Forgejo-first, DB-second, reconciliation syncs DB from Forgejo).

    For v1 ``forgejo_write_pending`` is unused (the v9 simplification
    dropped the 3-step protocol); kept in the schema so a future
    upgrade can re-enable it without a migration.
    """

    __tablename__ = "controller_events"

    event_id: Mapped[int] = mapped_column(AutoincrementPk, primary_key=True, autoincrement=True)
    workflow_id: Mapped[int] = mapped_column(
        AutoincrementPk,
        ForeignKey("workflows.workflow_id", ondelete="CASCADE"),
        nullable=False,
    )
    ts: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    event_type: Mapped[str] = mapped_column(String(64), nullable=False)
    from_state: Mapped[str | None] = mapped_column(String(32), nullable=True)
    to_state: Mapped[str | None] = mapped_column(String(32), nullable=True)
    attempt_id: Mapped[int | None] = mapped_column(
        AutoincrementPk,
        ForeignKey("workflow_attempts.attempt_id", ondelete="SET NULL"),
        nullable=True,
    )
    payload: Mapped[dict[str, Any] | None] = mapped_column(JsonColumn, nullable=True)

    # v3 write protocol fields (kept for v2-style upgrade path)
    forgejo_write_pending: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
    forgejo_fingerprint: Mapped[str | None] = mapped_column(String(128), nullable=True)
    forgejo_result: Mapped[dict[str, Any] | None] = mapped_column(JsonColumn, nullable=True)
    replay_attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0)

    __table_args__ = (
        Index(
            "ix_events_pending_forgejo",
            "workflow_id", "ts",
            postgresql_where=text("forgejo_write_pending = true"),
            sqlite_where=text("forgejo_write_pending = 1"),
        ),
        Index("ix_events_by_workflow_ts", "workflow_id", "ts"),
    )


# ─── flake_history ────────────────────────────────────────────────────


class FlakeHistory(Base):
    """CI flake learning — per-(owner, repo, context) flake-score
    accumulator. Plan v6: ship empty + learning heuristic; bootstrap
    via first-failure mandatory retry. flake_score = flake_count /
    observations; threshold 0.6.
    """

    __tablename__ = "flake_history"

    owner: Mapped[str] = mapped_column(String(128), primary_key=True)
    repo: Mapped[str] = mapped_column(String(128), primary_key=True)
    context: Mapped[str] = mapped_column(String(256), primary_key=True)

    observations: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    flake_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    flake_score: Mapped[float] = mapped_column(Float, nullable=False, default=0.0)
    last_seen: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    operator_override: Mapped[bool | None] = mapped_column(Boolean, nullable=True)


# ─── ci_observations ──────────────────────────────────────────────────


class CIObservation(Base):
    """Raw CI status observed for a head_sha at a point in time.

    Append-only; used for parser-rot detection + flake history input
    + post-mortem debugging. Cleaned up via periodic retention sweep
    (90 days).
    """

    __tablename__ = "ci_observations"

    observation_id: Mapped[int] = mapped_column(AutoincrementPk, primary_key=True, autoincrement=True)
    head_sha: Mapped[str] = mapped_column(String(64), nullable=False)
    observed_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    state: Mapped[str] = mapped_column(String(16), nullable=False)
    contexts: Mapped[dict[str, Any]] = mapped_column(JsonColumn, nullable=False)

    __table_args__ = (
        Index("ix_ci_obs_by_sha", "head_sha", "observed_at"),
    )