//! 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>, } /// 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, } /// 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, } impl ParallelDualWriteIndexer { pub fn new(pool: PgPool, opensearch: Option>) -> Self { Self { pool, opensearch } } /// Index chunk to both pgvector AND OpenSearch in parallel pub async fn index_parallel(&self, chunk: &IndexableChunk, embedding: &[f32]) -> Result { 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)>, ) -> Vec { 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")); } }