klieo-memory-pgvector 3.1.0

//! PostgreSQL connection configuration + lazy table bootstrap.

use crate::error::store_err;
use dashmap::DashMap;
use klieo_core::error::MemoryError;
use sqlx_postgres::{PgPool, PgPoolOptions};
use std::sync::Arc;
use tokio::sync::OnceCell;

const TABLE_VERSION_SUFFIX: &str = "_v1";
const DEFAULT_TABLE_PREFIX: &str = "klieo_facts";
const DEFAULT_MAX_CONNECTIONS: u32 = 5;
const MAX_PREFIX_LEN: usize = 48;

/// Configuration for [`MemoryPgvector`](crate::MemoryPgvector).
///
/// The connection `url` carries credentials and TLS mode
/// (`postgres://user:pass@host/db?sslmode=require`) — TLS is the operator's
/// responsibility via the URL, mirroring standard libpq behaviour.
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct PgvectorConfig {
    url: String,
    table_prefix: String,
    embedder_id: String,
    max_connections: u32,
}

impl PgvectorConfig {
    /// Build a config from a Postgres connection URL.
    pub fn new(url: impl Into<String>) -> Self {
        Self {
            url: url.into(),
            table_prefix: DEFAULT_TABLE_PREFIX.to_string(),
            embedder_id: "default".to_string(),
            max_connections: DEFAULT_MAX_CONNECTIONS,
        }
    }

    /// Override the table-name prefix (default `klieo_facts`). Must be a valid
    /// lowercase SQL identifier; rejected at connect time otherwise.
    pub fn with_table_prefix(mut self, prefix: impl Into<String>) -> Self {
        self.table_prefix = prefix.into();
        self
    }

    /// Set the embedder identifier this store is provisioned for; surfaced via
    /// `FilterableLongTermMemory::embedder_id()` so a composer can hard-fail on
    /// cross-embedder queries rather than drift silently.
    pub fn with_embedder_id(mut self, id: impl Into<String>) -> Self {
        self.embedder_id = id.into();
        self
    }

    /// Override the pool's maximum connection count (default 5).
    pub fn with_max_connections(mut self, n: u32) -> Self {
        self.max_connections = n;
        self
    }

    /// Embedder identifier configured for this store.
    pub fn embedder_id(&self) -> &str {
        &self.embedder_id
    }

    pub(crate) fn url(&self) -> &str {
        &self.url
    }
}

/// Connected pool handle, cheap to clone. `ensure_table` is cached per-process
/// per embedding dimension via a `DashMap<dim, OnceCell>` so the first
/// `remember` provisions the table + indexes and subsequent calls short-circuit.
#[derive(Clone)]
pub(crate) struct PgvectorHandle {
    pub(crate) pool: PgPool,
    pub(crate) table: String,
    bootstrapped: Arc<DashMap<u64, Arc<OnceCell<()>>>>,
}

impl PgvectorHandle {
    pub(crate) async fn connect(cfg: &PgvectorConfig) -> Result<Self, MemoryError> {
        let table = build_table_name(&cfg.table_prefix)?;
        let pool = PgPoolOptions::new()
            .max_connections(cfg.max_connections)
            .connect(cfg.url())
            .await
            .map_err(store_err)?;
        Ok(Self {
            pool,
            table,
            bootstrapped: Arc::new(DashMap::new()),
        })
    }

    pub(crate) async fn ensure_table(&self, vector_dim: u64) -> Result<(), MemoryError> {
        let cell = self
            .bootstrapped
            .entry(vector_dim)
            .or_insert_with(|| Arc::new(OnceCell::new()))
            .clone();
        cell.get_or_try_init(|| self.bootstrap_table(vector_dim))
            .await?;
        Ok(())
    }

    async fn bootstrap_table(&self, vector_dim: u64) -> Result<(), MemoryError> {
        // `table` is a validated identifier and `vector_dim` an integer, so the
        // DDL interpolation carries no untrusted input; all fact data flows
        // through bound parameters on the read/write paths.
        let table = &self.table;
        let statements = [
            "CREATE EXTENSION IF NOT EXISTS vector".to_string(),
            format!(
                "CREATE TABLE IF NOT EXISTS {table} (\
                 fact_id uuid PRIMARY KEY, \
                 text text NOT NULL, \
                 metadata jsonb NOT NULL DEFAULT 'null'::jsonb, \
                 embedding vector({vector_dim}) NOT NULL, \
                 scope_kind text NOT NULL, \
                 scope_value text NOT NULL)"
            ),
            format!(
                "CREATE INDEX IF NOT EXISTS {table}_embedding_hnsw \
                 ON {table} USING hnsw (embedding vector_cosine_ops)"
            ),
            format!(
                "CREATE INDEX IF NOT EXISTS {table}_scope ON {table} (scope_kind, scope_value)"
            ),
        ];
        for sql in statements {
            sqlx_core::query::query(&sql)
                .execute(&self.pool)
                .await
                .map_err(store_err)?;
        }
        // pgvector fixes the embedding column's dimension at creation, so a
        // pre-existing table from a different embedder makes `CREATE TABLE IF NOT
        // EXISTS` a silent no-op and inserts would fail at runtime. For a
        // `vector(N)` column pgvector stores `N` in `atttypmod`; verify it and
        // fail closed on a mismatch (-1 = unconstrained `vector`, left alone).
        let actual_dim: Option<i32> = sqlx_core::query_scalar::query_scalar(
            "SELECT a.atttypmod FROM pg_attribute a \
             JOIN pg_class c ON a.attrelid = c.oid \
             WHERE c.relname = $1 AND a.attname = 'embedding' AND NOT a.attisdropped",
        )
        .bind(table)
        .fetch_optional(&self.pool)
        .await
        .map_err(store_err)?;
        if let Some(actual) = actual_dim {
            if actual > 0 && actual as u64 != vector_dim {
                return Err(MemoryError::Embedding(format!(
                    "pgvector table `{table}` has embedding dimension {actual}, but the \
                     configured embedder produces {vector_dim}-dim vectors; a model change \
                     needs a new table or a re-index"
                )));
            }
        }
        Ok(())
    }
}

/// Build `{prefix}_v1`, rejecting a prefix that is not a safe lowercase SQL
/// identifier (the prefix is interpolated into DDL, so it must not carry
/// arbitrary SQL even though it is operator-supplied).
fn build_table_name(prefix: &str) -> Result<String, MemoryError> {
    if !is_valid_identifier(prefix) {
        return Err(MemoryError::Store(format!(
            "invalid table prefix `{prefix}`: expected a lowercase identifier \
             matching [a-z_][a-z0-9_]* of at most {MAX_PREFIX_LEN} chars"
        )));
    }
    Ok(format!("{prefix}{TABLE_VERSION_SUFFIX}"))
}

fn is_valid_identifier(s: &str) -> bool {
    if s.is_empty() || s.len() > MAX_PREFIX_LEN {
        return false;
    }
    let mut chars = s.chars();
    let first_ok = chars
        .next()
        .map(|c| c.is_ascii_lowercase() || c == '_')
        .unwrap_or(false);
    first_ok
        && s.chars()
            .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_')
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn valid_prefix_builds_versioned_table_name() {
        assert_eq!(build_table_name("klieo_facts").unwrap(), "klieo_facts_v1");
        assert_eq!(build_table_name("triage").unwrap(), "triage_v1");
    }

    #[test]
    fn injection_prefix_is_rejected() {
        assert!(build_table_name("facts; DROP TABLE users").is_err());
        assert!(build_table_name("Facts").is_err()); // uppercase
        assert!(build_table_name("1facts").is_err()); // leading digit
        assert!(build_table_name("").is_err());
    }

    #[test]
    fn overlong_prefix_is_rejected() {
        let long = "a".repeat(MAX_PREFIX_LEN + 1);
        assert!(build_table_name(&long).is_err());
    }
}