"""Repositories over `city_briefs` and `city_briefs_runs`. The repository layer is the only code that owns SQL for these two tables. The orchestrator and API layer call repo methods and never write raw SQL themselves — this keeps the surface narrow and audit-friendly. """ from __future__ import annotations import logging from datetime import date, datetime, timezone from typing import Any from uuid import UUID from app.db import db from app.models import STAGE_COLUMNS logger = logging.getLogger(__name__) def _ts() -> str: """Current timestamp as an ISO string usable in SQL.""" return datetime.now(timezone.utc).isoformat() class BriefRepo: """Reads + writes for `city_briefs`.""" async def fetch_by_show(self, show_id: str) -> dict | None: return await db.fetchone( "SELECT * FROM city_briefs WHERE show_id = %s", (show_id,), ) async def fetch_existing_published(self, show_id: str) -> dict | None: return await db.fetchone( "SELECT * FROM city_briefs WHERE show_id = %s AND status = 'published'", (show_id,), ) async def insert(self, payload: dict[str, Any]) -> dict: cols = ", ".join(payload.keys()) placeholders = ", ".join(f"%s" for _ in payload) rows = await db.fetch( f"INSERT INTO city_briefs ({cols}) VALUES ({placeholders}) RETURNING *", tuple(payload.values()), ) return rows[0] async def ensure_pending(self, *, show_id: str, show_date: date, city: str, country: str, venue: str | None = None, capacity: int | None = None, comedian_id: str = "max-amini") -> dict: """Create a pending brief row if none exists; return the row either way. `UNIQUE(show_id)` guarantees we cannot race a sibling into two rows. """ existing = await self.fetch_by_show(show_id) if existing: return existing payload = { "show_id": show_id, "show_date": show_date, "city": city, "country": country, "venue": venue, "capacity": capacity, "comedian_id": comedian_id, "status": "pending", } return await self.insert(payload) async def update(self, brief_id: UUID | str, payload: dict[str, Any]) -> dict: cols = ", ".join(f"{k} = %s" for k in payload.keys()) params: tuple[Any, ...] = (*payload.values(), brief_id) rows = await db.fetch( f"UPDATE city_briefs SET {cols} WHERE brief_id = %s RETURNING *", params, ) return rows[0] async def set_generating(self, brief_id: UUID | str) -> dict: return await self.update(brief_id, {"status": "generating", "generated_at": _ts()}) async def set_generated(self, brief_id: UUID | str, generated_topics: list, ranked_topics: list | None = None, verified_facts: list | None = None, removed_facts: list | None = None, city_intelligence: dict | None = None, recent_news: list | None = None, competitor_warnings: list | None = None, crowd_work: list | None = None, recent_transcripts: list | None = None, previous_show_date: date | None = None, previous_show_venue: str | None = None, previous_tickets_sold: int | None = None) -> dict: return await self.update(brief_id, { "status": "generated", "generated_at": _ts(), "generated_topics": generated_topics, "ranked_topics": ranked_topics, "verified_facts": verified_facts, "removed_facts": removed_facts, "city_intelligence": city_intelligence, "recent_news": recent_news, "competitor_warnings": competitor_warnings, "crowd_work": crowd_work, "recent_transcripts": recent_transcripts, "previous_show_date": previous_show_date, "previous_show_venue": previous_show_venue, "previous_tickets_sold": previous_tickets_sold, }) async def set_deploying(self, brief_id: UUID | str, public_url: str) -> dict: return await self.update(brief_id, {"status": "deploying", "public_url": public_url}) async def mark_published(self, brief_id: UUID | str, generation_duration_ms: int) -> dict: return await self.update(brief_id, { "status": "published", "published_at": _ts(), "generation_duration": generation_duration_ms, "error": None, }) async def mark_failed(self, brief_id: UUID | str, error: str | None) -> dict: return await self.update(brief_id, { "status": "failed", "error": error[:4000] if error else None, }) async def recent_briefs(self, limit: int = 20) -> list[dict]: """For the API's status dashboard / manual inspection.""" return await db.fetch( "SELECT brief_id, show_id, show_date, city, country, status, " "public_url, error, published_at, created_at " "FROM city_briefs ORDER BY created_at DESC LIMIT %s", (limit,), ) class RunRepo: """Reads + writes for `city_briefs_runs` — per-stage observability.""" async def insert(self, *, show_id: str | None = None, show_date: date | None = None) -> dict: if show_id is None: # A "run" without a show means the run-level row (e.g. discovery run). params: tuple[Any, ...] = () else: params = (show_id, show_date) if show_id is None: rows = await db.fetch( "INSERT INTO city_briefs_runs (show_id, show_date) " "VALUES (NULL, NULL) RETURNING *" ) else: rows = await db.fetch( "INSERT INTO city_briefs_runs (show_id, show_date) " "VALUES (%s, %s) RETURNING *", (show_id, show_date), ) return rows[0] async def fetch(self, brief_run_id: UUID | str) -> dict | None: return await db.fetchone( "SELECT * FROM city_briefs_runs WHERE brief_run_id = %s", (brief_run_id,), ) async def set_stage(self, brief_run_id: UUID | str, stage: str, state: str, error: str | None = None, public_url: str | None = None) -> dict | None: """Set a stage's state (+ optional error / public_url scratch). `stage` must be one of STAGE_COLUMNS. `state` must be one of 'pending','running','ok','failed','skipped'. """ if stage not in STAGE_COLUMNS: raise ValueError(f"unknown stage {stage!r}") if state not in ("pending", "running", "ok", "failed", "skipped"): raise ValueError(f"unknown stage state {state!r}") sets: list[str] = [f"{stage} = %s"] params: list[Any] = [state] if error is not None: err_col = f"{stage}_error" sets.append(f"{err_col} = %s") params.append(error[:4000] if error else None) if public_url is not None and stage == "vps_deployment": # The orchestrator stashes the deployed URL here for the verify stages. sets.append("url_verification = %s") params.append("pending") params.append(brief_run_id) rows = await db.fetch( f"UPDATE city_briefs_runs SET {', '.join(sets)} " f"WHERE brief_run_id = %s RETURNING *", tuple(params), ) return rows[0] if rows else None async def mark_complete(self, brief_run_id: UUID | str, status: str, duration_ms: int, error: str | None = None, retry_count: int | None = None) -> dict: payload: dict[str, Any] = { "completed_at": _ts(), "status": status, "duration": duration_ms, } if error is not None: payload["error"] = error[:4000] if error else None if retry_count is not None: payload["retry_count"] = retry_count cols = ", ".join(f"{k} = %s" for k in payload.keys()) params = (*payload.values(), brief_run_id) rows = await db.fetch( f"UPDATE city_briefs_runs SET {cols} WHERE brief_run_id = %s RETURNING *", params, ) return rows[0] async def increment_retry(self, brief_run_id: UUID | str) -> None: await db.execute( "UPDATE city_briefs_runs SET retry_count = retry_count + 1 " "WHERE brief_run_id = %s", (brief_run_id,), ) # Singletons consumed by orchestrator + API layer. brief_repo = BriefRepo() run_repo = RunRepo()