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
21static TABLE_NAME_RE: LazyLock<Regex> =
22    LazyLock::new(|| Regex::new(r"^[a-zA-Z_][a-zA-Z0-9_]*$").unwrap());
23
24/// Validate that a table name is safe for SQL interpolation.
25///
26/// Only allows: `^[a-zA-Z_][a-zA-Z0-9_]*$`
27/// This prevents SQL injection via table names.
28pub fn validate_table_name(table: &str) -> Result<(), String> {
29    if TABLE_NAME_RE.is_match(table) {
30        Ok(())
31    } else {
32        Err(format!(
33            "Invalid table name '{}': must match ^[a-zA-Z_][a-zA-Z0-9_]*$",
34            table
35        ))
36    }
37}
38
39/// Build CREATE TABLE SQL (pure function, convenient for testing)
40pub fn build_table_sql(table: &str, dim: usize) -> String {
41    format!(
42        "CREATE TABLE IF NOT EXISTS {} (id TEXT PRIMARY KEY, content TEXT, metadata JSONB, embedding vector({}))",
43        table, dim
44    )
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50
51    #[test]
52    fn test_build_table_sql() {
53        let sql = build_table_sql("docs", 1536);
54        assert!(sql.contains("CREATE TABLE"));
55        assert!(sql.contains("vector(1536)"));
56        assert!(sql.contains("docs"));
57    }
58
59    #[test]
60    fn test_build_table_sql_different_dim() {
61        let sql = build_table_sql("embeddings", 768);
62        assert!(sql.contains("vector(768)"));
63        assert!(sql.contains("embeddings"));
64    }
65
66    #[test]
67    fn test_build_table_sql_contains_metadata() {
68        let sql = build_table_sql("docs", 1536);
69        assert!(sql.contains("metadata JSONB"));
70        assert!(sql.contains("id TEXT PRIMARY KEY"));
71    }
72
73    #[test]
74    fn test_validate_table_name_valid() {
75        assert!(validate_table_name("users").is_ok());
76        assert!(validate_table_name("my_table").is_ok());
77        assert!(validate_table_name("_private").is_ok());
78        assert!(validate_table_name("Table123").is_ok());
79    }
80
81    #[test]
82    fn test_validate_table_name_invalid() {
83        // SQL injection attempts
84        assert!(validate_table_name("users; DROP TABLE users--").is_err());
85        assert!(validate_table_name("users; DROP TABLE users").is_err());
86        assert!(validate_table_name("123table").is_err()); // starts with digit
87        assert!(validate_table_name("user-table").is_err()); // contains hyphen
88        assert!(validate_table_name("user.table").is_err()); // contains dot
89        assert!(validate_table_name("").is_err()); // empty
90    }
91}
lc_vector_stores/pgvector.rs

lc_vector_stores/
pgvector.rs
