wowlab_sentinel/
db.rs

1use std::time::Duration;
2
3use sqlx::{PgConnection, PgPool, postgres::PgPoolOptions};
4use wowlab_types::sim::FastMap;
5
6use crate::{Config, ExposeSecret};
7
8pub(crate) trait DbClient {
9    const NAME: &'static str;
10}
11
12#[derive(Debug)]
13pub(crate) struct DbClientSpec {
14    pub name: &'static str,
15    pub max_connections: u32,
16    pub statement_timeout: Option<Duration>,
17}
18
19inventory::collect!(DbClientSpec);
20
21/// Declare a database client: a named pool with its own size and optional statement timeout.
22macro_rules! db_client {
23    ($marker:ident, $name:literal, max = $max:expr) => {
24        $crate::db_client!(@build $marker, $name, $max, ::core::option::Option::None);
25    };
26    ($marker:ident, $name:literal, max = $max:expr, statement_timeout_secs = $secs:expr) => {
27        $crate::db_client!(
28            @build $marker, $name, $max,
29            ::core::option::Option::Some(::core::time::Duration::from_secs($secs))
30        );
31    };
32    (@build $marker:ident, $name:literal, $max:expr, $timeout:expr) => {
33        pub(crate) struct $marker;
34
35        impl $crate::db::DbClient for $marker {
36            const NAME: &'static str = $name;
37        }
38
39        inventory::submit! {
40            $crate::db::DbClientSpec {
41                name: $name,
42                max_connections: $max,
43                statement_timeout: $timeout,
44            }
45        }
46    };
47}
48
49pub(crate) use db_client;
50
51async fn set_statement_timeout(
52    conn: &mut PgConnection,
53    timeout: Duration,
54) -> Result<(), sqlx::Error> {
55    let ms = u64::try_from(timeout.as_millis()).unwrap_or(u64::MAX);
56
57    sqlx::query(&format!("SET statement_timeout = {ms}"))
58        .execute(conn)
59        .await?;
60
61    Ok(())
62}
63
64#[derive(Debug)]
65pub(crate) struct DbRegistry {
66    pools: FastMap<&'static str, PgPool>,
67}
68
69impl DbRegistry {
70    pub(crate) async fn connect(config: &Config) -> Result<Self, sqlx::Error> {
71        let url = config.database_url.expose_secret();
72        let acquire_timeout = Duration::from_secs(config.db_acquire_timeout_secs);
73        let idle_timeout = Duration::from_secs(config.db_idle_timeout_secs);
74
75        let mut pools = FastMap::default();
76        let mut connected = Vec::new();
77
78        for spec in inventory::iter::<DbClientSpec> {
79            let mut options = PgPoolOptions::new()
80                .max_connections(spec.max_connections)
81                .acquire_timeout(acquire_timeout)
82                .idle_timeout(idle_timeout);
83
84            if let Some(timeout) = spec.statement_timeout {
85                options = options.after_connect(move |conn, _meta| {
86                    Box::pin(set_statement_timeout(conn, timeout))
87                });
88            }
89
90            pools.insert(spec.name, options.connect(url).await?);
91            connected.push((spec.name, spec.max_connections, spec.statement_timeout));
92        }
93
94        tracing::info!(?connected, "Connected database clients");
95
96        Ok(Self { pools })
97    }
98
99    pub(crate) fn get<C>(&self) -> &PgPool
100    where
101        C: DbClient,
102    {
103        self.pools
104            .get(C::NAME)
105            .expect("db client must be registered at startup")
106    }
107
108    pub(crate) async fn close(&self) {
109        for pool in self.pools.values() {
110            pool.close().await;
111        }
112    }
113}
wowlab_sentinel/db.rs

wowlab_sentinel/
db.rs
