"""PostgreSQL implementation for formal decode-content state.""" from __future__ import annotations from typing import Any from uuid import UUID import psycopg2.extras from decode_content.models import ( ContractSnapshot, DecodeJob, DecodeResult, IngestRecord, KnowledgeParticle, PayloadDraft, ScopeResult, ) Json = psycopg2.extras.Json psycopg2.extras.register_uuid() def _json(value: dict[str, Any] | None) -> Json: return Json(value or {}) def _stage_value(value: Any) -> str | None: if isinstance(value, list): return " / ".join(str(item) for item in value if item) if value: return str(value) return None class PostgresDecodeRepository: """Repository backed by the formal PostgreSQL schema. Like the acquisition repository, this class owns no transaction boundary; callers pass an open connection and commit/rollback outside. """ 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 _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 _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_decode_job( self, *, item_id: UUID, status: str = "pending", metadata: dict[str, Any] | None = None, ) -> DecodeJob: row = self._one( """ INSERT INTO decode_jobs(item_id, status, metadata, started_at) VALUES (%s, %s, %s, CASE WHEN %s = 'running' THEN now() ELSE NULL END) RETURNING * """, (item_id, status, _json(metadata), status), ) return DecodeJob.model_validate(row) def save_decode_result( self, *, item_id: UUID, decode_job_id: UUID | None = None, read_result: dict[str, Any] | None = None, gate_result: dict[str, Any] | None = None, framing_result: dict[str, Any] | None = None, status: str = "draft", ) -> DecodeResult: row = self._one( """ INSERT INTO decode_results( decode_job_id, item_id, read_result, gate_result, framing_result, status ) VALUES (%s, %s, %s, %s, %s, %s) RETURNING * """, ( decode_job_id, item_id, _json(read_result), _json(gate_result), _json(framing_result), status, ), ) if decode_job_id: self._one( """ UPDATE decode_jobs SET status = %s, finished_at = now(), error_message = NULL WHERE id = %s RETURNING * """, (status, decode_job_id), ) return DecodeResult.model_validate(row) def get_decode_result_for_item(self, item_id: UUID) -> DecodeResult: row = self._one( """ SELECT * FROM decode_results WHERE item_id = %s ORDER BY created_at DESC LIMIT 1 """, (item_id,), ) return DecodeResult.model_validate(row) def mark_running_decode_jobs_failed(self, item_id: UUID, error_message: str) -> None: with self.conn.cursor() as cur: cur.execute( """ UPDATE decode_jobs SET status = 'failed', finished_at = now(), error_message = %s WHERE item_id = %s AND status = 'running' """, (error_message, item_id), ) def save_knowledge_particle( self, *, item_id: UUID, particle_type: str, title: str, decode_result_id: UUID | None = None, parent_particle_id: UUID | None = None, content: dict[str, Any] | None = None, status: str = "draft", ) -> KnowledgeParticle: content = content or {} creation_stage = _stage_value(content.get("创作阶段")) if creation_stage is None: for step in content.get("steps") or []: creation_stage = _stage_value(step.get("创作阶段")) if creation_stage: break row = self._one( """ INSERT INTO knowledge_particles( decode_result_id, item_id, parent_particle_id, particle_type, title, business_stage, creation_stage, content, status ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING * """, ( decode_result_id, item_id, parent_particle_id, particle_type, title, _stage_value(content.get("业务阶段")), creation_stage, _json(content), status, ), ) return KnowledgeParticle.model_validate(row) def list_knowledge_particles(self, item_id: UUID | None = None) -> list[KnowledgeParticle]: if item_id is None: rows = self._all( "SELECT * FROM knowledge_particles ORDER BY created_at DESC", (), ) else: rows = self._all( """ SELECT * FROM knowledge_particles WHERE item_id = %s ORDER BY sort_order, created_at """, (item_id,), ) return [KnowledgeParticle.model_validate(row) for row in rows] def save_scope_result( self, *, scope_type: str, scope_value: str, particle_id: UUID | None = None, item_id: UUID | None = None, evidence: dict[str, Any] | None = None, status: str = "draft", ) -> ScopeResult: evidence = evidence or {} row = self._one( """ INSERT INTO scope_results( particle_id, item_id, scope_type, scope_value, is_reused, matched_scope_id, confidence, evidence, status ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING * """, ( particle_id, item_id, scope_type, scope_value, evidence.get("is_reused"), evidence.get("matched_scope_id"), evidence.get("confidence"), _json(evidence), status, ), ) return ScopeResult.model_validate(row) def list_scope_results(self, item_id: UUID | None = None) -> list[ScopeResult]: if item_id is None: rows = self._all("SELECT * FROM scope_results ORDER BY created_at DESC", ()) else: rows = self._all( """ SELECT * FROM scope_results WHERE item_id = %s ORDER BY created_at """, (item_id,), ) return [ScopeResult.model_validate(row) for row in rows] def save_payload_draft( self, *, payload: dict[str, Any], particle_id: UUID | None = None, item_id: UUID | None = None, review_status: str = "pending", ingest_ready: bool = False, status: str = "draft", ) -> PayloadDraft: row = self._one( """ INSERT INTO payload_drafts( particle_id, item_id, payload, review_status, ingest_ready, status ) VALUES (%s, %s, %s, %s, %s, %s) RETURNING * """, (particle_id, item_id, _json(payload), review_status, ingest_ready, status), ) return PayloadDraft.model_validate(row) def list_payload_drafts(self, item_id: UUID | None = None) -> list[PayloadDraft]: if item_id is None: rows = self._all( "SELECT * FROM payload_drafts ORDER BY created_at DESC", (), ) else: rows = self._all( """ SELECT * FROM payload_drafts WHERE item_id = %s ORDER BY created_at """, (item_id,), ) return [PayloadDraft.model_validate(row) for row in rows] def mark_payload_draft_ingested(self, payload_draft_id: UUID) -> PayloadDraft: row = self._one( """ UPDATE payload_drafts SET review_status = 'approved', ingest_ready = true, status = 'ingested', error_message = NULL WHERE id = %s RETURNING * """, (payload_draft_id,), ) return PayloadDraft.model_validate(row) def save_ingest_record( self, *, payload_draft_id: UUID | None = None, target_system: str, target_id: str | None = None, status: str = "pending", attempt_count: int = 1, response_payload: dict[str, Any] | None = None, error_message: str | None = None, ) -> IngestRecord: row = self._one( """ INSERT INTO ingest_records( payload_draft_id, target_system, target_id, status, attempt_count, response_payload, error_message ) VALUES (%s, %s, %s, %s, %s, %s, %s) RETURNING * """, ( payload_draft_id, target_system, target_id, status, attempt_count, _json(response_payload), error_message, ), ) return IngestRecord.model_validate(row) def list_ingest_records(self, item_id: UUID | None = None) -> list[IngestRecord]: if item_id is None: rows = self._all("SELECT * FROM ingest_records ORDER BY created_at DESC", ()) else: rows = self._all( """ SELECT ir.* FROM ingest_records ir JOIN payload_drafts pd ON pd.id = ir.payload_draft_id WHERE pd.item_id = %s ORDER BY ir.created_at """, (item_id,), ) return [IngestRecord.model_validate(row) for row in rows] def save_contract_snapshot( self, *, contract_name: str, contract_type: str, version_label: str | None = None, content_hash: str | None = None, source_path: str | None = None, snapshot: dict[str, Any] | None = None, pipeline_run_id: UUID | None = None, ) -> ContractSnapshot: row = self._one( """ INSERT INTO contract_artifacts( contract_name, contract_type, version_label, content_hash, source_path, snapshot ) VALUES (%s, %s, %s, %s, %s, %s) ON CONFLICT (contract_name, contract_type, source_path, content_hash) DO UPDATE SET version_label = COALESCE(EXCLUDED.version_label, contract_artifacts.version_label), source_path = COALESCE(EXCLUDED.source_path, contract_artifacts.source_path), snapshot = EXCLUDED.snapshot, updated_at = now() RETURNING id, contract_name, contract_type, version_label, content_hash, source_path, snapshot, created_at, updated_at """, ( contract_name, contract_type, version_label, content_hash or "", source_path, _json(snapshot), ), ) if pipeline_run_id is not None: self._one( """ INSERT INTO run_contract_artifacts(pipeline_run_id, contract_artifact_id) VALUES (%s, %s) ON CONFLICT (pipeline_run_id, contract_artifact_id) DO UPDATE SET created_at = run_contract_artifacts.created_at RETURNING * """, (pipeline_run_id, row["id"]), ) return ContractSnapshot.model_validate(row)