daimon_plugin_pgvector/
builder.rs

1//! Builder for [`PgVectorStore`].
2
3use daimon_core::{DaimonError, Result};
4use deadpool_postgres::{Config, Pool, Runtime};
5use tokio_postgres::NoTls;
6
7use crate::DistanceMetric;
8use crate::migrations;
9use crate::store::PgVectorStore;
10
11/// Builds a [`PgVectorStore`] with connection pooling and optional auto-migration.
12///
13/// # Example
14///
15/// ```ignore
16/// use daimon_plugin_pgvector::{PgVectorStoreBuilder, DistanceMetric};
17///
18/// let store = PgVectorStoreBuilder::new("host=localhost dbname=mydb", 1536)
19///     .table("embeddings")
20///     .distance_metric(DistanceMetric::Cosine)
21///     .hnsw_m(16)
22///     .hnsw_ef_construction(64)
23///     .auto_migrate(true)
24///     .build()
25///     .await?;
26/// ```
27pub struct PgVectorStoreBuilder {
28    connection_string: String,
29    dimensions: usize,
30    table: String,
31    distance_metric: DistanceMetric,
32    auto_migrate: bool,
33    hnsw_m: Option<usize>,
34    hnsw_ef_construction: Option<usize>,
35    pool_size: usize,
36}
37
38impl PgVectorStoreBuilder {
39    /// Creates a new builder.
40    ///
41    /// - `connection_string`: PostgreSQL connection string
42    ///   (e.g. `"host=localhost dbname=mydb user=postgres"` or
43    ///   `"postgresql://user:pass@host/db"`)
44    /// - `dimensions`: the fixed vector dimension count (must match your embedding model)
45    pub fn new(connection_string: impl Into<String>, dimensions: usize) -> Self {
46        Self {
47            connection_string: connection_string.into(),
48            dimensions,
49            table: "daimon_vectors".into(),
50            distance_metric: DistanceMetric::Cosine,
51            auto_migrate: true,
52            hnsw_m: None,
53            hnsw_ef_construction: None,
54            pool_size: 16,
55        }
56    }
57
58    /// Sets the table name. Default: `"daimon_vectors"`.
59    pub fn table(mut self, table: impl Into<String>) -> Self {
60        self.table = table.into();
61        self
62    }
63
64    /// Sets the distance metric. Default: [`DistanceMetric::Cosine`].
65    pub fn distance_metric(mut self, metric: DistanceMetric) -> Self {
66        self.distance_metric = metric;
67        self
68    }
69
70    /// Enables or disables automatic schema creation on first connection.
71    /// Default: `true`.
72    ///
73    /// When disabled, use the SQL from [`crate::migrations`] to set up
74    /// the schema manually.
75    pub fn auto_migrate(mut self, enabled: bool) -> Self {
76        self.auto_migrate = enabled;
77        self
78    }
79
80    /// Sets the HNSW `m` parameter (max connections per layer).
81    /// `None` uses the PostgreSQL default (16).
82    pub fn hnsw_m(mut self, m: usize) -> Self {
83        self.hnsw_m = Some(m);
84        self
85    }
86
87    /// Sets the HNSW `ef_construction` parameter (build-time search width).
88    /// `None` uses the PostgreSQL default (64).
89    pub fn hnsw_ef_construction(mut self, ef: usize) -> Self {
90        self.hnsw_ef_construction = Some(ef);
91        self
92    }
93
94    /// Sets the maximum number of connections in the pool. Default: `16`.
95    pub fn pool_size(mut self, size: usize) -> Self {
96        self.pool_size = size;
97        self
98    }
99
100    /// Builds the [`PgVectorStore`], optionally running migrations.
101    pub async fn build(self) -> Result<PgVectorStore> {
102        // Validate the table name *before* it is ever interpolated into SQL.
103        // PostgreSQL cannot bind identifiers as parameters, so the table name
104        // is formatted directly into every statement in `store.rs` and
105        // `migrations.rs`. An unvalidated name like `foo; DROP TABLE bar; --`
106        // would be a straightforward SQL injection. We only accept plain or
107        // schema-qualified identifiers matching `[A-Za-z_][A-Za-z0-9_]*`.
108        validate_table_name(&self.table)?;
109
110        let pool = self.create_pool()?;
111
112        if self.auto_migrate {
113            self.run_migrations(&pool).await?;
114        }
115
116        Ok(PgVectorStore {
117            pool,
118            table: self.table,
119            dimensions: self.dimensions,
120            distance_metric: self.distance_metric,
121        })
122    }
123
124    // (validation lives in the free `validate_table_name` fn below)
125
126    fn create_pool(&self) -> Result<Pool> {
127        let mut cfg = Config::new();
128        cfg.url = Some(self.connection_string.clone());
129        cfg.pool = Some(deadpool_postgres::PoolConfig {
130            max_size: self.pool_size,
131            ..Default::default()
132        });
133
134        cfg.create_pool(Some(Runtime::Tokio1), NoTls)
135            .map_err(|e| DaimonError::Other(format!("pgvector pool creation error: {e}")))
136    }
137
138    async fn run_migrations(&self, pool: &Pool) -> Result<()> {
139        let client = pool
140            .get()
141            .await
142            .map_err(|e| DaimonError::Other(format!("pgvector migration pool error: {e}")))?;
143
144        tracing::info!("pgvector: creating extension and table '{}'", self.table);
145
146        client
147            .execute(migrations::CREATE_EXTENSION, &[])
148            .await
149            .map_err(|e| DaimonError::Other(format!("pgvector CREATE EXTENSION error: {e}")))?;
150
151        let create_table = migrations::create_table_sql(&self.table, self.dimensions);
152        client
153            .execute(&create_table as &str, &[])
154            .await
155            .map_err(|e| DaimonError::Other(format!("pgvector CREATE TABLE error: {e}")))?;
156
157        let ops_class = self.distance_metric.ops_class();
158        let create_index = migrations::create_hnsw_index_sql(
159            &self.table,
160            ops_class,
161            self.hnsw_m,
162            self.hnsw_ef_construction,
163        );
164        client
165            .execute(&create_index as &str, &[])
166            .await
167            .map_err(|e| DaimonError::Other(format!("pgvector CREATE INDEX error: {e}")))?;
168
169        tracing::info!("pgvector: migration complete for '{}'", self.table);
170        Ok(())
171    }
172}
173
174/// Validates a (possibly schema-qualified) PostgreSQL table identifier.
175///
176/// Accepts a single identifier (`embeddings`) or a schema-qualified one
177/// (`public.embeddings`) where every dot-separated part matches
178/// `^[A-Za-z_][A-Za-z0-9_]*$`. Rejects anything containing quotes, whitespace,
179/// semicolons, or other punctuation that could break out of the identifier
180/// position in a formatted SQL statement.
181fn validate_table_name(table: &str) -> Result<()> {
182    fn is_valid_part(part: &str) -> bool {
183        let mut chars = part.chars();
184        match chars.next() {
185            Some(c) if c.is_ascii_alphabetic() || c == '_' => {}
186            _ => return false,
187        }
188        chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
189    }
190
191    if table.is_empty() {
192        return Err(DaimonError::Other(
193            "pgvector: table name must not be empty".to_string(),
194        ));
195    }
196
197    let parts: Vec<&str> = table.split('.').collect();
198    if parts.len() > 2 || !parts.iter().all(|p| is_valid_part(p)) {
199        return Err(DaimonError::Other(format!(
200            "pgvector: invalid table name '{table}': expected an identifier matching \
201             [A-Za-z_][A-Za-z0-9_]* (optionally schema-qualified as schema.table)"
202        )));
203    }
204
205    Ok(())
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    #[test]
213    fn test_valid_table_names() {
214        assert!(validate_table_name("embeddings").is_ok());
215        assert!(validate_table_name("daimon_vectors").is_ok());
216        assert!(validate_table_name("_private").is_ok());
217        assert!(validate_table_name("Table123").is_ok());
218        assert!(validate_table_name("public.embeddings").is_ok());
219    }
220
221    #[test]
222    fn test_invalid_table_names_rejected() {
223        assert!(validate_table_name("").is_err());
224        assert!(validate_table_name("foo; DROP TABLE bar").is_err());
225        assert!(validate_table_name("foo; DROP TABLE bar; --").is_err());
226        assert!(validate_table_name("\"foo\"").is_err());
227        assert!(validate_table_name("foo bar").is_err());
228        assert!(validate_table_name("1foo").is_err());
229        assert!(validate_table_name("foo.bar.baz").is_err());
230        assert!(validate_table_name("foo'").is_err());
231        assert!(validate_table_name("foo)").is_err());
232        assert!(validate_table_name("schema..table").is_err());
233    }
234}
daimon_plugin_pgvector/builder.rs

daimon_plugin_pgvector/
builder.rs
