poimen-memory/crates/mem-cli/src/parallel_dual_write.rs

//! Parallel Dual-Write Indexer (Refactored)
//!
//! pgvector (primary, must succeed) + OpenSearch (secondary, fire-and-forget)
//! Both execute concurrently via tokio::join!

use anyhow::{anyhow, Result};
use sha2::{Digest, Sha256};
use sqlx::PgPool;
use uuid::Uuid;
use pgvector::Vector;
use std::sync::Arc;
use crate::opensearch_client::OpenSearchClient;
use serde::{Deserialize, Serialize};

#[derive(Clone)]
pub struct ParallelDualWriteIndexer {
    pool: PgPool,
    opensearch: Option<Arc<OpenSearchClient>>,
}

/// Chunk to index
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexableChunk {
    pub chunk_id: String,
    pub content: String,
    pub source: String,
    pub project: String,
    pub level: String,
    pub breadcrumb: Vec<String>,
}

/// Result of parallel dual-write
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DualWriteResult {
    pub chunk_id: String,
    pub pgvector_success: bool,
    pub opensearch_success: bool,
    pub error: Option<String>,
}

impl ParallelDualWriteIndexer {
    pub fn new(pool: PgPool, opensearch: Option<Arc<OpenSearchClient>>) -> Self {
        Self { pool, opensearch }
    }

    /// Index chunk to both pgvector AND OpenSearch in parallel
    pub async fn index_parallel(&self, chunk: &IndexableChunk, embedding: &[f32]) -> Result<DualWriteResult> {
        let chunk_id = chunk.chunk_id.clone();

        // PARALLEL: Execute both writes concurrently
        let (pgvector_result, opensearch_result) = tokio::join!(
            self.write_pgvector(chunk, embedding),
            self.write_opensearch(chunk, embedding)
        );

        let pgvector_success = pgvector_result.is_ok();
        let opensearch_success = opensearch_result.is_ok();

        let error = if !pgvector_success {
            pgvector_result.err().map(|e| e.to_string())
        } else if !opensearch_success {
            opensearch_result.err().map(|e| e.to_string())
        } else {
            None
        };

        // Primary (pgvector) success = operation success
        if !pgvector_success {
            return Err(anyhow!("pgvector write failed: {:?}", error));
        }

        Ok(DualWriteResult {
            chunk_id,
            pgvector_success,
            opensearch_success,
            error,
        })
    }

    /// Write to pgvector (PRIMARY - must succeed)
    async fn write_pgvector(&self, chunk: &IndexableChunk, embedding: &[f32]) -> Result<()> {
        let vector = Vector::from(embedding.to_vec());
        let chunk_hash = self.compute_hash(&chunk.content);

        sqlx::query(
            "INSERT INTO memory_vector (id, project, level, text, embedding, breadcrumb, source, chunk_hash, indexed_at)
             VALUES ($1, $2, $3, $4, $5, $6, $7, $8, now())
             ON CONFLICT (id) DO UPDATE SET
               indexed_at = now(),
               embedding = $5"
        )
        .bind(&chunk.chunk_id)
        .bind(&chunk.project)
        .bind(&chunk.level)
        .bind(&chunk.content)
        .bind(&vector)
        .bind(chunk.breadcrumb.join(" > "))
        .bind(&chunk.source)
        .bind(&chunk_hash)
        .execute(&self.pool)
        .await?;

        tracing::debug!("pgvector indexed: {}", chunk.chunk_id);
        Ok(())
    }

    /// Write to OpenSearch (SECONDARY - fire-and-forget)
    async fn write_opensearch(&self, chunk: &IndexableChunk, _embedding: &[f32]) -> Result<()> {
        if self.opensearch.is_none() {
            return Ok(());
        }

        let opensearch = self.opensearch.clone().unwrap();
        let chunk_id = chunk.chunk_id.clone();
        let chunk = chunk.clone();

        // Spawn background task (non-blocking)
        tokio::spawn(async move {
            let result = opensearch.index_chunk(
                &chunk_id,
                &chunk.content,
                &chunk.source,
                &chunk.project,
                &chunk.level,
                &chunk.breadcrumb.join(" > "),
            ).await;

            match result {
                Ok(_) => tracing::debug!("OpenSearch indexed (async): {}", chunk_id),
                Err(e) => tracing::warn!("OpenSearch index failed (async, non-blocking): {}: {}", chunk_id, e),
            }
        });

        Ok(())
    }

    /// Batch parallel index (multiple chunks)
    pub async fn index_batch_parallel(
        &self,
        chunks: Vec<(&IndexableChunk, Vec<f32>)>,
    ) -> Vec<DualWriteResult> {
        let futures = chunks.into_iter().map(|(chunk, embedding)| {
            self.index_parallel(chunk, &embedding)
        });

        futures::future::join_all(futures)
            .await
            .into_iter()
            .filter_map(|r| r.ok())
            .collect()
    }

    /// Compute SHA256 hash
    fn compute_hash(&self, content: &str) -> String {
        let mut hasher = Sha256::new();
        hasher.update(content.as_bytes());
        format!("{:x}", hasher.finalize())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_indexable_chunk_structure() {
        let chunk = IndexableChunk {
            chunk_id: "c1".to_string(),
            content: "test".to_string(),
            source: "src".to_string(),
            project: "proj".to_string(),
            level: "L1".to_string(),
            breadcrumb: vec!["a".to_string()],
        };
        assert_eq!(chunk.chunk_id, "c1");
    }

    #[test]
    fn test_dual_write_result_structure() {
        let result = DualWriteResult {
            chunk_id: "c1".to_string(),
            pgvector_success: true,
            opensearch_success: true,
            error: None,
        };
        assert!(result.pgvector_success);
    }

    #[test]
    fn test_parallel_indexer_creation() {
        let pool = sqlx::postgres::PgPoolOptions::new()
            .max_connections(1)
            .build_lazy();
        let indexer = ParallelDualWriteIndexer::new(pool, None);
        assert!(indexer.opensearch.is_none());
    }

    #[test]
    fn test_hash_computation() {
        let pool = sqlx::postgres::PgPoolOptions::new()
            .max_connections(1)
            .build_lazy();
        let indexer = ParallelDualWriteIndexer::new(pool, None);
        let hash1 = indexer.compute_hash("test");
        let hash2 = indexer.compute_hash("test");
        assert_eq!(hash1, hash2);
    }

    #[test]
    fn test_hash_different_content() {
        let pool = sqlx::postgres::PgPoolOptions::new()
            .max_connections(1)
            .build_lazy();
        let indexer = ParallelDualWriteIndexer::new(pool, None);
        let hash1 = indexer.compute_hash("test1");
        let hash2 = indexer.compute_hash("test2");
        assert_ne!(hash1, hash2);
    }

    #[test]
    fn test_dual_write_result_pgvector_failed() {
        let result = DualWriteResult {
            chunk_id: "c1".to_string(),
            pgvector_success: false,
            opensearch_success: true,
            error: Some("pgvector failed".to_string()),
        };
        assert!(!result.pgvector_success);
        assert!(result.error.is_some());
    }

    #[test]
    fn test_dual_write_result_opensearch_failed() {
        let result = DualWriteResult {
            chunk_id: "c1".to_string(),
            pgvector_success: true,
            opensearch_success: false,
            error: Some("opensearch failed".to_string()),
        };
        assert!(result.pgvector_success);
        assert!(!result.opensearch_success);
    }

    #[test]
    fn test_breadcrumb_join() {
        let breadcrumb = vec!["a".to_string(), "b".to_string(), "c".to_string()];
        let joined = breadcrumb.join(" > ");
        assert_eq!(joined, "a > b > c");
    }

    #[test]
    fn test_chunk_source_tracking() {
        let chunk = IndexableChunk {
            chunk_id: "c1".to_string(),
            content: "test".to_string(),
            source: "transcript://session-123".to_string(),
            project: "poimen".to_string(),
            level: "L1".to_string(),
            breadcrumb: vec![],
        };
        assert!(chunk.source.contains("session"));
    }
}