swiftide_integrations/pgvector/
persist.rs

1//! Storage persistence implementation for vector embeddings.
2//!
3//! Implements the [`Persist`] trait for [`PgVector`], providing vector storage capabilities:
4//! - Database schema initialization and setup
5//! - Single-node storage operations
6//! - Optimized batch storage with configurable batch sizes
7//!
8//! NOTE: Persisting and retrieving metadata is not supported at the moment.
9//!
10//! The implementation ensures thread-safe concurrent access and handles
11//! connection management automatically.
12use crate::pgvector::PgVector;
13use anyhow::{Result, anyhow};
14use async_trait::async_trait;
15use sqlx::AssertSqlSafe;
16use std::sync::Arc;
17use swiftide_core::{
18    Persist,
19    indexing::{IndexingStream, TextNode},
20};
21
22#[async_trait]
23impl Persist for PgVector {
24    type Input = String;
25    type Output = String;
26    #[tracing::instrument(skip_all)]
27    async fn setup(&self) -> Result<()> {
28        // Get or initialize the connection pool
29        let pool = self.pool_get_or_initialize().await?;
30
31        if self.sql_stmt_bulk_insert.get().is_none() {
32            let sql = self.generate_unnest_upsert_sql()?;
33
34            self.sql_stmt_bulk_insert
35                .set(Arc::new(sql))
36                .map_err(|_| anyhow!("SQL bulk store statement is already set"))?;
37        }
38
39        let mut tx = pool.begin().await?;
40
41        // Create extension
42        let sql = "CREATE EXTENSION IF NOT EXISTS vector";
43        sqlx::query(sql).execute(&mut *tx).await?;
44
45        // Create table
46        let create_table_sql = self.generate_create_table_sql()?;
47        sqlx::query(AssertSqlSafe(create_table_sql))
48            .execute(&mut *tx)
49            .await?;
50
51        // Create HNSW index
52        let index_sql = self.create_index_sql()?;
53        sqlx::query(AssertSqlSafe(index_sql))
54            .execute(&mut *tx)
55            .await?;
56
57        tx.commit().await?;
58
59        Ok(())
60    }
61
62    #[tracing::instrument(skip_all)]
63    async fn store(&self, node: TextNode) -> Result<TextNode> {
64        let mut nodes = vec![node; 1];
65        self.store_nodes(&nodes).await?;
66
67        let node = nodes.swap_remove(0);
68
69        Ok(node)
70    }
71
72    #[tracing::instrument(skip_all)]
73    async fn batch_store(&self, nodes: Vec<TextNode>) -> IndexingStream<String> {
74        self.store_nodes(&nodes).await.map(|()| nodes).into()
75    }
76
77    fn batch_size(&self) -> Option<usize> {
78        Some(self.batch_size)
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use crate::pgvector::fixtures::TestContext;
85    use std::collections::HashSet;
86    use swiftide_core::{Persist, indexing::EmbeddedField};
87
88    #[test_log::test(tokio::test)]
89    async fn test_persist_setup_no_error_when_table_exists() {
90        let test_context = TestContext::setup_with_cfg(
91            vec!["filter"].into(),
92            HashSet::from([EmbeddedField::Combined]),
93        )
94        .await
95        .expect("Test setup failed");
96
97        test_context
98            .pgv_storage
99            .setup()
100            .await
101            .expect("PgVector setup should not fail when the table already exists");
102    }
103}
swiftide_integrations/pgvector/persist.rs

swiftide_integrations/pgvector/
persist.rs
