lean_ctx/core/
pgvector_store.rs

1//! Optional pgvector (PostgreSQL) backend for dense (embedding) search.
2//!
3//! This module is behind the `pgvector` feature flag. It mirrors the
4//! `qdrant_store` API (namespaced per-project tables, md5-derived point ids,
5//! delete-by-file + upsert incremental sync) but talks to a self-hosted
6//! PostgreSQL instance with the `vector` extension installed.
7//!
8//! Deliberately dependency-free: SQL is executed through the `psql` CLI —
9//! the same pattern the postgres context provider uses — so no async runtime
10//! or native driver enters the dependency tree. Row output is read as
11//! one JSON object per line (`json_build_object(...)::text`), which is robust
12//! against arbitrary characters in file paths and symbol names.
13
14use std::collections::HashSet;
15use std::io::Write as _;
16use std::path::Path;
17
18use serde::Deserialize;
19
20use crate::core::bm25_index::{BM25Index, CodeChunk};
21
22#[derive(Debug, Clone)]
23pub struct PgvectorConfig {
24    /// PostgreSQL connection string (postgres://user:pass@host:port/db).
25    pub url: String,
26    /// Connect timeout for each psql invocation (seconds).
27    pub timeout_secs: u64,
28    /// Table name prefix; the project namespace hash and dimensions are appended.
29    pub table_prefix: String,
30}
31
32impl PgvectorConfig {
33    pub fn from_env() -> Result<Self, String> {
34        let url = std::env::var("LEANCTX_PGVECTOR_URL")
35            .map_err(|_| "LEANCTX_PGVECTOR_URL is required for pgvector backend".to_string())?;
36        let url = url.trim().to_string();
37        if url.is_empty() {
38            return Err("LEANCTX_PGVECTOR_URL is required for pgvector backend".to_string());
39        }
40
41        let timeout_secs = std::env::var("LEANCTX_PGVECTOR_TIMEOUT_SECS")
42            .ok()
43            .and_then(|v| v.trim().parse::<u64>().ok())
44            .filter(|v| *v > 0)
45            .unwrap_or(10);
46
47        let table_prefix = std::env::var("LEANCTX_PGVECTOR_TABLE_PREFIX")
48            .ok()
49            .map(|v| v.trim().to_string())
50            .filter(|v| !v.is_empty())
51            .unwrap_or_else(|| "lctx_code_".to_string());
52        validate_pg_identifier_prefix(&table_prefix)?;
53
54        Ok(Self {
55            url,
56            timeout_secs,
57            table_prefix,
58        })
59    }
60}
61
62#[derive(Debug, Clone)]
63pub struct PgvectorStore {
64    cfg: PgvectorConfig,
65}
66
67#[derive(Debug, Clone)]
68pub struct PgvectorHit {
69    pub score: f32,
70    pub file_path: String,
71    pub symbol_name: String,
72    pub kind: crate::core::bm25_index::ChunkKind,
73    pub start_line: usize,
74    pub end_line: usize,
75}
76
77#[derive(Debug, Deserialize)]
78struct PgRow {
79    score: f32,
80    file_path: String,
81    symbol_name: String,
82    kind: String,
83    start_line: usize,
84    end_line: usize,
85}
86
87impl PgvectorStore {
88    pub fn from_env() -> Result<Self, String> {
89        let cfg = PgvectorConfig::from_env()?;
90        Ok(Self { cfg })
91    }
92
93    /// Namespaced table name for a project root at given dimensionality.
94    /// Mirrors `QdrantStore::collection_name` (prefix + namespace hash + dims).
95    pub fn table_name(&self, root: &Path, dimensions: usize) -> Result<String, String> {
96        let ns = crate::core::index_namespace::namespace_hash(root);
97        let name = format!("{}{}_d{}", self.cfg.table_prefix, ns, dimensions);
98        if name.len() > 63 {
99            return Err(format!(
100                "pgvector table name exceeds PostgreSQL's 63-byte identifier limit: {name}"
101            ));
102        }
103        Ok(name)
104    }
105
106    /// Ensure the extension + table exist. Returns `true` if the table was created.
107    pub fn ensure_table(&self, table: &str, dimensions: usize) -> Result<bool, String> {
108        let existed = self.table_exists(table)?;
109        if existed {
110            return Ok(false);
111        }
112        let sql = format!(
113            "CREATE EXTENSION IF NOT EXISTS vector;\n\
114             CREATE TABLE IF NOT EXISTS {table} (\n\
115               id BIGINT PRIMARY KEY,\n\
116               file_path TEXT NOT NULL,\n\
117               symbol_name TEXT NOT NULL,\n\
118               kind TEXT NOT NULL,\n\
119               start_line BIGINT NOT NULL,\n\
120               end_line BIGINT NOT NULL,\n\
121               embedding vector({dimensions}) NOT NULL\n\
122             );\n\
123             CREATE INDEX IF NOT EXISTS {table}_file_idx ON {table} (file_path);"
124        );
125        self.run_sql(&sql)?;
126        Ok(true)
127    }
128
129    /// Same incremental semantics as the qdrant backend: fresh table gets a
130    /// full upsert; otherwise changed files are replaced (delete + upsert).
131    pub fn sync_index(
132        &self,
133        table: &str,
134        index: &BM25Index,
135        aligned_embeddings: &[Vec<f32>],
136        changed_files: &[String],
137        created_new: bool,
138    ) -> Result<(), String> {
139        if index.chunks.len() != aligned_embeddings.len() {
140            return Err("embedding alignment length mismatch".to_string());
141        }
142
143        if created_new {
144            return self.upsert_filtered(table, index, aligned_embeddings, None);
145        }
146
147        if changed_files.is_empty() {
148            return Ok(());
149        }
150
151        let mut unique: Vec<String> = changed_files.to_vec();
152        unique.sort();
153        unique.dedup();
154
155        for file in &unique {
156            self.delete_by_file(table, file)?;
157        }
158
159        let changed_set: HashSet<&str> = unique.iter().map(String::as_str).collect();
160        self.upsert_filtered(table, index, aligned_embeddings, Some(&changed_set))
161    }
162
163    pub fn search(
164        &self,
165        table: &str,
166        query_vec: &[f32],
167        limit: usize,
168    ) -> Result<Vec<PgvectorHit>, String> {
169        let vec_literal = vector_literal(query_vec);
170        // Cosine distance operator `<=>`: similarity = 1 - distance.
171        let sql = format!(
172            "SELECT json_build_object(\
173               'score', 1 - (embedding <=> '{vec_literal}'::vector), \
174               'file_path', file_path, \
175               'symbol_name', symbol_name, \
176               'kind', kind, \
177               'start_line', start_line, \
178               'end_line', end_line\
179             )::text \
180             FROM {table} \
181             ORDER BY embedding <=> '{vec_literal}'::vector \
182             LIMIT {limit};"
183        );
184        let stdout = self.run_sql(&sql)?;
185
186        let mut out = Vec::new();
187        for line in stdout.lines() {
188            let line = line.trim();
189            if line.is_empty() {
190                continue;
191            }
192            let row: PgRow = serde_json::from_str(line)
193                .map_err(|e| format!("invalid pgvector row json: {e}"))?;
194            out.push(PgvectorHit {
195                score: row.score,
196                file_path: row.file_path,
197                symbol_name: row.symbol_name,
198                kind: crate::core::dense_backend::kind_from_str(&row.kind),
199                start_line: row.start_line,
200                end_line: row.end_line,
201            });
202        }
203        Ok(out)
204    }
205
206    fn table_exists(&self, table: &str) -> Result<bool, String> {
207        let literal = sql_string_literal(table)?;
208        let out = self.run_sql(&format!("SELECT to_regclass({literal}) IS NOT NULL;"))?;
209        Ok(out.trim() == "t")
210    }
211
212    /// Upsert all chunks (or only those whose file is in `changed_set`),
213    /// batched to keep individual statements bounded.
214    fn upsert_filtered(
215        &self,
216        table: &str,
217        index: &BM25Index,
218        aligned_embeddings: &[Vec<f32>],
219        changed_set: Option<&HashSet<&str>>,
220    ) -> Result<(), String> {
221        let mut batch: Vec<String> = Vec::new();
222        for (i, chunk) in index.chunks.iter().enumerate() {
223            if let Some(set) = changed_set
224                && !set.contains(chunk.file_path.as_str())
225            {
226                continue;
227            }
228            let vec = aligned_embeddings
229                .get(i)
230                .ok_or_else(|| "embedding alignment missing".to_string())?;
231            batch.push(values_row_for_chunk(chunk, vec)?);
232            if batch.len() >= UPSERT_BATCH_ROWS {
233                self.upsert_rows(table, &batch)?;
234                batch.clear();
235            }
236        }
237        if !batch.is_empty() {
238            self.upsert_rows(table, &batch)?;
239        }
240        Ok(())
241    }
242
243    fn upsert_rows(&self, table: &str, rows: &[String]) -> Result<(), String> {
244        let sql = format!(
245            "INSERT INTO {table} (id, file_path, symbol_name, kind, start_line, end_line, embedding)\n\
246             VALUES\n{}\n\
247             ON CONFLICT (id) DO UPDATE SET\n\
248               file_path = EXCLUDED.file_path,\n\
249               symbol_name = EXCLUDED.symbol_name,\n\
250               kind = EXCLUDED.kind,\n\
251               start_line = EXCLUDED.start_line,\n\
252               end_line = EXCLUDED.end_line,\n\
253               embedding = EXCLUDED.embedding;",
254            rows.join(",\n")
255        );
256        self.run_sql(&sql).map(|_| ())
257    }
258
259    fn delete_by_file(&self, table: &str, file_path: &str) -> Result<(), String> {
260        let literal = sql_string_literal(file_path)?;
261        self.run_sql(&format!("DELETE FROM {table} WHERE file_path = {literal};"))
262            .map(|_| ())
263    }
264
265    /// Run SQL through `psql` via a temp file (`-f`) so statement size is not
266    /// limited by ARG_MAX. Returns stdout (tuples-only, unaligned).
267    fn run_sql(&self, sql: &str) -> Result<String, String> {
268        let mut tmp = tempfile::NamedTempFile::new()
269            .map_err(|e| format!("pgvector: temp file failed: {e}"))?;
270        tmp.write_all(sql.as_bytes())
271            .map_err(|e| format!("pgvector: temp write failed: {e}"))?;
272        tmp.flush()
273            .map_err(|e| format!("pgvector: temp flush failed: {e}"))?;
274
275        let output = std::process::Command::new("psql")
276            .arg(&self.cfg.url)
277            .args(["-X", "-q", "-v", "ON_ERROR_STOP=1", "-t", "-A", "-f"])
278            .arg(tmp.path())
279            .env("PGCONNECT_TIMEOUT", self.cfg.timeout_secs.to_string())
280            .output()
281            .map_err(|e| {
282                format!("pgvector: failed to run psql (is the PostgreSQL client installed?): {e}")
283            })?;
284
285        if !output.status.success() {
286            let stderr = String::from_utf8_lossy(&output.stderr);
287            return Err(format!("pgvector: psql error: {}", stderr.trim()));
288        }
289        Ok(String::from_utf8_lossy(&output.stdout).into_owned())
290    }
291}
292
293const UPSERT_BATCH_ROWS: usize = 256;
294
295/// One `(id, 'file', 'symbol', 'kind', start, end, '[...]'::vector)` row.
296fn values_row_for_chunk(chunk: &CodeChunk, vector: &[f32]) -> Result<String, String> {
297    let id = point_id_for_chunk(chunk) as i64; // BIGINT: deterministic wrap of the u64 hash
298    let file = sql_string_literal(&chunk.file_path)?;
299    let symbol = sql_string_literal(&chunk.symbol_name)?;
300    let kind = sql_string_literal(crate::core::dense_backend::kind_to_str(&chunk.kind))?;
301    Ok(format!(
302        "({id}, {file}, {symbol}, {kind}, {}, {}, '{}'::vector)",
303        chunk.start_line,
304        chunk.end_line,
305        vector_literal(vector)
306    ))
307}
308
309/// Identical id scheme to `qdrant_store::point_id_for_chunk` so a project can
310/// switch backends without changing point identity semantics.
311fn point_id_for_chunk(chunk: &CodeChunk) -> u64 {
312    use md5::{Digest, Md5};
313    let mut h = Md5::new();
314    h.update(chunk.file_path.as_bytes());
315    h.update(chunk.start_line.to_le_bytes());
316    h.update(chunk.end_line.to_le_bytes());
317    h.update(chunk.symbol_name.as_bytes());
318    h.update(crate::core::dense_backend::kind_to_str(&chunk.kind).as_bytes());
319    let out = h.finalize();
320    u64::from_le_bytes(out[0..8].try_into().unwrap_or([0u8; 8]))
321}
322
323/// pgvector input format: `[0.1,0.2,...]`.
324fn vector_literal(vector: &[f32]) -> String {
325    let mut s = String::with_capacity(vector.len() * 10 + 2);
326    s.push('[');
327    for (i, v) in vector.iter().enumerate() {
328        if i > 0 {
329            s.push(',');
330        }
331        // `{v}` for f32 is locale-independent and round-trips through Postgres real parsing.
332        s.push_str(&format!("{v}"));
333    }
334    s.push(']');
335    s
336}
337
338/// Standard SQL single-quoted literal with `'` doubling. Rejects NUL bytes
339/// (PostgreSQL cannot store them in TEXT anyway) and backslashes are safe
340/// because standard_conforming_strings is on by default since PostgreSQL 9.1.
341fn sql_string_literal(s: &str) -> Result<String, String> {
342    if s.contains('\0') {
343        return Err("pgvector: NUL byte in string".to_string());
344    }
345    Ok(format!("'{}'", s.replace('\'', "''")))
346}
347
348/// The table prefix is interpolated into SQL identifiers; enforce the same
349/// whitelist the postgres provider uses for schema names.
350fn validate_pg_identifier_prefix(name: &str) -> Result<(), String> {
351    let valid_start = name
352        .chars()
353        .next()
354        .is_some_and(|c| c.is_ascii_alphabetic() || c == '_');
355    let valid_rest = name
356        .chars()
357        .skip(1)
358        .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '$');
359    if name.is_empty() || name.len() > 40 || !valid_start || !valid_rest {
360        return Err(format!(
361            "Invalid LEANCTX_PGVECTOR_TABLE_PREFIX: {name:?} (allowed: [A-Za-z_][A-Za-z0-9_$]*, max 40 chars)"
362        ));
363    }
364    Ok(())
365}
366
367#[cfg(test)]
368mod tests {
369    use super::*;
370    use crate::core::bm25_index::ChunkKind;
371
372    fn chunk(file: &str, name: &str, start: usize, end: usize, kind: ChunkKind) -> CodeChunk {
373        CodeChunk {
374            file_path: file.to_string(),
375            symbol_name: name.to_string(),
376            kind,
377            start_line: start,
378            end_line: end,
379            content: "fn x() {}".to_string(),
380            tokens: vec![],
381            token_count: 0,
382        }
383    }
384
385    #[test]
386    fn point_id_matches_qdrant_scheme_and_is_stable() {
387        let c = chunk("src/main.rs", "main", 1, 10, ChunkKind::Function);
388        assert_eq!(point_id_for_chunk(&c), point_id_for_chunk(&c));
389        let c2 = chunk("src/main.rs", "main", 2, 10, ChunkKind::Function);
390        assert_ne!(point_id_for_chunk(&c), point_id_for_chunk(&c2));
391    }
392
393    #[test]
394    fn sql_string_literal_escapes_quotes() {
395        assert_eq!(sql_string_literal("a'b").unwrap(), "'a''b'");
396        assert_eq!(sql_string_literal("plain").unwrap(), "'plain'");
397        assert!(sql_string_literal("nul\0byte").is_err());
398    }
399
400    #[test]
401    fn vector_literal_is_bracketed_csv() {
402        assert_eq!(vector_literal(&[0.5, -1.0, 2.0]), "[0.5,-1,2]");
403        assert_eq!(vector_literal(&[]), "[]");
404    }
405
406    #[test]
407    fn values_row_contains_escaped_fields() {
408        let c = chunk("src/a'b.rs", "fn'x", 3, 9, ChunkKind::Method);
409        let row = values_row_for_chunk(&c, &[0.25, 0.75]).unwrap();
410        assert!(row.contains("'src/a''b.rs'"));
411        assert!(row.contains("'fn''x'"));
412        assert!(row.contains("'Method'"));
413        assert!(row.contains("'[0.25,0.75]'::vector"));
414    }
415
416    #[test]
417    fn table_prefix_validation() {
418        assert!(validate_pg_identifier_prefix("lctx_code_").is_ok());
419        assert!(validate_pg_identifier_prefix("with space").is_err());
420        assert!(validate_pg_identifier_prefix("1leading_digit").is_err());
421        assert!(validate_pg_identifier_prefix("drop;table").is_err());
422        assert!(validate_pg_identifier_prefix("").is_err());
423    }
424
425    #[test]
426    fn config_requires_url() {
427        let _env = crate::core::data_dir::test_env_lock();
428        crate::test_env::remove_var("LEANCTX_PGVECTOR_URL");
429        assert!(PgvectorConfig::from_env().is_err());
430
431        crate::test_env::set_var("LEANCTX_PGVECTOR_URL", "postgres://localhost/lctx");
432        crate::test_env::remove_var("LEANCTX_PGVECTOR_TABLE_PREFIX");
433        crate::test_env::remove_var("LEANCTX_PGVECTOR_TIMEOUT_SECS");
434        let cfg = PgvectorConfig::from_env().unwrap();
435        assert_eq!(cfg.url, "postgres://localhost/lctx");
436        assert_eq!(cfg.table_prefix, "lctx_code_");
437        assert_eq!(cfg.timeout_secs, 10);
438        crate::test_env::remove_var("LEANCTX_PGVECTOR_URL");
439    }
440
441    /// Real round-trip against a live PostgreSQL+pgvector instance.
442    /// Run explicitly (needs LEANCTX_PGVECTOR_URL + psql client on PATH):
443    ///   LEANCTX_PGVECTOR_URL=postgres://... cargo test --lib pgvector_e2e -- --ignored
444    #[test]
445    #[ignore = "requires live PostgreSQL with pgvector extension (set LEANCTX_PGVECTOR_URL)"]
446    fn pgvector_e2e_round_trip() {
447        let store = PgvectorStore::from_env().expect("LEANCTX_PGVECTOR_URL must be set");
448        let table = "lctx_e2e_round_trip_d3".to_string();
449        let _ = store.run_sql(&format!("DROP TABLE IF EXISTS {table};"));
450
451        // Fresh table + full upsert.
452        assert!(
453            store.ensure_table(&table, 3).unwrap(),
454            "table should be new"
455        );
456        assert!(
457            !store.ensure_table(&table, 3).unwrap(),
458            "second call sees it"
459        );
460
461        let mut index = BM25Index::new();
462        index
463            .chunks
464            .push(chunk("src/a.rs", "alpha", 1, 5, ChunkKind::Function));
465        index
466            .chunks
467            .push(chunk("src/b.rs", "beta", 10, 20, ChunkKind::Struct));
468        let embeddings = vec![vec![1.0, 0.0, 0.0], vec![0.0, 1.0, 0.0]];
469
470        store
471            .sync_index(&table, &index, &embeddings, &[], true)
472            .unwrap();
473
474        let hits = store.search(&table, &[1.0, 0.0, 0.0], 2).unwrap();
475        assert_eq!(hits.len(), 2);
476        assert_eq!(hits[0].file_path, "src/a.rs");
477        assert_eq!(hits[0].symbol_name, "alpha");
478        assert!(hits[0].score > 0.99, "cosine sim of identical vec ~ 1.0");
479        assert_eq!(hits[0].kind, ChunkKind::Function);
480
481        // Incremental: b.rs changes (chunk moves), delete-by-file + re-upsert.
482        index.chunks[1] = chunk("src/b.rs", "beta", 30, 40, ChunkKind::Struct);
483        store
484            .sync_index(
485                &table,
486                &index,
487                &embeddings,
488                &["src/b.rs".to_string()],
489                false,
490            )
491            .unwrap();
492
493        let hits = store.search(&table, &[0.0, 1.0, 0.0], 2).unwrap();
494        assert_eq!(hits[0].file_path, "src/b.rs");
495        assert_eq!(hits[0].start_line, 30, "stale row was replaced");
496        assert_eq!(hits[0].end_line, 40);
497
498        // Escaping survives the round trip.
499        index
500            .chunks
501            .push(chunk("src/it's.rs", "q'uote", 2, 3, ChunkKind::Method));
502        let embeddings = vec![
503            vec![1.0, 0.0, 0.0],
504            vec![0.0, 1.0, 0.0],
505            vec![0.0, 0.0, 1.0],
506        ];
507        store
508            .sync_index(
509                &table,
510                &index,
511                &embeddings,
512                &["src/it's.rs".to_string()],
513                false,
514            )
515            .unwrap();
516        let hits = store.search(&table, &[0.0, 0.0, 1.0], 1).unwrap();
517        assert_eq!(hits[0].file_path, "src/it's.rs");
518        assert_eq!(hits[0].symbol_name, "q'uote");
519
520        store.run_sql(&format!("DROP TABLE {table};")).unwrap();
521    }
522
523    #[test]
524    fn table_name_is_namespaced_and_bounded() {
525        let _env = crate::core::data_dir::test_env_lock();
526        crate::test_env::set_var("LEANCTX_PGVECTOR_URL", "postgres://localhost/lctx");
527        let store = PgvectorStore::from_env().unwrap();
528        let name = store
529            .table_name(Path::new("/tmp/some-project"), 384)
530            .unwrap();
531        assert!(name.starts_with("lctx_code_"));
532        assert!(name.ends_with("_d384"));
533        assert!(name.len() <= 63);
534        crate::test_env::remove_var("LEANCTX_PGVECTOR_URL");
535    }
536}
lean_ctx/core/pgvector_store.rs

lean_ctx/core/
pgvector_store.rs
