poimen-memory/crates/mem-store/src/db_repo.rs

/// 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<Postgres>,
}

impl DbPool {
    /// Create new DB pool from connection string
    pub async fn new(database_url: &str) -> Result<Self, DbError> {
        let pool = Pool::<Postgres>::connect(database_url)
            .await
            .map_err(|e| DbError::ConnectionFailed(e.to_string()))?;
        
        Ok(DbPool { pool })
    }
    
    /// Get pool for queries
    pub fn pool(&self) -> &Pool<Postgres> {
        &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<Postgres>,
}

impl PersistentEntityRepo {
    pub fn new(pool: Pool<Postgres>) -> Self {
        Self { pool }
    }
    
    /// Save entity to database (idempotent)
    pub async fn save(&self, entity: &Entity) -> Result<String, DbError> {
        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<Option<Entity>, 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<Vec<Entity>, 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<Postgres>,
}

impl PersistentEdgeRepo {
    pub fn new(pool: Pool<Postgres>) -> Self {
        Self { pool }
    }
    
    /// Save edge to database (idempotent)
    pub async fn save(&self, edge: &Edge) -> Result<String, DbError> {
        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<Option<Edge>, 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<Vec<Edge>, 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<Utc>,
    pub reviewed_at: Option<DateTime<Utc>>,
    pub reviewed_by: Option<String>,  // User ID who reviewed
    pub rejection_reason: Option<String>,
}

/// Review queue repository
pub struct ReviewQueueRepo {
    pool: Pool<Postgres>,
}

impl ReviewQueueRepo {
    pub fn new(pool: Pool<Postgres>) -> Self {
        Self { pool }
    }
    
    /// Add item to review queue
    pub async fn enqueue(&self, entry: &ReviewQueueEntry) -> Result<String, DbError> {
        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<Vec<ReviewQueueEntry>, 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<f32>,
    pub contradiction_score: Option<f32>,
    pub status: String,           // "extracted" | "approved" | "rejected"
    pub extracted_at: DateTime<Utc>,
    pub extracted_by: String,     // User or "system"
}

pub struct ExtractionAuditRepo {
    pool: Pool<Postgres>,
}

impl ExtractionAuditRepo {
    pub fn new(pool: Pool<Postgres>) -> Self {
        Self { pool }
    }
    
    /// Log an extraction attempt (immutable append)
    pub async fn log_extraction(&self, entry: &ExtractionAuditEntry) -> Result<String, DbError> {
        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<Vec<ExtractionAuditEntry>, 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);
    }
}