"""PostgreSQL implementation of the formal acquisition repository.""" from __future__ import annotations from typing import Any from uuid import UUID import psycopg2.extras from acquisition.domain import ( AcquisitionJob, AcquisitionRun, CandidateItem, ItemClassification, MediaAsset, Query, QueryBatch, ) Json = psycopg2.extras.Json psycopg2.extras.register_uuid() class PostgresAcquisitionRepository: """Repository backed by the formal cloud PostgreSQL schema. The repository does not commit by itself; callers own transaction scope via core.db_session.transaction or an equivalent connection boundary. """ def __init__(self, conn: Any): self.conn = conn def _one(self, sql: str, params: tuple[Any, ...]) -> dict[str, Any]: with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: cur.execute(sql, params) row = cur.fetchone() if row is None: raise RuntimeError("expected one row, got none") return dict(row) def _all(self, sql: str, params: tuple[Any, ...]) -> list[dict[str, Any]]: with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: cur.execute(sql, params) return [dict(row) for row in cur.fetchall()] def create_query_batch( self, *, name: str, source_type: str = "manual", generation_method: str | None = None, target_platforms: list[str] | None = None, status: str = "draft", metadata: dict[str, Any] | None = None, ) -> QueryBatch: row = self._one( """ INSERT INTO query_batches( name, source_type, generation_method, target_platforms, status, metadata ) VALUES (%s, %s, %s, %s, %s, %s) RETURNING * """, ( name, source_type, generation_method, target_platforms or [], status, Json(metadata or {}), ), ) return QueryBatch.model_validate(row) def add_query( self, *, batch_id: UUID | None, query_text: str, axes: dict[str, Any] | None = None, keep: bool | None = None, filter_reason: str | None = None, status: str = "draft", sort_order: int = 0, metadata: dict[str, Any] | None = None, ) -> Query: row = self._one( """ INSERT INTO queries( batch_id, query_text, axes, keep, filter_reason, status, sort_order, metadata ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s) RETURNING * """, ( batch_id, query_text, Json(axes or {}), keep, filter_reason, status, sort_order, Json(metadata or {}), ), ) return Query.model_validate(row) def list_queries_for_batch( self, batch_id: UUID, *, keep: bool | None = None, ) -> list[Query]: if keep is None: rows = self._all( """ SELECT * FROM queries WHERE batch_id = %s ORDER BY sort_order, created_at """, (batch_id,), ) else: rows = self._all( """ SELECT * FROM queries WHERE batch_id = %s AND keep IS NOT DISTINCT FROM %s ORDER BY sort_order, created_at """, (batch_id, keep), ) return [Query.model_validate(row) for row in rows] def get_query_batch(self, batch_id: UUID) -> QueryBatch: row = self._one("SELECT * FROM query_batches WHERE id = %s", (batch_id,)) return QueryBatch.model_validate(row) def create_acquisition_run( self, *, batch_id: UUID | None = None, run_key: str | None = None, status: str = "pending", note: str | None = None, metadata: dict[str, Any] | None = None, ) -> AcquisitionRun: row = self._one( """ INSERT INTO acquisition_runs(batch_id, run_key, status, note, metadata) VALUES (%s, %s, %s, %s, %s) ON CONFLICT (run_key) DO UPDATE SET status = EXCLUDED.status, note = EXCLUDED.note, metadata = EXCLUDED.metadata RETURNING * """, (batch_id, run_key, status, note, Json(metadata or {})), ) return AcquisitionRun.model_validate(row) def ensure_acquisition_job( self, *, run_id: UUID, query_id: UUID | None, platform: str, search_limit: int | None = None, display_limit: int | None = None, status: str = "pending", metadata: dict[str, Any] | None = None, ) -> AcquisitionJob: row = self._one( """ INSERT INTO acquisition_jobs( run_id, query_id, platform, search_limit, display_limit, status, metadata ) VALUES (%s, %s, %s, %s, %s, %s, %s) ON CONFLICT (run_id, query_id, platform) DO UPDATE SET search_limit = EXCLUDED.search_limit, display_limit = EXCLUDED.display_limit, metadata = EXCLUDED.metadata RETURNING * """, ( run_id, query_id, platform, search_limit, display_limit, status, Json(metadata or {}), ), ) return AcquisitionJob.model_validate(row) def update_acquisition_job( self, job_id: UUID, *, status: str, attempt_count: int | None = None, error_message: str | None = None, metadata: dict[str, Any] | None = None, ) -> AcquisitionJob: row = self._one( """ UPDATE acquisition_jobs SET status = %s, attempt_count = COALESCE(%s, attempt_count), error_message = %s, metadata = CASE WHEN %s THEN %s ELSE metadata END WHERE id = %s RETURNING * """, ( status, attempt_count, error_message, metadata is not None, Json(metadata or {}), job_id, ), ) return AcquisitionJob.model_validate(row) def upsert_candidate_item( self, *, platform: str, job_id: UUID | None = None, query_id: UUID | None = None, platform_item_id: str | None = None, canonical_url: str | None = None, content_type: str | None = None, title: str | None = None, author_name: str | None = None, raw_summary: str | None = None, status: str = "candidate", source_payload: dict[str, Any] | None = None, metadata: dict[str, Any] | None = None, error_message: str | None = None, ) -> CandidateItem: existing_id = None if platform_item_id: row = self._one_or_none( """ SELECT id FROM candidate_items WHERE platform = %s AND platform_item_id = %s ORDER BY created_at DESC LIMIT 1 """, (platform, platform_item_id), ) existing_id = row["id"] if row else None if existing_id: row = self._one( """ UPDATE candidate_items SET job_id = %s, query_id = %s, canonical_url = %s, content_type = %s, title = %s, author_name = %s, raw_summary = %s, status = %s, source_payload = %s, metadata = %s, error_message = %s WHERE id = %s RETURNING * """, ( job_id, query_id, canonical_url, content_type, title, author_name, raw_summary, status, Json(source_payload or {}), Json(metadata or {}), error_message, existing_id, ), ) else: row = self._one( """ INSERT INTO candidate_items( job_id, query_id, platform, platform_item_id, canonical_url, content_type, title, author_name, raw_summary, status, source_payload, metadata, error_message ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING * """, ( job_id, query_id, platform, platform_item_id, canonical_url, content_type, title, author_name, raw_summary, status, Json(source_payload or {}), Json(metadata or {}), error_message, ), ) return CandidateItem.model_validate(row) def _one_or_none(self, sql: str, params: tuple[Any, ...]) -> dict[str, Any] | None: with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: cur.execute(sql, params) row = cur.fetchone() return dict(row) if row else None def add_media_asset( self, *, item_id: UUID, media_type: str, source_url: str | None = None, oss_url: str | None = None, cdn_url: str | None = None, position: int = 0, status: str = "pending", source_payload: dict[str, Any] | None = None, metadata: dict[str, Any] | None = None, ) -> MediaAsset: row = self._one( """ INSERT INTO media_assets( item_id, media_type, source_url, oss_url, cdn_url, position, status, source_payload, metadata ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING * """, ( item_id, media_type, source_url, oss_url, cdn_url, position, status, Json(source_payload or {}), Json(metadata or {}), ), ) return MediaAsset.model_validate(row) def add_item_classification( self, *, item_id: UUID, is_creation_knowledge: bool | None = None, label: str | None = None, confidence: float | None = None, reason: str | None = None, model_name: str | None = None, prompt_version: str | None = None, result_payload: dict[str, Any] | None = None, status: str = "pending", error_message: str | None = None, ) -> ItemClassification: row = self._one( """ INSERT INTO item_classifications( item_id, is_creation_knowledge, label, confidence, reason, model_name, prompt_version, result_payload, status, error_message ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING * """, ( item_id, is_creation_knowledge, label, confidence, reason, model_name, prompt_version, Json(result_payload or {}), status, error_message, ), ) return ItemClassification.model_validate(row) def get_run_summary(self, run_id: UUID) -> dict[str, Any]: summary = self._one( """ SELECT ar.id, ar.run_key, ar.batch_id, ar.status, ar.started_at, ar.finished_at, COUNT(DISTINCT q.id)::int AS query_count, COUNT(DISTINCT aj.id)::int AS job_count, COUNT(DISTINCT ci.id)::int AS candidate_count, COUNT(DISTINCT ic.id) FILTER ( WHERE ic.is_creation_knowledge IS TRUE )::int AS creation_hit_count FROM acquisition_runs ar LEFT JOIN acquisition_jobs aj ON aj.run_id = ar.id LEFT JOIN queries q ON q.id = aj.query_id LEFT JOIN candidate_items ci ON ci.job_id = aj.id LEFT JOIN item_classifications ic ON ic.item_id = ci.id WHERE ar.id = %s GROUP BY ar.id """, (run_id,), ) summary["queries"] = self._all( """ SELECT q.id AS query_id, q.query_text, COUNT(DISTINCT aj.id)::int AS job_count, COUNT(DISTINCT ci.id)::int AS candidate_count, COUNT(DISTINCT ic.id) FILTER ( WHERE ic.is_creation_knowledge IS TRUE )::int AS creation_hit_count, jsonb_object_agg( aj.platform, jsonb_build_object( 'status', aj.status, 'attempt_count', aj.attempt_count, 'display_limit', aj.display_limit, 'search_limit', aj.search_limit, 'error_message', aj.error_message ) ) FILTER (WHERE aj.id IS NOT NULL) AS platforms FROM queries q JOIN acquisition_jobs aj ON aj.query_id = q.id LEFT JOIN candidate_items ci ON ci.job_id = aj.id LEFT JOIN item_classifications ic ON ic.item_id = ci.id WHERE aj.run_id = %s GROUP BY q.id, q.query_text, q.sort_order ORDER BY q.sort_order, q.query_text """, (run_id,), ) return summary def get_query_detail(self, *, run_id: UUID, query_id: UUID) -> dict[str, Any]: query = self._one("SELECT * FROM queries WHERE id = %s", (query_id,)) jobs = self._all( """ SELECT * FROM acquisition_jobs WHERE run_id = %s AND query_id = %s ORDER BY platform """, (run_id, query_id), ) items = self._all( """ SELECT ci.* FROM candidate_items ci JOIN acquisition_jobs aj ON aj.id = ci.job_id WHERE aj.run_id = %s AND ci.query_id = %s ORDER BY ci.platform, ci.created_at """, (run_id, query_id), ) item_ids = [row["id"] for row in items] media: list[dict[str, Any]] = [] classifications: list[dict[str, Any]] = [] if item_ids: media = self._all( """ SELECT * FROM media_assets WHERE item_id = ANY(%s) ORDER BY item_id, position """, (item_ids,), ) classifications = self._all( """ SELECT * FROM item_classifications WHERE item_id = ANY(%s) ORDER BY created_at DESC """, (item_ids,), ) return { "query": query, "jobs": jobs, "items": items, "media_assets": media, "classifications": classifications, } def list_creation_candidate_items( self, *, run_id: UUID | None = None, limit: int = 100, ) -> list[CandidateItem]: if run_id is None: rows = self._all( """ SELECT ci.* FROM candidate_items ci JOIN item_classifications ic ON ic.item_id = ci.id WHERE ic.is_creation_knowledge IS TRUE ORDER BY ci.created_at LIMIT %s """, (limit,), ) else: rows = self._all( """ SELECT ci.* FROM candidate_items ci JOIN acquisition_jobs aj ON aj.id = ci.job_id JOIN item_classifications ic ON ic.item_id = ci.id WHERE aj.run_id = %s AND ic.is_creation_knowledge IS TRUE ORDER BY ci.created_at LIMIT %s """, (run_id, limit), ) return [CandidateItem.model_validate(row) for row in rows] def get_candidate_item(self, item_id: UUID) -> CandidateItem: row = self._one("SELECT * FROM candidate_items WHERE id = %s", (item_id,)) return CandidateItem.model_validate(row) def list_media_assets_for_item(self, item_id: UUID) -> list[MediaAsset]: rows = self._all( """ SELECT * FROM media_assets WHERE item_id = %s ORDER BY position """, (item_id,), ) return [MediaAsset.model_validate(row) for row in rows]