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        Ok(())
139    }
140}
141
142/// Build `{prefix}_v1`, rejecting a prefix that is not a safe lowercase SQL
143/// identifier (the prefix is interpolated into DDL, so it must not carry
144/// arbitrary SQL even though it is operator-supplied).
145fn build_table_name(prefix: &str) -> Result<String, MemoryError> {
146    if !is_valid_identifier(prefix) {
147        return Err(MemoryError::Store(format!(
148            "invalid table prefix `{prefix}`: expected a lowercase identifier \
149             matching [a-z_][a-z0-9_]* of at most {MAX_PREFIX_LEN} chars"
150        )));
151    }
152    Ok(format!("{prefix}{TABLE_VERSION_SUFFIX}"))
153}
154
155fn is_valid_identifier(s: &str) -> bool {
156    if s.is_empty() || s.len() > MAX_PREFIX_LEN {
157        return false;
158    }
159    let mut chars = s.chars();
160    let first_ok = chars
161        .next()
162        .map(|c| c.is_ascii_lowercase() || c == '_')
163        .unwrap_or(false);
164    first_ok
165        && s.chars()
166            .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_')
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172
173    #[test]
174    fn valid_prefix_builds_versioned_table_name() {
175        assert_eq!(build_table_name("klieo_facts").unwrap(), "klieo_facts_v1");
176        assert_eq!(build_table_name("triage").unwrap(), "triage_v1");
177    }
178
179    #[test]
180    fn injection_prefix_is_rejected() {
181        assert!(build_table_name("facts; DROP TABLE users").is_err());
182        assert!(build_table_name("Facts").is_err()); // uppercase
183        assert!(build_table_name("1facts").is_err()); // leading digit
184        assert!(build_table_name("").is_err());
185    }
186
187    #[test]
188    fn overlong_prefix_is_rejected() {
189        let long = "a".repeat(MAX_PREFIX_LEN + 1);
190        assert!(build_table_name(&long).is_err());
191    }
192}
klieo_memory_pgvector/client.rs

klieo_memory_pgvector/
client.rs
