52 lines
1.7 KiB
Rust
52 lines
1.7 KiB
Rust
use std::time::Duration;
|
|
|
|
use diesel::pg::PgConnection;
|
|
use diesel::r2d2::{ConnectionManager, CustomizeConnection, Pool};
|
|
use diesel::RunQueryDsl;
|
|
|
|
pub type PgPool = Pool<ConnectionManager<PgConnection>>;
|
|
|
|
pub const DEFAULT_MAX_POOL_SIZE: u32 = 2;
|
|
|
|
#[derive(Debug)]
|
|
struct SchemaCustomizer;
|
|
|
|
impl CustomizeConnection<PgConnection, diesel::r2d2::Error> for SchemaCustomizer {
|
|
fn on_acquire(&self, conn: &mut PgConnection) -> Result<(), diesel::r2d2::Error> {
|
|
diesel::sql_query(
|
|
"SELECT set_config('search_path', (
|
|
SELECT string_agg(schema_name, ', ')
|
|
FROM (
|
|
SELECT 'tenant' AS schema_name WHERE EXISTS (
|
|
SELECT 1 FROM pg_namespace WHERE nspname = 'tenant'
|
|
)
|
|
UNION ALL
|
|
SELECT 'shared' AS schema_name WHERE EXISTS (
|
|
SELECT 1 FROM pg_namespace WHERE nspname = 'shared'
|
|
)
|
|
UNION ALL
|
|
SELECT 'public' AS schema_name
|
|
) AS schemas
|
|
), false)",
|
|
)
|
|
.execute(conn)
|
|
.map(|_| ())
|
|
.map_err(diesel::r2d2::Error::QueryError)
|
|
}
|
|
}
|
|
|
|
pub fn init_pool(database_url: &str) -> anyhow::Result<PgPool> {
|
|
init_pool_with_size(database_url, DEFAULT_MAX_POOL_SIZE)
|
|
}
|
|
|
|
pub fn init_pool_with_size(database_url: &str, max_size: u32) -> anyhow::Result<PgPool> {
|
|
let manager = ConnectionManager::<PgConnection>::new(database_url);
|
|
let pool_size = max_size.max(1);
|
|
let pool = Pool::builder()
|
|
.max_size(pool_size)
|
|
.connection_timeout(Duration::from_secs(10))
|
|
.connection_customizer(Box::new(SchemaCustomizer))
|
|
.build(manager)?;
|
|
Ok(pool)
|
|
}
|