axum_kit/
postgres.rs

1use anyhow::Result;
2use serde::Deserialize;
3use sqlx::{postgres::PgPoolOptions, Executor, PgPool};
4use std::{sync::OnceLock, time::Duration};
5
6#[derive(Debug, Deserialize)]
7pub struct PostgresConfig {
8    pub url: String,
9    pub max_connections: u32,
10    pub min_connections: u32,
11    pub acquire_timeout: u64,
12    pub idle_timeout: u64,
13    pub max_lifetime: u64,
14}
15
16static PG_POOL: OnceLock<PgPool> = OnceLock::new();
17static PG_TIMEZONE: OnceLock<String> = OnceLock::new();
18
19pub async fn init(config: &PostgresConfig) -> Result<()> {
20    let tz = iana_time_zone::get_timezone().unwrap_or_else(|_| "UTC".to_string());
21    PG_TIMEZONE
22        .set(tz)
23        .map_err(|_| anyhow::anyhow!("Failed to set PG_TIMEZONE"))?;
24    let pool = PgPoolOptions::new()
25        .after_connect(|conn, _meta| {
26            Box::pin(async move {
27                conn.execute(format!("SET TIME ZONE '{}';", pg_session_timezone()).as_str())
28                    .await?;
29                Ok(())
30            })
31        })
32        .max_connections(config.max_connections)
33        .min_connections(config.min_connections)
34        .acquire_timeout(Duration::from_secs(config.acquire_timeout))
35        .idle_timeout(Some(Duration::from_secs(config.idle_timeout)))
36        .max_lifetime(Some(Duration::from_secs(config.max_lifetime)))
37        .connect(config.url.as_str())
38        .await?;
39    PG_POOL
40        .set(pool)
41        .map_err(|_| anyhow::anyhow!("Failed to set OnceLock<PgPool>"))
42}
43
44pub fn conn() -> &'static PgPool {
45    PG_POOL.get().expect("OnceLock<PgPool> not initialized")
46}
47
48pub fn pg_session_timezone() -> &'static str {
49    PG_TIMEZONE.get().expect("PG_TIMEZONE not initialized")
50}
axum_kit/postgres.rs

axum_kit/
postgres.rs
