lc_vector_stores/
pgvector.rs

1//! PGVector vector store (PostgreSQL + pgvector extension)
2//!
3//! This module provides helper utilities for PGVector. The full `PGVectorStore`
4//! implementation requires `sqlx` and `pgvector` crates, which must be added
5//! by the user to their own `Cargo.toml` due to potential conflicts with
6//! `rusqlite` (libsqlite3-sys linkage).
7//!
8//! To use `PGVectorStore`, add these to your `Cargo.toml`:
9//! ```toml
10//! sqlx = { version = "0.7", features = ["runtime-tokio", "postgres"] }
11//! pgvector = { version = "0.3", features = ["sqlx"] }
12//! ```
13//!
14//! Then enable the `pgvector-storage` feature and include the implementation
15//! from the project's `src/vector_stores/pgvector.rs`.
16
17use std::sync::LazyLock;
18
19use regex::Regex;
20
21use crate::VectorStoreError;
22
23static TABLE_NAME_RE: LazyLock<Regex> =
24    LazyLock::new(|| Regex::new(r"^[a-zA-Z_][a-zA-Z0-9_]*$").unwrap());
25
26/// Validate that a table name is safe for SQL interpolation.
27///
28/// Only allows: `^[a-zA-Z_][a-zA-Z0-9_]*$`
29/// This prevents SQL injection via table names.
30pub fn validate_table_name(table: &str) -> Result<(), VectorStoreError> {
31    if TABLE_NAME_RE.is_match(table) {
32        Ok(())
33    } else {
34        Err(VectorStoreError::ConfigError(format!(
35            "Invalid table name '{}': must match ^[a-zA-Z_][a-zA-Z0-9_]*$",
36            table
37        )))
38    }
39}
40
41/// Build CREATE TABLE SQL (pure function, convenient for testing)
42pub fn build_table_sql(table: &str, dim: usize) -> String {
43    format!(
44        "CREATE TABLE IF NOT EXISTS {} (id TEXT PRIMARY KEY, content TEXT, metadata JSONB, embedding vector({}))",
45        table, dim
46    )
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52
53    #[test]
54    fn test_build_table_sql() {
55        let sql = build_table_sql("docs", 1536);
56        assert!(sql.contains("CREATE TABLE"));
57        assert!(sql.contains("vector(1536)"));
58        assert!(sql.contains("docs"));
59    }
60
61    #[test]
62    fn test_build_table_sql_different_dim() {
63        let sql = build_table_sql("embeddings", 768);
64        assert!(sql.contains("vector(768)"));
65        assert!(sql.contains("embeddings"));
66    }
67
68    #[test]
69    fn test_build_table_sql_contains_metadata() {
70        let sql = build_table_sql("docs", 1536);
71        assert!(sql.contains("metadata JSONB"));
72        assert!(sql.contains("id TEXT PRIMARY KEY"));
73    }
74
75    #[test]
76    fn test_validate_table_name_valid() {
77        assert!(validate_table_name("users").is_ok());
78        assert!(validate_table_name("my_table").is_ok());
79        assert!(validate_table_name("_private").is_ok());
80        assert!(validate_table_name("Table123").is_ok());
81    }
82
83    #[test]
84    fn test_validate_table_name_invalid() {
85        // SQL injection attempts
86        assert!(validate_table_name("users; DROP TABLE users--").is_err());
87        assert!(validate_table_name("users; DROP TABLE users").is_err());
88        assert!(validate_table_name("123table").is_err()); // starts with digit
89        assert!(validate_table_name("user-table").is_err()); // contains hyphen
90        assert!(validate_table_name("user.table").is_err()); // contains dot
91        assert!(validate_table_name("").is_err()); // empty
92    }
93}
lc_vector_stores/pgvector.rs

lc_vector_stores/
pgvector.rs
