webserver_rs/web_core/
db.rs

1use crate::config::Database as DataBaseInfo;
2use sea_orm::{ConnectOptions, Database, DatabaseConnection, DbErr};
3use std::{collections::HashMap, sync::OnceLock};
4
5static DBCONN_POOL: OnceLock<HashMap<String, DatabaseConnection>> = OnceLock::new();
6
7async fn init_db_pool(db_info: Vec<DataBaseInfo>) -> Result<(), DbErr> {
8    let mut map = HashMap::new();
9    for db in db_info {
10        let opt = ConnectOptions::new(db.protocol)
11            .sqlx_logging(false)
12            .to_owned();
13        map.insert(db.name, Database::connect(opt).await?);
14    }
15    _ = DBCONN_POOL.set(map);
16    Ok(())
17}
18
19#[allow(dead_code)]
20#[cfg_attr(
21    docsrs,
22    doc(cfg(any(feature = "mysql", feature = "sqlite", feature = "postgres")))
23)]
24#[cfg(any(feature = "mysql", feature = "sqlite", feature = "postgres"))]
25/// Get database connection by indexing `name` of `[[database]]` set in `config.toml`
26pub fn get(key: &str) -> Result<&'static DatabaseConnection, std::io::Error> {
27    DBCONN_POOL
28        .get()
29        .ok_or(std::io::Error::new(
30            std::io::ErrorKind::NotFound,
31            "Database Connection Pool wasn't initialized when the program startup",
32        ))?
33        .get(key)
34        .ok_or(std::io::Error::new(
35            std::io::ErrorKind::NotFound,
36            format!("`{key}` is not found in Database Connection Pool"),
37        ))
38}
39
40#[allow(dead_code)]
41pub(crate) async fn init_db_if_enable(db_info: Vec<DataBaseInfo>) -> Result<(), DbErr> {
42    if cfg!(feature = "mysql") || cfg!(feature = "sqlite") || cfg!(feature = "postgres") {
43        init_db_pool(db_info).await
44    } else {
45        Ok(())
46    }
47}