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

klieo_memory_pgvector/
client.rs
