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

"""Worker dequeue helper.

The worker picks one pending ``workflow_attempts`` row, marks it
``in_progress`` with lock metadata, and returns it. The atomic
"pick + mark" is the core multi-machine coordination primitive.

Per plan v5/v6:
- Postgres: ``SELECT … FOR UPDATE SKIP LOCKED LIMIT 1`` is the
  canonical primitive. Multiple workers race; only one acquires
  each row.
- SQLite: no SKIP LOCKED. Workers serialise via ``BEGIN IMMEDIATE``
  + ``UPDATE … WHERE status='pending'`` (RETURNING clause; SQLite
  3.35+). For single-machine SQLite tests this is fine; multi-machine
  SQLite over NFS is explicitly NOT supported (per plan v9: Postgres
  is required for multi-machine).

Pickup-count guard (v6 blocker): the dequeue UPDATE bumps
``pickup_count``. A separate query (master-side) finds attempts
where ``pickup_count >= MAX_PICKUPS`` and transitions the workflow
to STUCK.
"""
from __future__ import annotations

import os
from dataclasses import dataclass
from datetime import datetime, timezone

from sqlalchemy import bindparam, text
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session

from .models import WorkflowAttempt


# Per-attempt pickup limit. If a worker dequeues an attempt and
# fails (status='failed' outcome='worker-internal-error'), it can be
# re-pended; pickup_count++. Beyond MAX_PICKUPS, the master transitions
# the workflow to STUCK with reason='attempt-pickup-exhausted'.
DEFAULT_MAX_PICKUPS = int(os.environ.get("CONTROLLER_MAX_ATTEMPT_PICKUPS", "3"))


@dataclass
class DequeueResult:
    """What the worker gets back from dequeue.

    ``acquired`` is True iff a row was claimed. On False, no rows
    available (worker sleeps + retries) — ``reason`` distinguishes
    'empty queue' from 'all candidate rows exceeded pickup limit'.
    """

    acquired: bool
    attempt_id: int | None = None
    workflow_id: int | None = None
    role: str | None = None
    tier: int | None = None
    pickup_count: int | None = None
    reason: str | None = None


def _now() -> datetime:
    return datetime.now(timezone.utc)


def dequeue_one(
    session: Session,
    *,
    instance_id: str,
    roles: list[str],
    max_pickups: int = DEFAULT_MAX_PICKUPS,
) -> DequeueResult:
    """Atomically pick the oldest pending attempt this worker can
    handle and stamp it as in_progress with the worker's identity.

    ``instance_id`` format: ``{hostname}/{pid}/{worker_uuid}``
    (v9: slash-delimited so IPv6 hostnames don't trip parsing).
    ``roles`` is the list of role values the worker accepts.

    Returns ``DequeueResult(acquired=False)`` if no eligible row.
    Returns ``DequeueResult(acquired=True, attempt_id=..., ...)``
    on success — the row is now locked under ``instance_id`` and
    the worker holds the heartbeat responsibility.

    Dialect routing:
    - Postgres: ``SELECT … FOR UPDATE SKIP LOCKED`` (atomic dequeue
      across racing workers).
    - SQLite: ``BEGIN IMMEDIATE`` was already applied by the caller's
      session_scope; SQLite's locking serializes the SELECT + UPDATE.

    The session is committed by the caller (``session_scope``).
    """
    if not roles:
        return DequeueResult(acquired=False, reason="no_roles_configured")

    dialect = session.bind.dialect.name if session.bind else "sqlite"
    now = _now()

    if dialect == "postgresql":
        return _dequeue_postgres(session, instance_id, roles, max_pickups, now)
    # SQLite (or any other; fall through with a generic implementation)
    return _dequeue_sqlite(session, instance_id, roles, max_pickups, now)


def _dequeue_postgres(
    session: Session,
    instance_id: str,
    roles: list[str],
    max_pickups: int,
    now: datetime,
) -> DequeueResult:
    """Postgres path: SELECT FOR UPDATE SKIP LOCKED + UPDATE."""
    # Two-statement transactional dequeue. We use raw SQL because
    # SQLAlchemy ORM's auto-flush behavior would re-query.
    select_sql = text("""
        SELECT attempt_id, workflow_id, role, tier, pickup_count
          FROM workflow_attempts
         WHERE status = 'pending'
           AND role = ANY(:roles)
           AND pickup_count < :max_pickups
         ORDER BY created_at ASC
         LIMIT 1
         FOR UPDATE SKIP LOCKED
    """).bindparams(bindparam("roles"), bindparam("max_pickups"))
    row = session.execute(
        select_sql, {"roles": roles, "max_pickups": max_pickups}
    ).first()
    if row is None:
        return DequeueResult(acquired=False, reason="no_pending_eligible")

    update_sql = text("""
        UPDATE workflow_attempts
           SET status = 'in_progress',
               locked_by_instance = :instance,
               locked_at = :now,
               lock_heartbeat_at = :now,
               pickup_count = pickup_count + 1,
               started_at = :now
         WHERE attempt_id = :attempt_id
    """)
    session.execute(
        update_sql,
        {"instance": instance_id, "now": now, "attempt_id": row.attempt_id},
    )
    return DequeueResult(
        acquired=True,
        attempt_id=row.attempt_id,
        workflow_id=row.workflow_id,
        role=row.role,
        tier=row.tier,
        pickup_count=(row.pickup_count or 0) + 1,
    )


def _dequeue_sqlite(
    session: Session,
    instance_id: str,
    roles: list[str],
    max_pickups: int,
    now: datetime,
) -> DequeueResult:
    """SQLite path: serialise on BEGIN IMMEDIATE + UPDATE.

    SQLite 3.35+ supports RETURNING, which makes this atomic in one
    statement. Older SQLite would need a SELECT-then-UPDATE, accepting
    a tight TOCTOU window — fine for tests but the test DB pins
    sqlite ≥3.35 in practice (Python 3.13 ships 3.43+).
    """
    # SQLite doesn't have ANY(:roles); emit an IN-clause via parameter
    # expansion.
    role_clauses = ",".join(f":role_{i}" for i in range(len(roles)))
    update_sql = text(f"""
        UPDATE workflow_attempts
           SET status = 'in_progress',
               locked_by_instance = :instance,
               locked_at = :now,
               lock_heartbeat_at = :now,
               pickup_count = pickup_count + 1,
               started_at = :now
         WHERE attempt_id = (
             SELECT attempt_id FROM workflow_attempts
              WHERE status = 'pending'
                AND role IN ({role_clauses})
                AND pickup_count < :max_pickups
              ORDER BY created_at ASC
              LIMIT 1
         )
         RETURNING attempt_id, workflow_id, role, tier, pickup_count
    """)
    params = {
        "instance": instance_id,
        "now": now,
        "max_pickups": max_pickups,
        **{f"role_{i}": r for i, r in enumerate(roles)},
    }
    row = session.execute(update_sql, params).first()
    if row is None:
        return DequeueResult(acquired=False, reason="no_pending_eligible")
    return DequeueResult(
        acquired=True,
        attempt_id=row.attempt_id,
        workflow_id=row.workflow_id,
        role=row.role,
        tier=row.tier,
        pickup_count=row.pickup_count,
    )