use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use sqlx::postgres::PgPoolOptions;
use sqlx::PgPool;
mod chunks;
mod cleanup;
mod core;
mod embeddings;
mod encoding;
mod graph;
mod index_state;
mod migration;
pub mod migrations;
mod search;
#[cfg(test)]
mod tests;
#[derive(Clone)]
pub struct PostgresStore {
pub(crate) pool: PgPool,
pub(crate) vec_available: Arc<AtomicBool>,
}
impl PostgresStore {
pub async fn connect(url: &str) -> anyhow::Result<Self> {
let cfg = crate::config::DatabaseConfig::default();
let pool = Self::tuned_pool(url, &cfg).await?;
Self::from_pool(pool).await
}
async fn tuned_pool(url: &str, cfg: &crate::config::DatabaseConfig) -> anyhow::Result<PgPool> {
use std::time::Duration;
let session_sql = format!(
"SET statement_timeout = {}; SET lock_timeout = {}; \
SET idle_in_transaction_session_timeout = {}; SET work_mem = '{}'",
cfg.statement_timeout_ms,
cfg.lock_timeout_ms,
cfg.idle_in_transaction_timeout_ms,
cfg.work_mem,
);
let pool = PgPoolOptions::new()
.max_connections(cfg.pool_size as u32)
.acquire_timeout(Duration::from_millis(cfg.connection_timeout_ms))
.max_lifetime(Duration::from_secs(cfg.max_connection_lifetime_secs))
.idle_timeout(Duration::from_secs(cfg.idle_connection_timeout_secs))
.after_connect(move |conn, _meta| {
let session_sql = session_sql.clone();
Box::pin(async move {
use sqlx::Executor;
conn.execute(session_sql.as_str()).await?;
Ok(())
})
})
.connect(url)
.await?;
Ok(pool)
}
pub async fn from_pool(pool: PgPool) -> anyhow::Result<Self> {
migrations::run(pool.clone()).await?;
let vec_available = probe_vector_extension(&pool).await?;
Ok(Self {
pool,
vec_available: Arc::new(AtomicBool::new(vec_available)),
})
}
#[cfg(test)]
#[allow(dead_code)] pub(crate) fn with_vec_available(pool: PgPool, vec_available: bool) -> Self {
Self {
pool,
vec_available: Arc::new(AtomicBool::new(vec_available)),
}
}
pub fn has_vector_extension(&self) -> bool {
self.vec_available.load(Ordering::Relaxed)
}
}
async fn probe_vector_extension(pool: &PgPool) -> anyhow::Result<bool> {
let found: Option<i32> =
sqlx::query_scalar("SELECT 1 FROM pg_extension WHERE extname = 'vector'")
.fetch_optional(pool)
.await?;
Ok(found.is_some())
}
#[allow(dead_code)]
fn _assert_postgres_object_safe(s: &PostgresStore) {
let _: &dyn crate::db::traits::StoreCore = s;
let _: &dyn crate::db::traits::StoreChunks = s;
let _: &dyn crate::db::traits::StoreSearch = s;
let _: &dyn crate::db::traits::StoreGraph = s;
let _: &dyn crate::db::traits::StoreEmbeddings = s;
let _: &dyn crate::db::traits::StoreMigration = s;
let _: &dyn crate::db::traits::StoreCleanup = s;
let _: &dyn crate::db::traits::StoreIndexState = s;
let _: &dyn crate::db::traits::StoreEncoding = s;
let _: &(dyn crate::db::traits::Store + Send + Sync) = s;
}
#[allow(dead_code)]
fn _assert_arc_store_send_sync(s: std::sync::Arc<dyn crate::db::traits::Store + Send + Sync>) {
fn req<T: Send + Sync + 'static>(_: &T) {}
req(&s);
}
#[allow(dead_code)]
fn _assert_postgres_is_arc_store(s: PostgresStore) {
_assert_arc_store_send_sync(std::sync::Arc::new(s));
}