Skip to content

SDK - Expose config for sqlx pool #1389

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 27, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 54 additions & 6 deletions pgml-sdks/pgml/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,12 @@
//!
//! With this SDK, you can seamlessly manage various database tables related to documents, text chunks, text splitters, LLM (Language Model) models, and embeddings. By leveraging the SDK's capabilities, you can efficiently index LLM embeddings using PgVector for fast and accurate queries.

use anyhow::Context;
use once_cell::sync::Lazy;
use parking_lot::RwLock;
use sqlx::{postgres::PgPoolOptions, PgPool};
use std::collections::HashMap;
use std::env;
use std::{collections::HashMap, time::Duration};
use tokio::runtime::{Builder, Runtime};
use tracing::Level;
use tracing_subscriber::FmtSubscriber;
Expand Down Expand Up @@ -67,13 +68,60 @@ async fn get_or_initialize_pool(database_url: &Option<String>) -> anyhow::Result
if let Some(pool) = pools.get(&url) {
Ok(pool.clone())
} else {
let timeout = std::env::var("PGML_CHECKOUT_TIMEOUT")
.unwrap_or_else(|_| "5000".to_string())
.parse::<u64>()
.expect("Error parsing PGML_CHECKOUT_TIMEOUT, expected an integer");
let acquire_timeout = std::env::var("PGML_CHECKOUT_TIMEOUT")
.ok()
.map(|v| v.parse::<u64>())
.transpose()
.context("Error parsing PGML_CHECKOUT_TIMEOUT, expected an integer")?
.map(anyhow::Ok)
.unwrap_or_else(|| {
Ok(std::env::var("PGML_POOL_ACQUIRE_TIMEOUT")
.ok()
.map(|v| v.parse::<u64>())
.transpose()
.context("Error parsing PGML_POOL_ACQUIRE_TIMEOUT, expected an integer")?
.unwrap_or(30000))
})?;
let acquire_timeout = Duration::from_millis(acquire_timeout);

let max_lifetime = std::env::var("PGML_POOL_MAX_LIFETIME")
.ok()
.map(|v| {
anyhow::Ok(Duration::from_millis(v.parse::<u64>().context(
"Error parsing PGML_POOL_MAX_LIFETIME, expected an integer",
)?))
})
.transpose()?;

let idle_timeout = std::env::var("PGML_POOL_IDLE_TIMEOUT")
.ok()
.map(|v| {
anyhow::Ok(Duration::from_millis(v.parse::<u64>().context(
"Error parsing PGML_POOL_IDLE_TIMEOUT, expected an integer",
)?))
})
.transpose()?;

let max_connections = std::env::var("PGML_POOL_MAX_CONNECTIONS")
.ok()
.map(|v| v.parse::<u32>())
.transpose()
.context("Error parsing PGML_POOL_MAX_CONNECTIONS, expected an integer")?
.unwrap_or(10);

let min_connections = std::env::var("PGML_POOL_MIN_CONNECTIONS")
.ok()
.map(|v| v.parse::<u32>())
.transpose()
.context("Error parsing PGML_POOL_MIN_CONNECTIONS, expected an integer")?
.unwrap_or(0);

let pool = PgPoolOptions::new()
.acquire_timeout(std::time::Duration::from_millis(timeout))
.max_connections(max_connections)
.min_connections(min_connections)
.acquire_timeout(acquire_timeout)
.max_lifetime(max_lifetime)
.idle_timeout(idle_timeout)
.connect_lazy(&url)?;

pools.insert(url.to_string(), pool.clone());
Expand Down
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy