|
|
@@ -0,0 +1,551 @@
|
|
|
+"""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]
|
SamLee