/// PostgreSQL repository implementation for Phase 2.6 DB Integration. /// /// Connects ingest pipeline to persistent storage. /// Handles transactions, error recovery, and audit logging. use sqlx::{Pool, Postgres, Row, Transaction, Error as SqlxError}; use serde::{Deserialize, Serialize}; use chrono::{DateTime, Utc}; use mem_core::entity::Entity; use mem_core::edge::Edge; use crate::entity_repo::EntityRepoOps; use crate::edge_repo::EdgeRepoOps; /// Database connection error types #[derive(Debug, Clone)] pub enum DbError { ConnectionFailed(String), QueryFailed(String), TransactionFailed(String), DuplicateKey(String), NotFound(String), InvalidData(String), } impl std::fmt::Display for DbError { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { match self { DbError::ConnectionFailed(msg) => write!(f, "Connection failed: {}", msg), DbError::QueryFailed(msg) => write!(f, "Query failed: {}", msg), DbError::TransactionFailed(msg) => write!(f, "Transaction failed: {}", msg), DbError::DuplicateKey(msg) => write!(f, "Duplicate key: {}", msg), DbError::NotFound(msg) => write!(f, "Not found: {}", msg), DbError::InvalidData(msg) => write!(f, "Invalid data: {}", msg), } } } impl std::error::Error for DbError {} /// PostgreSQL repository pool pub struct DbPool { pool: Pool, } impl DbPool { /// Create new DB pool from connection string pub async fn new(database_url: &str) -> Result { let pool = Pool::::connect(database_url) .await .map_err(|e| DbError::ConnectionFailed(e.to_string()))?; Ok(DbPool { pool }) } /// Get pool for queries pub fn pool(&self) -> &Pool { &self.pool } /// Test connection pub async fn health_check(&self) -> Result<(), DbError> { sqlx::query("SELECT 1") .fetch_one(&self.pool) .await .map_err(|e| DbError::ConnectionFailed(e.to_string()))?; Ok(()) } } /// Persistent entity repository pub struct PersistentEntityRepo { pool: Pool, } impl PersistentEntityRepo { pub fn new(pool: Pool) -> Self { Self { pool } } /// Save entity to database (idempotent) pub async fn save(&self, entity: &Entity) -> Result { let query = r#" INSERT INTO memory_entity (id, entity_type, name, description, embedding, created_at, updated_at) VALUES ($1, $2, $3, $4, $5, $6, $7) ON CONFLICT(id) DO UPDATE SET name = EXCLUDED.name, description = EXCLUDED.description, updated_at = EXCLUDED.updated_at RETURNING id; "#; let id = sqlx::query_scalar::<_, String>(query) .bind(&entity.id) .bind(&entity.entity_type) .bind(&entity.name) .bind(&entity.description) .bind(&entity.embedding) .bind(Utc::now()) .bind(Utc::now()) .fetch_one(&self.pool) .await .map_err(|e| { if e.to_string().contains("duplicate") { DbError::DuplicateKey(format!("Entity {} already exists", entity.id)) } else { DbError::QueryFailed(e.to_string()) } })?; Ok(id) } /// Get entity by ID pub async fn get(&self, id: &str) -> Result, DbError> { let query = r#" SELECT id, entity_type, name, description, embedding, created_at, updated_at FROM memory_entity WHERE id = $1 AND deleted_at IS NULL; "#; let row = sqlx::query(query) .bind(id) .fetch_optional(&self.pool) .await .map_err(|e| DbError::QueryFailed(e.to_string()))?; Ok(row.map(|r| Entity { id: r.get("id"), entity_type: r.get("entity_type"), name: r.get("name"), description: r.get("description"), embedding: r.get("embedding"), created_at: r.get("created_at"), updated_at: r.get("updated_at"), })) } /// List entities with pagination pub async fn list(&self, limit: i64, offset: i64) -> Result, DbError> { let query = r#" SELECT id, entity_type, name, description, embedding, created_at, updated_at FROM memory_entity WHERE deleted_at IS NULL ORDER BY created_at DESC LIMIT $1 OFFSET $2; "#; let rows = sqlx::query(query) .bind(limit) .bind(offset) .fetch_all(&self.pool) .await .map_err(|e| DbError::QueryFailed(e.to_string()))?; Ok(rows.iter().map(|r| Entity { id: r.get("id"), entity_type: r.get("entity_type"), name: r.get("name"), description: r.get("description"), embedding: r.get("embedding"), created_at: r.get("created_at"), updated_at: r.get("updated_at"), }).collect()) } /// Soft delete entity pub async fn delete(&self, id: &str) -> Result<(), DbError> { let query = r#" UPDATE memory_entity SET deleted_at = $1 WHERE id = $2; "#; sqlx::query(query) .bind(Utc::now()) .bind(id) .execute(&self.pool) .await .map_err(|e| DbError::QueryFailed(e.to_string()))?; Ok(()) } } /// Persistent edge repository pub struct PersistentEdgeRepo { pool: Pool, } impl PersistentEdgeRepo { pub fn new(pool: Pool) -> Self { Self { pool } } /// Save edge to database (idempotent) pub async fn save(&self, edge: &Edge) -> Result { let query = r#" INSERT INTO memory_edge (id, source_id, target_id, relation_type, fact, strength, t_valid, t_invalid, t_created, t_expired) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) ON CONFLICT(id) DO UPDATE SET strength = EXCLUDED.strength, t_invalid = EXCLUDED.t_invalid, t_expired = EXCLUDED.t_expired RETURNING id; "#; let id = sqlx::query_scalar::<_, String>(query) .bind(&edge.id) .bind(&edge.source_id) .bind(&edge.target_id) .bind(&edge.relation_type) .bind(&edge.fact) .bind(edge.strength) .bind(edge.t_valid) .bind(edge.t_invalid) .bind(edge.t_created) .bind(edge.t_expired) .fetch_one(&self.pool) .await .map_err(|e| { if e.to_string().contains("duplicate") { DbError::DuplicateKey(format!("Edge {} already exists", edge.id)) } else { DbError::QueryFailed(e.to_string()) } })?; Ok(id) } /// Get edge by ID pub async fn get(&self, id: &str) -> Result, DbError> { let query = r#" SELECT id, source_id, target_id, relation_type, fact, strength, t_valid, t_invalid, t_created, t_expired FROM memory_edge WHERE id = $1 AND t_expired IS NULL; "#; let row = sqlx::query(query) .bind(id) .fetch_optional(&self.pool) .await .map_err(|e| DbError::QueryFailed(e.to_string()))?; Ok(row.map(|r| Edge { id: r.get("id"), source_id: r.get("source_id"), target_id: r.get("target_id"), relation_type: r.get("relation_type"), fact: r.get("fact"), strength: r.get("strength"), t_valid: r.get("t_valid"), t_invalid: r.get("t_invalid"), t_created: r.get("t_created"), t_expired: r.get("t_expired"), })) } /// List edges for a source entity pub async fn list_from(&self, source_id: &str, limit: i64) -> Result, DbError> { let query = r#" SELECT id, source_id, target_id, relation_type, fact, strength, t_valid, t_invalid, t_created, t_expired FROM memory_edge WHERE source_id = $1 AND t_expired IS NULL AND t_invalid IS NULL ORDER BY t_created DESC LIMIT $2; "#; let rows = sqlx::query(query) .bind(source_id) .bind(limit) .fetch_all(&self.pool) .await .map_err(|e| DbError::QueryFailed(e.to_string()))?; Ok(rows.iter().map(|r| Edge { id: r.get("id"), source_id: r.get("source_id"), target_id: r.get("target_id"), relation_type: r.get("relation_type"), fact: r.get("fact"), strength: r.get("strength"), t_valid: r.get("t_valid"), t_invalid: r.get("t_invalid"), t_created: r.get("t_created"), t_expired: r.get("t_expired"), }).collect()) } /// Mark edge as contradicted (soft delete) pub async fn invalidate(&self, id: &str) -> Result<(), DbError> { let query = r#" UPDATE memory_edge SET t_invalid = $1 WHERE id = $2; "#; sqlx::query(query) .bind(Utc::now()) .bind(id) .execute(&self.pool) .await .map_err(|e| DbError::QueryFailed(e.to_string()))?; Ok(()) } } /// Review queue entry for human verification #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ReviewQueueEntry { pub id: String, pub extraction_type: String, // "entity" | "edge" | "contradiction" pub content: serde_json::Value, // Full extracted data pub status: String, // "pending" | "approved" | "rejected" pub created_at: DateTime, pub reviewed_at: Option>, pub reviewed_by: Option, // User ID who reviewed pub rejection_reason: Option, } /// Review queue repository pub struct ReviewQueueRepo { pool: Pool, } impl ReviewQueueRepo { pub fn new(pool: Pool) -> Self { Self { pool } } /// Add item to review queue pub async fn enqueue(&self, entry: &ReviewQueueEntry) -> Result { let query = r#" INSERT INTO review_queue (id, extraction_type, content, status, created_at) VALUES ($1, $2, $3, $4, $5) RETURNING id; "#; let id = sqlx::query_scalar::<_, String>(query) .bind(&entry.id) .bind(&entry.extraction_type) .bind(&entry.content) .bind(&entry.status) .bind(Utc::now()) .fetch_one(&self.pool) .await .map_err(|e| DbError::QueryFailed(e.to_string()))?; Ok(id) } /// Get pending items for review pub async fn list_pending(&self, limit: i64) -> Result, DbError> { let query = r#" SELECT id, extraction_type, content, status, created_at, reviewed_at, reviewed_by, rejection_reason FROM review_queue WHERE status = 'pending' ORDER BY created_at ASC LIMIT $1; "#; let rows = sqlx::query(query) .bind(limit) .fetch_all(&self.pool) .await .map_err(|e| DbError::QueryFailed(e.to_string()))?; Ok(rows.iter().map(|r| ReviewQueueEntry { id: r.get("id"), extraction_type: r.get("extraction_type"), content: r.get("content"), status: r.get("status"), created_at: r.get("created_at"), reviewed_at: r.get("reviewed_at"), reviewed_by: r.get("reviewed_by"), rejection_reason: r.get("rejection_reason"), }).collect()) } /// Approve review queue entry pub async fn approve(&self, id: &str, reviewed_by: &str) -> Result<(), DbError> { let query = r#" UPDATE review_queue SET status = 'approved', reviewed_at = $1, reviewed_by = $2 WHERE id = $3; "#; sqlx::query(query) .bind(Utc::now()) .bind(reviewed_by) .bind(id) .execute(&self.pool) .await .map_err(|e| DbError::QueryFailed(e.to_string()))?; Ok(()) } /// Reject review queue entry pub async fn reject(&self, id: &str, reviewed_by: &str, reason: &str) -> Result<(), DbError> { let query = r#" UPDATE review_queue SET status = 'rejected', reviewed_at = $1, reviewed_by = $2, rejection_reason = $3 WHERE id = $4; "#; sqlx::query(query) .bind(Utc::now()) .bind(reviewed_by) .bind(reason) .bind(id) .execute(&self.pool) .await .map_err(|e| DbError::QueryFailed(e.to_string()))?; Ok(()) } } /// Extraction Audit Repository (Immutable log for audit trail) #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ExtractionAuditEntry { pub id: String, pub extraction_type: String, // "entity" | "edge" pub extraction_id: String, // ID of extracted entity/edge pub source_content: String, // Original text pub extracted_data: serde_json::Value, pub llm_confidence: Option, pub contradiction_score: Option, pub status: String, // "extracted" | "approved" | "rejected" pub extracted_at: DateTime, pub extracted_by: String, // User or "system" } pub struct ExtractionAuditRepo { pool: Pool, } impl ExtractionAuditRepo { pub fn new(pool: Pool) -> Self { Self { pool } } /// Log an extraction attempt (immutable append) pub async fn log_extraction(&self, entry: &ExtractionAuditEntry) -> Result { let query = r#" INSERT INTO extraction_audit (id, extraction_type, extraction_id, source_content, extracted_data, llm_confidence, contradiction_score, status, extracted_at, extracted_by) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) RETURNING id; "#; let id = sqlx::query_scalar::<_, String>(query) .bind(&entry.id) .bind(&entry.extraction_type) .bind(&entry.extraction_id) .bind(&entry.source_content) .bind(&entry.extracted_data) .bind(entry.llm_confidence) .bind(entry.contradiction_score) .bind(&entry.status) .bind(entry.extracted_at) .bind(&entry.extracted_by) .fetch_one(&self.pool) .await .map_err(|e| DbError::QueryFailed(e.to_string()))?; Ok(id) } /// Get audit trail for an extracted item pub async fn get_history(&self, extraction_id: &str) -> Result, DbError> { let query = r#" SELECT id, extraction_type, extraction_id, source_content, extracted_data, llm_confidence, contradiction_score, status, extracted_at, extracted_by FROM extraction_audit WHERE extraction_id = $1 ORDER BY extracted_at DESC; "#; let rows = sqlx::query(query) .bind(extraction_id) .fetch_all(&self.pool) .await .map_err(|e| DbError::QueryFailed(e.to_string()))?; Ok(rows.iter().map(|r| ExtractionAuditEntry { id: r.get("id"), extraction_type: r.get("extraction_type"), extraction_id: r.get("extraction_id"), source_content: r.get("source_content"), extracted_data: r.get("extracted_data"), llm_confidence: r.get("llm_confidence"), contradiction_score: r.get("contradiction_score"), status: r.get("status"), extracted_at: r.get("extracted_at"), extracted_by: r.get("extracted_by"), }).collect()) } } #[cfg(test)] mod tests { use super::*; #[test] fn test_db_error_display() { let err = DbError::ConnectionFailed("test".to_string()); assert!(err.to_string().contains("Connection failed")); } #[test] fn test_review_queue_entry_creation() { let entry = ReviewQueueEntry { id: "test-1".to_string(), extraction_type: "entity".to_string(), content: serde_json::json!({"name": "test"}), status: "pending".to_string(), created_at: Utc::now(), reviewed_at: None, reviewed_by: None, rejection_reason: None, }; assert_eq!(entry.extraction_type, "entity"); } #[test] fn test_dead_letter_entry_creation() { let entry = DeadLetterEntry { id: "dlq-1".to_string(), original_content: "test content".to_string(), error_message: "extraction failed".to_string(), error_type: "extraction_failed".to_string(), retry_count: 0, max_retries: 3, created_at: Utc::now(), last_retry_at: None, }; assert_eq!(entry.retry_count, 0); assert!(entry.retry_count < entry.max_retries); } }