klieo_memory_pgvector/
client.rs

1//! PostgreSQL connection configuration + lazy table bootstrap.
2
3use crate::error::store_err;
4use dashmap::DashMap;
5use klieo_core::error::MemoryError;
6use sqlx_postgres::{PgPool, PgPoolOptions};
7use std::sync::Arc;
8use tokio::sync::OnceCell;
9
10const TABLE_VERSION_SUFFIX: &str = "_v1";
11const DEFAULT_TABLE_PREFIX: &str = "klieo_facts";
12const DEFAULT_MAX_CONNECTIONS: u32 = 5;
13const MAX_PREFIX_LEN: usize = 48;
14
15/// Configuration for [`MemoryPgvector`](crate::MemoryPgvector).
16///
17/// The connection `url` carries credentials and TLS mode
18/// (`postgres://user:pass@host/db?sslmode=require`) — TLS is the operator's
19/// responsibility via the URL, mirroring standard libpq behaviour.
20#[derive(Clone, Debug)]
21#[non_exhaustive]
22pub struct PgvectorConfig {
23    url: String,
24    table_prefix: String,
25    embedder_id: String,
26    max_connections: u32,
27}
28
29impl PgvectorConfig {
30    /// Build a config from a Postgres connection URL.
31    pub fn new(url: impl Into<String>) -> Self {
32        Self {
33            url: url.into(),
34            table_prefix: DEFAULT_TABLE_PREFIX.to_string(),
35            embedder_id: "default".to_string(),
36            max_connections: DEFAULT_MAX_CONNECTIONS,
37        }
38    }
39
40    /// Override the table-name prefix (default `klieo_facts`). Must be a valid
41    /// lowercase SQL identifier; rejected at connect time otherwise.
42    pub fn with_table_prefix(mut self, prefix: impl Into<String>) -> Self {
43        self.table_prefix = prefix.into();
44        self
45    }
46
47    /// Set the embedder identifier this store is provisioned for; surfaced via
48    /// `FilterableLongTermMemory::embedder_id()` so a composer can hard-fail on
49    /// cross-embedder queries rather than drift silently.
50    pub fn with_embedder_id(mut self, id: impl Into<String>) -> Self {
51        self.embedder_id = id.into();
52        self
53    }
54
55    /// Override the pool's maximum connection count (default 5).
56    pub fn with_max_connections(mut self, n: u32) -> Self {
57        self.max_connections = n;
58        self
59    }
60
61    /// Embedder identifier configured for this store.
62    pub fn embedder_id(&self) -> &str {
63        &self.embedder_id
64    }
65
66    pub(crate) fn url(&self) -> &str {
67        &self.url
68    }
69}
70
71/// Connected pool handle, cheap to clone. `ensure_table` is cached per-process
72/// per embedding dimension via a `DashMap<dim, OnceCell>` so the first
73/// `remember` provisions the table + indexes and subsequent calls short-circuit.
74#[derive(Clone)]
75pub(crate) struct PgvectorHandle {
76    pub(crate) pool: PgPool,
77    pub(crate) table: String,
78    bootstrapped: Arc<DashMap<u64, Arc<OnceCell<()>>>>,
79}
80
81impl PgvectorHandle {
82    pub(crate) async fn connect(cfg: &PgvectorConfig) -> Result<Self, MemoryError> {
83        let table = build_table_name(&cfg.table_prefix)?;
84        let pool = PgPoolOptions::new()
85            .max_connections(cfg.max_connections)
86            .connect(cfg.url())
87            .await
88            .map_err(store_err)?;
89        Ok(Self {
90            pool,
91            table,
92            bootstrapped: Arc::new(DashMap::new()),
93        })
94    }
95
96    pub(crate) async fn ensure_table(&self, vector_dim: u64) -> Result<(), MemoryError> {
97        let cell = self
98            .bootstrapped
99            .entry(vector_dim)
100            .or_insert_with(|| Arc::new(OnceCell::new()))
101            .clone();
102        cell.get_or_try_init(|| self.bootstrap_table(vector_dim))
103            .await?;
104        Ok(())
105    }
106
107    async fn bootstrap_table(&self, vector_dim: u64) -> Result<(), MemoryError> {
108        // `table` is a validated identifier and `vector_dim` an integer, so the
109        // DDL interpolation carries no untrusted input; all fact data flows
110        // through bound parameters on the read/write paths.
111        let table = &self.table;
112        let statements = [
113            "CREATE EXTENSION IF NOT EXISTS vector".to_string(),
114            format!(
115                "CREATE TABLE IF NOT EXISTS {table} (\
116                 fact_id uuid PRIMARY KEY, \
117                 text text NOT NULL, \
118                 metadata jsonb NOT NULL DEFAULT 'null'::jsonb, \
119                 embedding vector({vector_dim}) NOT NULL, \
120                 scope_kind text NOT NULL, \
121                 scope_value text NOT NULL)"
122            ),
123            format!(
124                "CREATE INDEX IF NOT EXISTS {table}_embedding_hnsw \
125                 ON {table} USING hnsw (embedding vector_cosine_ops)"
126            ),
127            format!(
128                "CREATE INDEX IF NOT EXISTS {table}_scope ON {table} (scope_kind, scope_value)"
129            ),
130        ];
131        for sql in statements {
132            sqlx_core::query::query(&sql)
133                .execute(&self.pool)
134                .await
135                .map_err(store_err)?;
136        }
137        // pgvector fixes the embedding column's dimension at creation, so a
138        // pre-existing table from a different embedder makes `CREATE TABLE IF NOT
139        // EXISTS` a silent no-op and inserts would fail at runtime. For a
140        // `vector(N)` column pgvector stores `N` in `atttypmod`; verify it and
141        // fail closed on a mismatch (-1 = unconstrained `vector`, left alone).
142        let actual_dim: Option<i32> = sqlx_core::query_scalar::query_scalar(
143            "SELECT a.atttypmod FROM pg_attribute a \
144             JOIN pg_class c ON a.attrelid = c.oid \
145             WHERE c.relname = $1 AND a.attname = 'embedding' AND NOT a.attisdropped",
146        )
147        .bind(table)
148        .fetch_optional(&self.pool)
149        .await
150        .map_err(store_err)?;
151        if let Some(actual) = actual_dim {
152            if actual > 0 && actual as u64 != vector_dim {
153                return Err(MemoryError::Embedding(format!(
154                    "pgvector table `{table}` has embedding dimension {actual}, but the \
155                     configured embedder produces {vector_dim}-dim vectors; a model change \
156                     needs a new table or a re-index"
157                )));
158            }
159        }
160        Ok(())
161    }
162}
163
164/// Build `{prefix}_v1`, rejecting a prefix that is not a safe lowercase SQL
165/// identifier (the prefix is interpolated into DDL, so it must not carry
166/// arbitrary SQL even though it is operator-supplied).
167fn build_table_name(prefix: &str) -> Result<String, MemoryError> {
168    if !is_valid_identifier(prefix) {
169        return Err(MemoryError::Store(format!(
170            "invalid table prefix `{prefix}`: expected a lowercase identifier \
171             matching [a-z_][a-z0-9_]* of at most {MAX_PREFIX_LEN} chars"
172        )));
173    }
174    Ok(format!("{prefix}{TABLE_VERSION_SUFFIX}"))
175}
176
177fn is_valid_identifier(s: &str) -> bool {
178    if s.is_empty() || s.len() > MAX_PREFIX_LEN {
179        return false;
180    }
181    let mut chars = s.chars();
182    let first_ok = chars
183        .next()
184        .map(|c| c.is_ascii_lowercase() || c == '_')
185        .unwrap_or(false);
186    first_ok
187        && s.chars()
188            .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_')
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194
195    #[test]
196    fn valid_prefix_builds_versioned_table_name() {
197        assert_eq!(build_table_name("klieo_facts").unwrap(), "klieo_facts_v1");
198        assert_eq!(build_table_name("triage").unwrap(), "triage_v1");
199    }
200
201    #[test]
202    fn injection_prefix_is_rejected() {
203        assert!(build_table_name("facts; DROP TABLE users").is_err());
204        assert!(build_table_name("Facts").is_err()); // uppercase
205        assert!(build_table_name("1facts").is_err()); // leading digit
206        assert!(build_table_name("").is_err());
207    }
208
209    #[test]
210    fn overlong_prefix_is_rejected() {
211        let long = "a".repeat(MAX_PREFIX_LEN + 1);
212        assert!(build_table_name(&long).is_err());
213    }
214}
klieo_memory_pgvector/client.rs

klieo_memory_pgvector/
client.rs
