feat: control-plane DB substrate for atomic assign/lease (Closes #613)
Add SQLite single-writer MVP control plane per the allocator ADR: sessions, work_items (issue/pr only), atomic assign+lease, mutation gates, terminal-lock index, and provider-neutral incident_links. DB coordinates concurrency; Gitea remains assignable work; raw Sentry/GlitchTip incidents are never work items. #600 and #612 build on this substrate.
This commit is contained in:
@@ -0,0 +1,944 @@
|
||||
"""Control-plane DB substrate for multi-session coordination (#613).
|
||||
|
||||
Implements the durable coordination store described by
|
||||
``docs/architecture/mcp-allocator-control-plane-observability-adr.md``:
|
||||
|
||||
* DB coordinates live concurrency (sessions, atomic assignment+lease,
|
||||
heartbeats, terminal-lock index, events, ``incident_links``).
|
||||
* Gitea remains the durable work record and the only assignable work unit
|
||||
(``issue`` / ``pr`` — never raw Sentry/GlitchTip incidents).
|
||||
* SQLite is the single-writer MVP backend; the API is backend-shaped so a
|
||||
shared Postgres or single allocator daemon can replace the connection
|
||||
layer later without changing call sites.
|
||||
|
||||
This module is the **substrate** for #600 (allocator policy/tool) and #612
|
||||
(incident bridge). It does **not** implement ``gitea_allocate_next_work``
|
||||
routing policy or provider adapters.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Iterator, Sequence
|
||||
|
||||
SCHEMA_VERSION = 1
|
||||
|
||||
# Assignable work kinds only — raw monitoring incidents are never work items.
|
||||
WORK_KINDS = frozenset({"issue", "pr"})
|
||||
|
||||
DEFAULT_LEASE_TTL_SECONDS = 4 * 3600
|
||||
|
||||
# Environment: path for SQLite MVP (single-writer).
|
||||
DB_PATH_ENV = "GITEA_CONTROL_PLANE_DB"
|
||||
DEFAULT_DB_PATH = os.path.expanduser("~/.cache/gitea-tools/control-plane/control_plane.sqlite3")
|
||||
|
||||
_SCHEMA_SQL = """
|
||||
PRAGMA foreign_keys = ON;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS schema_meta (
|
||||
key TEXT PRIMARY KEY,
|
||||
value TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS sessions (
|
||||
session_id TEXT PRIMARY KEY,
|
||||
role TEXT NOT NULL,
|
||||
profile TEXT,
|
||||
namespace TEXT,
|
||||
pid INTEGER,
|
||||
started_at TEXT NOT NULL,
|
||||
last_heartbeat_at TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'active'
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS work_items (
|
||||
work_item_id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
remote TEXT NOT NULL,
|
||||
org TEXT NOT NULL,
|
||||
repo TEXT NOT NULL,
|
||||
kind TEXT NOT NULL CHECK (kind IN ('issue', 'pr')),
|
||||
number INTEGER NOT NULL,
|
||||
state TEXT NOT NULL DEFAULT 'open',
|
||||
priority INTEGER NOT NULL DEFAULT 0,
|
||||
current_head_sha TEXT,
|
||||
updated_at TEXT NOT NULL,
|
||||
UNIQUE (remote, org, repo, kind, number)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS leases (
|
||||
lease_id TEXT PRIMARY KEY,
|
||||
work_item_id INTEGER NOT NULL REFERENCES work_items(work_item_id),
|
||||
session_id TEXT NOT NULL REFERENCES sessions(session_id),
|
||||
role TEXT NOT NULL,
|
||||
phase TEXT NOT NULL DEFAULT 'claimed',
|
||||
expires_at TEXT NOT NULL,
|
||||
heartbeat_at TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'active'
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS assignments (
|
||||
assignment_id TEXT PRIMARY KEY,
|
||||
work_item_id INTEGER NOT NULL REFERENCES work_items(work_item_id),
|
||||
session_id TEXT NOT NULL REFERENCES sessions(session_id),
|
||||
lease_id TEXT NOT NULL REFERENCES leases(lease_id),
|
||||
allowed_actions TEXT NOT NULL,
|
||||
forbidden_actions TEXT NOT NULL,
|
||||
expected_head_sha TEXT,
|
||||
role TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS terminal_locks (
|
||||
terminal_lock_id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
remote TEXT NOT NULL,
|
||||
org TEXT NOT NULL,
|
||||
repo TEXT NOT NULL,
|
||||
terminal_pr INTEGER NOT NULL,
|
||||
review_id TEXT,
|
||||
decision TEXT,
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
cleanup_state TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
UNIQUE (remote, org, repo, terminal_pr)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS events (
|
||||
event_id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
work_item_id INTEGER REFERENCES work_items(work_item_id),
|
||||
event_type TEXT NOT NULL,
|
||||
message TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
-- Provider-neutral incident ↔ Gitea link model (ADR §8–9). Not assignable work.
|
||||
CREATE TABLE IF NOT EXISTS incident_links (
|
||||
link_id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
provider TEXT NOT NULL,
|
||||
provider_base_url TEXT,
|
||||
provider_org TEXT,
|
||||
provider_project TEXT,
|
||||
provider_issue_id TEXT NOT NULL,
|
||||
provider_short_id TEXT,
|
||||
provider_permalink TEXT,
|
||||
fingerprint TEXT,
|
||||
gitea_org TEXT NOT NULL,
|
||||
gitea_repo TEXT NOT NULL,
|
||||
gitea_issue_number INTEGER NOT NULL,
|
||||
linked_pr_numbers TEXT,
|
||||
first_seen TEXT,
|
||||
last_seen TEXT,
|
||||
event_count INTEGER,
|
||||
status TEXT NOT NULL DEFAULT 'open',
|
||||
release_resolved_at TEXT,
|
||||
last_sync_at TEXT,
|
||||
UNIQUE (provider, provider_base_url, provider_org, provider_project, provider_issue_id)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_leases_work_status ON leases(work_item_id, status);
|
||||
CREATE INDEX IF NOT EXISTS idx_assignments_session ON assignments(session_id, status);
|
||||
CREATE INDEX IF NOT EXISTS idx_incident_gitea ON incident_links(gitea_org, gitea_repo, gitea_issue_number);
|
||||
"""
|
||||
|
||||
|
||||
def _utc_now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def _ts(dt: datetime | None = None) -> str:
|
||||
value = dt or _utc_now()
|
||||
if value.tzinfo is None:
|
||||
value = value.replace(tzinfo=timezone.utc)
|
||||
return value.astimezone(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
|
||||
|
||||
|
||||
def _parse_ts(value: str | None) -> datetime | None:
|
||||
if not value:
|
||||
return None
|
||||
text = value.strip()
|
||||
if text.endswith("Z"):
|
||||
text = text[:-1] + "+00:00"
|
||||
try:
|
||||
return datetime.fromisoformat(text)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def default_db_path() -> str:
|
||||
raw = (os.environ.get(DB_PATH_ENV) or DEFAULT_DB_PATH).strip()
|
||||
return raw or DEFAULT_DB_PATH
|
||||
|
||||
|
||||
class ControlPlaneError(RuntimeError):
|
||||
"""Base error for control-plane substrate failures."""
|
||||
|
||||
|
||||
class InvalidWorkKindError(ControlPlaneError):
|
||||
"""Raised when a non-assignable work kind is requested."""
|
||||
|
||||
|
||||
class LeaseRequiredError(ControlPlaneError):
|
||||
"""Raised when a mutation is attempted without a valid assignment/lease."""
|
||||
|
||||
|
||||
class ForeignLeaseError(ControlPlaneError):
|
||||
"""Raised when another session holds the active lease."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AssignmentResult:
|
||||
"""Result of an atomic assign+lease transaction."""
|
||||
|
||||
outcome: str # assigned | wait | no_safe_work
|
||||
assignment_id: str | None = None
|
||||
lease_id: str | None = None
|
||||
session_id: str | None = None
|
||||
role: str | None = None
|
||||
work_kind: str | None = None
|
||||
work_number: int | None = None
|
||||
remote: str | None = None
|
||||
org: str | None = None
|
||||
repo: str | None = None
|
||||
expected_head_sha: str | None = None
|
||||
allowed_actions: tuple[str, ...] = ()
|
||||
forbidden_actions: tuple[str, ...] = ()
|
||||
expires_at: str | None = None
|
||||
owner_session_id: str | None = None
|
||||
reason: str = ""
|
||||
|
||||
def as_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"outcome": self.outcome,
|
||||
"assignment_id": self.assignment_id,
|
||||
"lease_id": self.lease_id,
|
||||
"session_id": self.session_id,
|
||||
"role": self.role,
|
||||
"work_kind": self.work_kind,
|
||||
"work_number": self.work_number,
|
||||
"remote": self.remote,
|
||||
"org": self.org,
|
||||
"repo": self.repo,
|
||||
"expected_head_sha": self.expected_head_sha,
|
||||
"allowed_actions": list(self.allowed_actions),
|
||||
"forbidden_actions": list(self.forbidden_actions),
|
||||
"expires_at": self.expires_at,
|
||||
"owner_session_id": self.owner_session_id,
|
||||
"reason": self.reason,
|
||||
}
|
||||
|
||||
|
||||
class ControlPlaneDB:
|
||||
"""SQLite single-writer MVP control-plane store.
|
||||
|
||||
Use one process as writer for multi-session safety claims, or migrate to
|
||||
Postgres / a single allocator daemon for multi-host production (ADR §6).
|
||||
"""
|
||||
|
||||
def __init__(self, db_path: str | None = None) -> None:
|
||||
self.db_path = (db_path or default_db_path()).strip()
|
||||
parent = os.path.dirname(self.db_path)
|
||||
if parent:
|
||||
os.makedirs(parent, mode=0o700, exist_ok=True)
|
||||
self._lock = threading.RLock()
|
||||
self._init_schema()
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
# Default isolation (DEFERRED) so explicit BEGIN IMMEDIATE works.
|
||||
conn = sqlite3.connect(self.db_path, timeout=30)
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
# Serialize writers even across threads in one process.
|
||||
conn.execute("PRAGMA journal_mode = WAL")
|
||||
return conn
|
||||
|
||||
@contextmanager
|
||||
def _tx(self, immediate: bool = True) -> Iterator[sqlite3.Connection]:
|
||||
with self._lock:
|
||||
conn = self._connect()
|
||||
try:
|
||||
if immediate:
|
||||
conn.execute("BEGIN IMMEDIATE")
|
||||
else:
|
||||
conn.execute("BEGIN")
|
||||
yield conn
|
||||
conn.commit()
|
||||
except Exception:
|
||||
try:
|
||||
conn.rollback()
|
||||
except sqlite3.Error:
|
||||
pass
|
||||
raise
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _init_schema(self) -> None:
|
||||
# executescript auto-commits; run schema outside an open txn, then meta.
|
||||
with self._lock:
|
||||
conn = self._connect()
|
||||
try:
|
||||
conn.executescript(_SCHEMA_SQL)
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO schema_meta(key, value) VALUES (?, ?)",
|
||||
("schema_version", str(SCHEMA_VERSION)),
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO schema_meta(key, value) VALUES (?, ?)",
|
||||
(
|
||||
"architecture",
|
||||
"DB coordinates; Gitea records; Sentry/GlitchTip observe; "
|
||||
"bridge is only path from observations to Gitea work",
|
||||
),
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
# ── sessions ──────────────────────────────────────────────────────────
|
||||
|
||||
def upsert_session(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
role: str,
|
||||
profile: str | None = None,
|
||||
namespace: str | None = None,
|
||||
pid: int | None = None,
|
||||
status: str = "active",
|
||||
) -> dict[str, Any]:
|
||||
now = _ts()
|
||||
with self._tx() as conn:
|
||||
existing = conn.execute(
|
||||
"SELECT session_id FROM sessions WHERE session_id = ?",
|
||||
(session_id,),
|
||||
).fetchone()
|
||||
if existing:
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE sessions
|
||||
SET role = ?, profile = ?, namespace = ?, pid = ?,
|
||||
last_heartbeat_at = ?, status = ?
|
||||
WHERE session_id = ?
|
||||
""",
|
||||
(role, profile, namespace, pid, now, status, session_id),
|
||||
)
|
||||
else:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO sessions(
|
||||
session_id, role, profile, namespace, pid,
|
||||
started_at, last_heartbeat_at, status
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(session_id, role, profile, namespace, pid, now, now, status),
|
||||
)
|
||||
row = conn.execute(
|
||||
"SELECT * FROM sessions WHERE session_id = ?",
|
||||
(session_id,),
|
||||
).fetchone()
|
||||
return dict(row)
|
||||
|
||||
def heartbeat_session(self, session_id: str) -> None:
|
||||
with self._tx() as conn:
|
||||
conn.execute(
|
||||
"UPDATE sessions SET last_heartbeat_at = ? WHERE session_id = ?",
|
||||
(_ts(), session_id),
|
||||
)
|
||||
|
||||
# ── work items ────────────────────────────────────────────────────────
|
||||
|
||||
def upsert_work_item(
|
||||
self,
|
||||
*,
|
||||
remote: str,
|
||||
org: str,
|
||||
repo: str,
|
||||
kind: str,
|
||||
number: int,
|
||||
state: str = "open",
|
||||
priority: int = 0,
|
||||
current_head_sha: str | None = None,
|
||||
) -> int:
|
||||
kind_norm = (kind or "").strip().lower()
|
||||
if kind_norm not in WORK_KINDS:
|
||||
raise InvalidWorkKindError(
|
||||
f"work kind '{kind}' is not assignable; only {sorted(WORK_KINDS)} "
|
||||
f"are allowed (raw monitoring incidents are never work items)"
|
||||
)
|
||||
now = _ts()
|
||||
with self._tx() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO work_items(
|
||||
remote, org, repo, kind, number, state, priority,
|
||||
current_head_sha, updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(remote, org, repo, kind, number) DO UPDATE SET
|
||||
state = excluded.state,
|
||||
priority = excluded.priority,
|
||||
current_head_sha = excluded.current_head_sha,
|
||||
updated_at = excluded.updated_at
|
||||
""",
|
||||
(
|
||||
remote,
|
||||
org,
|
||||
repo,
|
||||
kind_norm,
|
||||
int(number),
|
||||
state,
|
||||
int(priority),
|
||||
current_head_sha,
|
||||
now,
|
||||
),
|
||||
)
|
||||
row = conn.execute(
|
||||
"""
|
||||
SELECT work_item_id FROM work_items
|
||||
WHERE remote = ? AND org = ? AND repo = ? AND kind = ? AND number = ?
|
||||
""",
|
||||
(remote, org, repo, kind_norm, int(number)),
|
||||
).fetchone()
|
||||
return int(row["work_item_id"])
|
||||
|
||||
# ── atomic assign + lease ─────────────────────────────────────────────
|
||||
|
||||
def assign_and_lease(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
role: str,
|
||||
remote: str,
|
||||
org: str,
|
||||
repo: str,
|
||||
kind: str,
|
||||
number: int,
|
||||
expected_head_sha: str | None = None,
|
||||
allowed_actions: Sequence[str] | None = None,
|
||||
forbidden_actions: Sequence[str] | None = None,
|
||||
lease_ttl_seconds: int = DEFAULT_LEASE_TTL_SECONDS,
|
||||
phase: str = "claimed",
|
||||
now: datetime | None = None,
|
||||
) -> AssignmentResult:
|
||||
"""Atomically create assignment + lease for one Gitea work item.
|
||||
|
||||
Concurrent sessions cannot both receive the same open work item.
|
||||
Returns ``wait`` with owner_session_id when a live foreign lease exists.
|
||||
"""
|
||||
kind_norm = (kind or "").strip().lower()
|
||||
if kind_norm not in WORK_KINDS:
|
||||
raise InvalidWorkKindError(
|
||||
f"cannot assign kind '{kind}'; only {sorted(WORK_KINDS)}"
|
||||
)
|
||||
|
||||
allowed = tuple(allowed_actions or ("implement", "comment"))
|
||||
forbidden = tuple(
|
||||
forbidden_actions
|
||||
or ("approve", "merge", "self_select_without_assignment")
|
||||
)
|
||||
moment = now or _utc_now()
|
||||
now_s = _ts(moment)
|
||||
expires = _ts(moment + timedelta(seconds=int(lease_ttl_seconds)))
|
||||
|
||||
with self._tx(immediate=True) as conn:
|
||||
# Ensure session exists
|
||||
sess = conn.execute(
|
||||
"SELECT session_id FROM sessions WHERE session_id = ?",
|
||||
(session_id,),
|
||||
).fetchone()
|
||||
if not sess:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO sessions(
|
||||
session_id, role, profile, namespace, pid,
|
||||
started_at, last_heartbeat_at, status
|
||||
) VALUES (?, ?, NULL, NULL, ?, ?, ?, 'active')
|
||||
""",
|
||||
(session_id, role, os.getpid(), now_s, now_s),
|
||||
)
|
||||
|
||||
# Upsert work item
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO work_items(
|
||||
remote, org, repo, kind, number, state, priority,
|
||||
current_head_sha, updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, 'open', 0, ?, ?)
|
||||
ON CONFLICT(remote, org, repo, kind, number) DO UPDATE SET
|
||||
current_head_sha = COALESCE(excluded.current_head_sha, work_items.current_head_sha),
|
||||
updated_at = excluded.updated_at
|
||||
""",
|
||||
(remote, org, repo, kind_norm, int(number), expected_head_sha, now_s),
|
||||
)
|
||||
work = conn.execute(
|
||||
"""
|
||||
SELECT * FROM work_items
|
||||
WHERE remote = ? AND org = ? AND repo = ? AND kind = ? AND number = ?
|
||||
""",
|
||||
(remote, org, repo, kind_norm, int(number)),
|
||||
).fetchone()
|
||||
work_item_id = int(work["work_item_id"])
|
||||
state = (work["state"] or "open").lower()
|
||||
if state in ("merged", "closed"):
|
||||
return AssignmentResult(
|
||||
outcome="no_safe_work",
|
||||
reason=f"work item {kind_norm}#{number} is {state}; never assign",
|
||||
)
|
||||
|
||||
# Expire stale leases on this work item first
|
||||
self._expire_stale_leases_conn(conn, work_item_id=work_item_id, now_s=now_s)
|
||||
|
||||
active = conn.execute(
|
||||
"""
|
||||
SELECT * FROM leases
|
||||
WHERE work_item_id = ? AND status = 'active'
|
||||
ORDER BY expires_at DESC
|
||||
LIMIT 1
|
||||
""",
|
||||
(work_item_id,),
|
||||
).fetchone()
|
||||
if active:
|
||||
owner = active["session_id"]
|
||||
if owner == session_id:
|
||||
# Owner-resume: refresh heartbeat and return existing assignment
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE leases
|
||||
SET heartbeat_at = ?, expires_at = ?, phase = ?
|
||||
WHERE lease_id = ?
|
||||
""",
|
||||
(now_s, expires, phase, active["lease_id"]),
|
||||
)
|
||||
asn = conn.execute(
|
||||
"""
|
||||
SELECT * FROM assignments
|
||||
WHERE lease_id = ? AND status = 'active'
|
||||
ORDER BY created_at DESC LIMIT 1
|
||||
""",
|
||||
(active["lease_id"],),
|
||||
).fetchone()
|
||||
if asn:
|
||||
return AssignmentResult(
|
||||
outcome="assigned",
|
||||
assignment_id=asn["assignment_id"],
|
||||
lease_id=active["lease_id"],
|
||||
session_id=session_id,
|
||||
role=asn["role"],
|
||||
work_kind=kind_norm,
|
||||
work_number=int(number),
|
||||
remote=remote,
|
||||
org=org,
|
||||
repo=repo,
|
||||
expected_head_sha=asn["expected_head_sha"],
|
||||
allowed_actions=tuple(json.loads(asn["allowed_actions"])),
|
||||
forbidden_actions=tuple(json.loads(asn["forbidden_actions"])),
|
||||
expires_at=expires,
|
||||
owner_session_id=session_id,
|
||||
reason="owner-resume: refreshed existing lease",
|
||||
)
|
||||
return AssignmentResult(
|
||||
outcome="wait",
|
||||
owner_session_id=owner,
|
||||
reason=(
|
||||
f"foreign active lease held by session {owner} on "
|
||||
f"{kind_norm}#{number}"
|
||||
),
|
||||
)
|
||||
|
||||
lease_id = f"lease-{uuid.uuid4().hex[:16]}"
|
||||
assignment_id = f"asn-{uuid.uuid4().hex[:16]}"
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO leases(
|
||||
lease_id, work_item_id, session_id, role, phase,
|
||||
expires_at, heartbeat_at, status
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, 'active')
|
||||
""",
|
||||
(lease_id, work_item_id, session_id, role, phase, expires, now_s),
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO assignments(
|
||||
assignment_id, work_item_id, session_id, lease_id,
|
||||
allowed_actions, forbidden_actions, expected_head_sha,
|
||||
role, status, created_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'active', ?)
|
||||
""",
|
||||
(
|
||||
assignment_id,
|
||||
work_item_id,
|
||||
session_id,
|
||||
lease_id,
|
||||
json.dumps(list(allowed)),
|
||||
json.dumps(list(forbidden)),
|
||||
expected_head_sha or work["current_head_sha"],
|
||||
role,
|
||||
now_s,
|
||||
),
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO events(work_item_id, event_type, message, created_at)
|
||||
VALUES (?, 'assigned', ?, ?)
|
||||
""",
|
||||
(
|
||||
work_item_id,
|
||||
f"session {session_id} assigned {kind_norm}#{number} lease={lease_id}",
|
||||
now_s,
|
||||
),
|
||||
)
|
||||
return AssignmentResult(
|
||||
outcome="assigned",
|
||||
assignment_id=assignment_id,
|
||||
lease_id=lease_id,
|
||||
session_id=session_id,
|
||||
role=role,
|
||||
work_kind=kind_norm,
|
||||
work_number=int(number),
|
||||
remote=remote,
|
||||
org=org,
|
||||
repo=repo,
|
||||
expected_head_sha=expected_head_sha or work["current_head_sha"],
|
||||
allowed_actions=allowed,
|
||||
forbidden_actions=forbidden,
|
||||
expires_at=expires,
|
||||
owner_session_id=session_id,
|
||||
reason="atomic assign+lease created",
|
||||
)
|
||||
|
||||
def _expire_stale_leases_conn(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
*,
|
||||
work_item_id: int | None = None,
|
||||
now_s: str | None = None,
|
||||
) -> int:
|
||||
now_s = now_s or _ts()
|
||||
if work_item_id is None:
|
||||
rows = conn.execute(
|
||||
"SELECT lease_id FROM leases WHERE status = 'active' AND expires_at <= ?",
|
||||
(now_s,),
|
||||
).fetchall()
|
||||
else:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT lease_id FROM leases
|
||||
WHERE status = 'active' AND expires_at <= ? AND work_item_id = ?
|
||||
""",
|
||||
(now_s, work_item_id),
|
||||
).fetchall()
|
||||
for row in rows:
|
||||
lid = row["lease_id"]
|
||||
conn.execute(
|
||||
"UPDATE leases SET status = 'expired' WHERE lease_id = ?",
|
||||
(lid,),
|
||||
)
|
||||
conn.execute(
|
||||
"UPDATE assignments SET status = 'expired' WHERE lease_id = ?",
|
||||
(lid,),
|
||||
)
|
||||
return len(rows)
|
||||
|
||||
def expire_stale_leases(self) -> int:
|
||||
with self._tx() as conn:
|
||||
return self._expire_stale_leases_conn(conn)
|
||||
|
||||
def heartbeat_lease(self, lease_id: str, *, session_id: str) -> dict[str, Any]:
|
||||
now_s = _ts()
|
||||
with self._tx() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT * FROM leases WHERE lease_id = ?",
|
||||
(lease_id,),
|
||||
).fetchone()
|
||||
if not row:
|
||||
raise ControlPlaneError(f"unknown lease_id {lease_id}")
|
||||
if row["session_id"] != session_id:
|
||||
raise ForeignLeaseError(
|
||||
f"lease {lease_id} owned by {row['session_id']}, not {session_id}"
|
||||
)
|
||||
if row["status"] != "active":
|
||||
raise ControlPlaneError(f"lease {lease_id} status is {row['status']}")
|
||||
exp = _parse_ts(row["expires_at"])
|
||||
if exp and exp <= _utc_now():
|
||||
conn.execute(
|
||||
"UPDATE leases SET status = 'expired' WHERE lease_id = ?",
|
||||
(lease_id,),
|
||||
)
|
||||
raise ControlPlaneError(f"lease {lease_id} already expired")
|
||||
# Extend TTL on heartbeat
|
||||
new_exp = _ts(_utc_now() + timedelta(seconds=DEFAULT_LEASE_TTL_SECONDS))
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE leases SET heartbeat_at = ?, expires_at = ?
|
||||
WHERE lease_id = ?
|
||||
""",
|
||||
(now_s, new_exp, lease_id),
|
||||
)
|
||||
conn.execute(
|
||||
"UPDATE sessions SET last_heartbeat_at = ? WHERE session_id = ?",
|
||||
(now_s, session_id),
|
||||
)
|
||||
return {
|
||||
"lease_id": lease_id,
|
||||
"heartbeat_at": now_s,
|
||||
"expires_at": new_exp,
|
||||
"session_id": session_id,
|
||||
}
|
||||
|
||||
def release_lease(self, lease_id: str, *, session_id: str) -> None:
|
||||
with self._tx() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT * FROM leases WHERE lease_id = ?",
|
||||
(lease_id,),
|
||||
).fetchone()
|
||||
if not row:
|
||||
return
|
||||
if row["session_id"] != session_id:
|
||||
raise ForeignLeaseError(
|
||||
f"cannot release lease {lease_id} owned by {row['session_id']}"
|
||||
)
|
||||
conn.execute(
|
||||
"UPDATE leases SET status = 'released' WHERE lease_id = ?",
|
||||
(lease_id,),
|
||||
)
|
||||
conn.execute(
|
||||
"UPDATE assignments SET status = 'released' WHERE lease_id = ?",
|
||||
(lease_id,),
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO events(work_item_id, event_type, message, created_at)
|
||||
VALUES (?, 'lease_released', ?, ?)
|
||||
""",
|
||||
(row["work_item_id"], f"session {session_id} released {lease_id}", _ts()),
|
||||
)
|
||||
|
||||
def require_valid_assignment(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
remote: str,
|
||||
org: str,
|
||||
repo: str,
|
||||
kind: str,
|
||||
number: int,
|
||||
action: str,
|
||||
) -> dict[str, Any]:
|
||||
"""Gate a mutation: require active assignment+lease for this session/work."""
|
||||
kind_norm = (kind or "").strip().lower()
|
||||
if kind_norm not in WORK_KINDS:
|
||||
raise InvalidWorkKindError(f"invalid kind '{kind}'")
|
||||
with self._tx(immediate=False) as conn:
|
||||
work = conn.execute(
|
||||
"""
|
||||
SELECT * FROM work_items
|
||||
WHERE remote = ? AND org = ? AND repo = ? AND kind = ? AND number = ?
|
||||
""",
|
||||
(remote, org, repo, kind_norm, int(number)),
|
||||
).fetchone()
|
||||
if not work:
|
||||
raise LeaseRequiredError(
|
||||
f"no work_item for {kind_norm}#{number}; assign first"
|
||||
)
|
||||
self._expire_stale_leases_conn(conn, work_item_id=int(work["work_item_id"]))
|
||||
asn = conn.execute(
|
||||
"""
|
||||
SELECT a.*, l.status AS lease_status, l.expires_at, l.lease_id
|
||||
FROM assignments a
|
||||
JOIN leases l ON l.lease_id = a.lease_id
|
||||
WHERE a.work_item_id = ?
|
||||
AND a.session_id = ?
|
||||
AND a.status = 'active'
|
||||
AND l.status = 'active'
|
||||
ORDER BY a.created_at DESC
|
||||
LIMIT 1
|
||||
""",
|
||||
(int(work["work_item_id"]), session_id),
|
||||
).fetchone()
|
||||
if not asn:
|
||||
raise LeaseRequiredError(
|
||||
f"session {session_id} has no active assignment/lease for "
|
||||
f"{kind_norm}#{number}"
|
||||
)
|
||||
exp = _parse_ts(asn["expires_at"])
|
||||
if exp and exp <= _utc_now():
|
||||
raise LeaseRequiredError(f"lease {asn['lease_id']} expired")
|
||||
allowed = json.loads(asn["allowed_actions"])
|
||||
forbidden = json.loads(asn["forbidden_actions"])
|
||||
if action in forbidden:
|
||||
raise LeaseRequiredError(
|
||||
f"action '{action}' is forbidden on assignment {asn['assignment_id']}"
|
||||
)
|
||||
if allowed and action not in allowed and action != "heartbeat":
|
||||
raise LeaseRequiredError(
|
||||
f"action '{action}' not in allowed_actions {allowed}"
|
||||
)
|
||||
return dict(asn)
|
||||
|
||||
# ── terminal locks ────────────────────────────────────────────────────
|
||||
|
||||
def set_terminal_lock(
|
||||
self,
|
||||
*,
|
||||
remote: str,
|
||||
org: str,
|
||||
repo: str,
|
||||
terminal_pr: int,
|
||||
review_id: str | None = None,
|
||||
decision: str | None = None,
|
||||
status: str = "active",
|
||||
cleanup_state: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
now_s = _ts()
|
||||
with self._tx() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO terminal_locks(
|
||||
remote, org, repo, terminal_pr, review_id, decision,
|
||||
status, cleanup_state, created_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(remote, org, repo, terminal_pr) DO UPDATE SET
|
||||
review_id = excluded.review_id,
|
||||
decision = excluded.decision,
|
||||
status = excluded.status,
|
||||
cleanup_state = excluded.cleanup_state
|
||||
""",
|
||||
(
|
||||
remote,
|
||||
org,
|
||||
repo,
|
||||
int(terminal_pr),
|
||||
review_id,
|
||||
decision,
|
||||
status,
|
||||
cleanup_state,
|
||||
now_s,
|
||||
),
|
||||
)
|
||||
row = conn.execute(
|
||||
"""
|
||||
SELECT * FROM terminal_locks
|
||||
WHERE remote = ? AND org = ? AND repo = ? AND terminal_pr = ?
|
||||
""",
|
||||
(remote, org, repo, int(terminal_pr)),
|
||||
).fetchone()
|
||||
return dict(row)
|
||||
|
||||
def get_active_terminal_lock(
|
||||
self, *, remote: str, org: str, repo: str
|
||||
) -> dict[str, Any] | None:
|
||||
with self._tx(immediate=False) as conn:
|
||||
row = conn.execute(
|
||||
"""
|
||||
SELECT * FROM terminal_locks
|
||||
WHERE remote = ? AND org = ? AND repo = ? AND status = 'active'
|
||||
ORDER BY created_at DESC LIMIT 1
|
||||
""",
|
||||
(remote, org, repo),
|
||||
).fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
# ── incident_links (provider-neutral; not assignable) ─────────────────
|
||||
|
||||
def upsert_incident_link(
|
||||
self,
|
||||
*,
|
||||
provider: str,
|
||||
provider_issue_id: str,
|
||||
gitea_org: str,
|
||||
gitea_repo: str,
|
||||
gitea_issue_number: int,
|
||||
provider_base_url: str | None = None,
|
||||
provider_org: str | None = None,
|
||||
provider_project: str | None = None,
|
||||
provider_short_id: str | None = None,
|
||||
provider_permalink: str | None = None,
|
||||
fingerprint: str | None = None,
|
||||
linked_pr_numbers: Sequence[int] | None = None,
|
||||
first_seen: str | None = None,
|
||||
last_seen: str | None = None,
|
||||
event_count: int | None = None,
|
||||
status: str = "open",
|
||||
) -> dict[str, Any]:
|
||||
now_s = _ts()
|
||||
with self._tx() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO incident_links(
|
||||
provider, provider_base_url, provider_org, provider_project,
|
||||
provider_issue_id, provider_short_id, provider_permalink,
|
||||
fingerprint, gitea_org, gitea_repo, gitea_issue_number,
|
||||
linked_pr_numbers, first_seen, last_seen, event_count,
|
||||
status, last_sync_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(provider, provider_base_url, provider_org, provider_project, provider_issue_id)
|
||||
DO UPDATE SET
|
||||
gitea_issue_number = excluded.gitea_issue_number,
|
||||
gitea_org = excluded.gitea_org,
|
||||
gitea_repo = excluded.gitea_repo,
|
||||
provider_permalink = COALESCE(excluded.provider_permalink, incident_links.provider_permalink),
|
||||
fingerprint = COALESCE(excluded.fingerprint, incident_links.fingerprint),
|
||||
linked_pr_numbers = COALESCE(excluded.linked_pr_numbers, incident_links.linked_pr_numbers),
|
||||
last_seen = COALESCE(excluded.last_seen, incident_links.last_seen),
|
||||
event_count = COALESCE(excluded.event_count, incident_links.event_count),
|
||||
status = excluded.status,
|
||||
last_sync_at = excluded.last_sync_at
|
||||
""",
|
||||
(
|
||||
provider,
|
||||
provider_base_url,
|
||||
provider_org,
|
||||
provider_project,
|
||||
provider_issue_id,
|
||||
provider_short_id,
|
||||
provider_permalink,
|
||||
fingerprint,
|
||||
gitea_org,
|
||||
gitea_repo,
|
||||
int(gitea_issue_number),
|
||||
json.dumps(list(linked_pr_numbers or [])),
|
||||
first_seen,
|
||||
last_seen,
|
||||
event_count,
|
||||
status,
|
||||
now_s,
|
||||
),
|
||||
)
|
||||
row = conn.execute(
|
||||
"""
|
||||
SELECT * FROM incident_links
|
||||
WHERE provider = ?
|
||||
AND IFNULL(provider_base_url, '') = IFNULL(?, '')
|
||||
AND IFNULL(provider_org, '') = IFNULL(?, '')
|
||||
AND IFNULL(provider_project, '') = IFNULL(?, '')
|
||||
AND provider_issue_id = ?
|
||||
""",
|
||||
(
|
||||
provider,
|
||||
provider_base_url,
|
||||
provider_org,
|
||||
provider_project,
|
||||
provider_issue_id,
|
||||
),
|
||||
).fetchone()
|
||||
return dict(row)
|
||||
|
||||
def get_incident_link_for_gitea_issue(
|
||||
self, *, gitea_org: str, gitea_repo: str, gitea_issue_number: int
|
||||
) -> dict[str, Any] | None:
|
||||
with self._tx(immediate=False) as conn:
|
||||
row = conn.execute(
|
||||
"""
|
||||
SELECT * FROM incident_links
|
||||
WHERE gitea_org = ? AND gitea_repo = ? AND gitea_issue_number = ?
|
||||
ORDER BY link_id DESC LIMIT 1
|
||||
""",
|
||||
(gitea_org, gitea_repo, int(gitea_issue_number)),
|
||||
).fetchone()
|
||||
return dict(row) if row else None
|
||||
Reference in New Issue
Block a user