lc_vector_stores/

sqlite_store.rs

// lc-vector-stores/src/sqlite_store.rs
//! SQLite document store implementation

use async_trait::async_trait;
use lc_shared::splitter::{RecursiveCharacterSplitter, TextSplitter};
use rusqlite::Connection;
use std::sync::Arc;
use std::sync::Mutex;
use uuid::Uuid;

use crate::document_store::{ChunkDocument, ChunkedDocumentStoreTrait, DocumentStore};
use crate::{Document, VectorStoreError};

/// Parses the metadata column; on corruption, logs and falls back to an empty map instead of silently swallowing the error.
fn parse_metadata_or_default(raw: &str) -> std::collections::HashMap<String, serde_json::Value> {
    serde_json::from_str(raw).unwrap_or_else(|e| {
        log::warn!(
            "SQLite metadata in store is corrupted, falling back to empty map: {}",
            e
        );
        std::collections::HashMap::new()
    })
}

/// SQLite document store configuration
#[derive(Debug, Clone)]
pub struct SQLiteStoreConfig {
    /// SQLite database file path
    pub db_path: String,
}

impl Default for SQLiteStoreConfig {
    fn default() -> Self {
        Self {
            db_path: "langchainrust.db".to_string(),
        }
    }
}

impl SQLiteStoreConfig {
    /// Creates a config from a database file path.
    pub fn new(path: impl Into<String>) -> Self {
        Self {
            db_path: path.into(),
        }
    }
}

/// SQLite document store implementation
pub struct SQLiteDocumentStore {
    conn: Arc<Mutex<Connection>>,
}

impl SQLiteDocumentStore {
    /// Opens the SQLite database and creates the required tables and indexes.
    pub fn new(config: SQLiteStoreConfig) -> Result<Self, VectorStoreError> {
        let conn = Connection::open(&config.db_path)
            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        conn.execute_batch(
            "CREATE TABLE IF NOT EXISTS documents (
                id TEXT PRIMARY KEY, content TEXT NOT NULL,
                metadata TEXT NOT NULL DEFAULT '{}'
            );
            CREATE TABLE IF NOT EXISTS chunks (
                chunk_id TEXT PRIMARY KEY, parent_id TEXT NOT NULL,
                content TEXT NOT NULL, segment INTEGER NOT NULL,
                metadata TEXT NOT NULL DEFAULT '{}'
            );
            CREATE INDEX IF NOT EXISTS idx_chunks_parent ON chunks(parent_id);
            CREATE INDEX IF NOT EXISTS idx_chunks_segment ON chunks(parent_id, segment);",
        )
        .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        Ok(Self {
            conn: Arc::new(Mutex::new(conn)),
        })
    }

    /// Explicitly triggers persistence to disk.
    ///
    /// C3: the fake default `ChunkedDocumentStoreTrait::save` has been removed from the trait;
    /// SQLite persists automatically on every write, so this method is a compatibility no-op
    /// (aligned with Redis's `save_to_disk`).
    pub async fn save(&self) -> Result<(), VectorStoreError> {
        Ok(())
    }
}

#[async_trait]
impl DocumentStore for SQLiteDocumentStore {
    async fn add_document(&self, document: Document) -> Result<String, VectorStoreError> {
        let id = document
            .id
            .clone()
            .unwrap_or_else(|| Uuid::new_v4().to_string());
        let meta = serde_json::to_string(&document.metadata).unwrap_or_else(|_| "{}".to_string());
        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        conn.execute(
            "INSERT OR REPLACE INTO documents (id, content, metadata) VALUES (?1, ?2, ?3)",
            rusqlite::params![id, document.content, meta],
        )
        .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        Ok(id)
    }

    async fn add_documents(
        &self,
        documents: Vec<Document>,
    ) -> Result<Vec<String>, VectorStoreError> {
        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let mut ids = Vec::new();
        for doc in documents {
            let id = doc.id.clone().unwrap_or_else(|| Uuid::new_v4().to_string());
            let meta = serde_json::to_string(&doc.metadata).unwrap_or_else(|_| "{}".to_string());
            conn.execute(
                "INSERT OR REPLACE INTO documents (id, content, metadata) VALUES (?1, ?2, ?3)",
                rusqlite::params![id, doc.content, meta],
            )
            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
            ids.push(id);
        }
        Ok(ids)
    }

    async fn get_document(&self, id: &str) -> Result<Option<Document>, VectorStoreError> {
        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let mut stmt = conn
            .prepare("SELECT id, content, metadata FROM documents WHERE id = ?1")
            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let result = stmt.query_row(rusqlite::params![id], |row| {
            let id: String = row.get(0)?;
            let content: String = row.get(1)?;
            let meta_str: String = row.get(2)?;
            Ok(Document {
                id: Some(id),
                content,
                metadata: parse_metadata_or_default(&meta_str),
            })
        });
        match result {
            Ok(doc) => Ok(Some(doc)),
            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
            Err(e) => Err(VectorStoreError::StorageError(e.to_string())),
        }
    }

    async fn delete_document(&self, id: &str) -> Result<(), VectorStoreError> {
        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        conn.execute(
            "DELETE FROM chunks WHERE parent_id = ?1",
            rusqlite::params![id],
        )
        .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        conn.execute("DELETE FROM documents WHERE id = ?1", rusqlite::params![id])
            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        Ok(())
    }

    async fn count(&self) -> usize {
        // M-23: `count` returns usize (cannot `?`), so a poisoned sync lock
        // surfaces as a logged error + 0, matching the error handling below.
        let conn = match self.conn.lock() {
            Ok(guard) => guard,
            Err(e) => {
                log::error!("SQLite lock poisoned in count(documents), returning 0: {}", e);
                return 0;
            }
        };
        match conn.query_row("SELECT COUNT(*) FROM documents", [], |r| r.get(0)) {
            Ok(count) => count,
            // M3: the trait returns usize so errors cannot propagate — stop silently swallowing
            // errors and returning 0; log an error to surface storage failures, consistent with
            // the "no silent row drops" pattern already used in this file.
            Err(e) => {
                log::error!("SQLite count(documents) query failed, returning 0: {}", e);
                0
            }
        }
    }

    async fn clear(&self) -> Result<(), VectorStoreError> {
        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        conn.execute_batch("DELETE FROM chunks; DELETE FROM documents;")
            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        Ok(())
    }
}

#[async_trait]
impl ChunkedDocumentStoreTrait for SQLiteDocumentStore {
    async fn add_parent_document(
        &self,
        document: Document,
        chunk_size: usize,
    ) -> Result<(String, Vec<String>), VectorStoreError> {
        let splitter = RecursiveCharacterSplitter::new(chunk_size, chunk_size / 10);
        let chunks_text = splitter.split_text(&document.content);
        let parent_id = document
            .id
            .clone()
            .unwrap_or_else(|| Uuid::new_v4().to_string());
        let meta = serde_json::to_string(&document.metadata).unwrap_or_else(|_| "{}".to_string());

        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        conn.execute(
            "INSERT OR REPLACE INTO documents (id, content, metadata) VALUES (?1, ?2, ?3)",
            rusqlite::params![parent_id, document.content, meta],
        )
        .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;

        // H5: replace the parent's prior chunk set first so a re-ingest that shrank /
        // changed segmentation leaves no stale `{parent}:chunk:{n}` rows behind
        // (INSERT OR REPLACE only overwrites equal ids; higher segment ids from a
        // previous, larger set would otherwise linger and stay retrievable).
        conn.execute(
            "DELETE FROM chunks WHERE parent_id = ?1",
            rusqlite::params![parent_id],
        )
        .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;

        let mut chunk_ids = Vec::new();
        for (i, text) in chunks_text.iter().enumerate() {
            let cid = format!("{}:chunk:{}", parent_id, i);
            conn.execute("INSERT OR REPLACE INTO chunks (chunk_id, parent_id, content, segment, metadata) VALUES (?1, ?2, ?3, ?4, ?5)",
                rusqlite::params![cid, parent_id, text, i, "{}"])
                .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
            chunk_ids.push(cid);
        }
        Ok((parent_id, chunk_ids))
    }

    async fn add_parent_documents(
        &self,
        documents: Vec<Document>,
        chunk_size: usize,
    ) -> Result<Vec<(String, Vec<String>)>, VectorStoreError> {
        let mut results = Vec::new();
        for doc in documents {
            results.push(self.add_parent_document(doc, chunk_size).await?);
        }
        Ok(results)
    }

    async fn get_parent_document(
        &self,
        parent_id: &str,
    ) -> Result<Option<Document>, VectorStoreError> {
        self.get_document(parent_id).await
    }

    async fn get_chunk(&self, chunk_id: &str) -> Result<Option<ChunkDocument>, VectorStoreError> {
        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let mut stmt = conn.prepare("SELECT chunk_id, parent_id, content, segment, metadata FROM chunks WHERE chunk_id = ?1")
            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let result = stmt.query_row(rusqlite::params![chunk_id], |row| {
            Ok(ChunkDocument {
                chunk_id: row.get(0)?,
                parent_id: row.get(1)?,
                content: row.get(2)?,
                segment: row.get(3)?,
                metadata: parse_metadata_or_default(&row.get::<_, String>(4)?),
            })
        });
        match result {
            Ok(chunk) => Ok(Some(chunk)),
            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
            Err(e) => Err(VectorStoreError::StorageError(e.to_string())),
        }
    }

    async fn get_chunk_document(
        &self,
        chunk_id: &str,
    ) -> Result<Option<Document>, VectorStoreError> {
        Ok(self.get_chunk(chunk_id).await?.map(|c| c.to_document()))
    }

    async fn get_chunks_for_parent(
        &self,
        parent_id: &str,
    ) -> Result<Vec<ChunkDocument>, VectorStoreError> {
        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let mut stmt = conn.prepare("SELECT chunk_id, parent_id, content, segment, metadata FROM chunks WHERE parent_id = ?1 ORDER BY segment")
            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let chunks = stmt
            .query_map(rusqlite::params![parent_id], |row| {
                Ok(ChunkDocument {
                    chunk_id: row.get(0)?,
                    parent_id: row.get(1)?,
                    content: row.get(2)?,
                    segment: row.get(3)?,
                    metadata: parse_metadata_or_default(&row.get::<_, String>(4)?),
                })
            })
            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?
            .filter_map(|r| match r {
                Ok(chunk) => Some(chunk),
                Err(e) => {
                    // stop silently dropping rows: log to surface storage degradation
                    log::error!(
                        "SQLite store: a row failed to deserialize and was dropped from the query result: {}",
                        e
                    );
                    None
                }
            })
            .collect();
        Ok(chunks)
    }

    async fn get_chunk_documents_for_parent(
        &self,
        parent_id: &str,
    ) -> Result<Vec<Document>, VectorStoreError> {
        Ok(self
            .get_chunks_for_parent(parent_id)
            .await?
            .into_iter()
            .map(|c| c.to_document())
            .collect())
    }

    async fn delete_parent_document(&self, parent_id: &str) -> Result<(), VectorStoreError> {
        self.delete_document(parent_id).await
    }

    async fn parent_count(&self) -> usize {
        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        match conn.query_row("SELECT COUNT(*) FROM documents", [], |r| r.get(0)) {
            Ok(count) => count,
            // M3: do not silently swallow errors and return 0; log an error to surface storage failures.
            Err(e) => {
                log::error!("SQLite parent_count query failed, returning 0: {}", e);
                0
            }
        }
    }

    async fn chunk_count(&self) -> usize {
        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        match conn.query_row("SELECT COUNT(*) FROM chunks", [], |r| r.get(0)) {
            Ok(count) => count,
            // M3: do not silently swallow errors and return 0; log an error to surface storage failures.
            Err(e) => {
                log::error!("SQLite chunk_count query failed, returning 0: {}", e);
                0
            }
        }
    }

    async fn get_all_chunks(&self) -> Result<Vec<ChunkDocument>, VectorStoreError> {
        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let mut stmt = conn
            .prepare("SELECT chunk_id, parent_id, content, segment, metadata FROM chunks")
            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let chunks = stmt
            .query_map([], |row| {
                Ok(ChunkDocument {
                    chunk_id: row.get(0)?,
                    parent_id: row.get(1)?,
                    content: row.get(2)?,
                    segment: row.get(3)?,
                    metadata: parse_metadata_or_default(&row.get::<_, String>(4)?),
                })
            })
            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?
            .filter_map(|r| match r {
                Ok(chunk) => Some(chunk),
                Err(e) => {
                    // stop silently dropping rows: log to surface storage degradation
                    log::error!(
                        "SQLite store: a row failed to deserialize and was dropped from the query result: {}",
                        e
                    );
                    None
                }
            })
            .collect();
        Ok(chunks)
    }

    async fn clear(&self) -> Result<(), VectorStoreError> {
        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        conn.execute_batch("DELETE FROM chunks; DELETE FROM documents;")
            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        Ok(())
    }

    // Blocking methods
    fn add_parent_document_blocking(
        &self,
        document: Document,
        chunk_size: usize,
    ) -> Result<(String, Vec<String>), VectorStoreError> {
        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let splitter = RecursiveCharacterSplitter::new(chunk_size, chunk_size / 10);
        let chunks_text = splitter.split_text(&document.content);
        let parent_id = document
            .id
            .clone()
            .unwrap_or_else(|| Uuid::new_v4().to_string());
        let meta = serde_json::to_string(&document.metadata).unwrap_or_else(|_| "{}".to_string());

        conn.execute(
            "INSERT OR REPLACE INTO documents (id, content, metadata) VALUES (?1, ?2, ?3)",
            rusqlite::params![parent_id, document.content, meta],
        )
        .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;

        // H5: drop the parent's prior chunk set so a shrunken re-ingest leaves none behind.
        conn.execute(
            "DELETE FROM chunks WHERE parent_id = ?1",
            rusqlite::params![parent_id],
        )
        .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;

        let mut chunk_ids = Vec::new();
        for (i, text) in chunks_text.iter().enumerate() {
            let cid = format!("{}:chunk:{}", parent_id, i);
            conn.execute("INSERT OR REPLACE INTO chunks (chunk_id, parent_id, content, segment, metadata) VALUES (?1, ?2, ?3, ?4, ?5)",
                rusqlite::params![cid, parent_id, text, i, "{}"])
                .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
            chunk_ids.push(cid);
        }
        Ok((parent_id, chunk_ids))
    }

    fn get_parent_document_blocking(
        &self,
        parent_id: &str,
    ) -> Result<Option<Document>, VectorStoreError> {
        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let mut stmt = conn
            .prepare("SELECT id, content, metadata FROM documents WHERE id = ?1")
            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let result = stmt.query_row(rusqlite::params![parent_id], |row| {
            Ok(Document {
                id: Some(row.get(0)?),
                content: row.get(1)?,
                metadata: parse_metadata_or_default(&row.get::<_, String>(2)?),
            })
        });
        match result {
            Ok(doc) => Ok(Some(doc)),
            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
            Err(e) => Err(VectorStoreError::StorageError(e.to_string())),
        }
    }

    fn get_chunk_blocking(
        &self,
        chunk_id: &str,
    ) -> Result<Option<ChunkDocument>, VectorStoreError> {
        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let mut stmt = conn.prepare("SELECT chunk_id, parent_id, content, segment, metadata FROM chunks WHERE chunk_id = ?1")
            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let result = stmt.query_row(rusqlite::params![chunk_id], |row| {
            Ok(ChunkDocument {
                chunk_id: row.get(0)?,
                parent_id: row.get(1)?,
                content: row.get(2)?,
                segment: row.get(3)?,
                metadata: parse_metadata_or_default(&row.get::<_, String>(4)?),
            })
        });
        match result {
            Ok(chunk) => Ok(Some(chunk)),
            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
            Err(e) => Err(VectorStoreError::StorageError(e.to_string())),
        }
    }

    fn blocking_get_chunks_for_parent(
        &self,
        parent_id: &str,
    ) -> Result<Vec<ChunkDocument>, VectorStoreError> {
        let conn = self.conn.lock().map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let mut stmt = conn.prepare("SELECT chunk_id, parent_id, content, segment, metadata FROM chunks WHERE parent_id = ?1 ORDER BY segment")
            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?;
        let chunks = stmt
            .query_map(rusqlite::params![parent_id], |row| {
                Ok(ChunkDocument {
                    chunk_id: row.get(0)?,
                    parent_id: row.get(1)?,
                    content: row.get(2)?,
                    segment: row.get(3)?,
                    metadata: parse_metadata_or_default(&row.get::<_, String>(4)?),
                })
            })
            .map_err(|e| VectorStoreError::StorageError(e.to_string()))?
            .filter_map(|r| match r {
                Ok(chunk) => Some(chunk),
                Err(e) => {
                    // stop silently dropping rows: log to surface storage degradation
                    log::error!(
                        "SQLite store: a row failed to deserialize and was dropped from the query result: {}",
                        e
                    );
                    None
                }
            })
            .collect();
        Ok(chunks)
    }
}