tibba_sql/
lib.rs

1// Copyright 2025 Tree xie.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use serde::Deserialize;
16use snafu::Snafu;
17use sqlx::MySqlPool;
18use sqlx::pool::PoolOptions;
19use std::sync::Arc;
20use std::sync::atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering};
21use std::time::Duration;
22use tibba_config::Config;
23use tibba_error::Error as BaseError;
24use tibba_util::parse_uri;
25use tracing::info;
26use validator::Validate;
27
28#[derive(Debug, Clone, Default, Validate)]
29pub struct DatabaseConfig {
30    pub origin_url: String,
31    #[validate(length(min = 10))]
32    pub url: String,
33    #[validate(range(min = 2, max = 1000))]
34    pub max_connections: u32,
35    #[validate(range(min = 0, max = 10))]
36    pub min_connections: u32,
37    pub connect_timeout: Duration,
38    pub idle_timeout: Duration,
39    pub max_lifetime: Duration,
40    pub test_before_acquire: bool,
41    pub password: Option<String>,
42}
43
44fn default_max_connections() -> u32 {
45    10
46}
47fn default_min_connections() -> u32 {
48    2
49}
50
51#[derive(Deserialize, Debug, Clone)]
52struct DatabaseQuery {
53    #[serde(default = "default_max_connections")]
54    pub max_connections: u32,
55    #[serde(default = "default_min_connections")]
56    pub min_connections: u32,
57    #[serde(default)]
58    #[serde(with = "humantime_serde")]
59    pub connect_timeout: Option<Duration>,
60    #[serde(default)]
61    #[serde(with = "humantime_serde")]
62    pub idle_timeout: Option<Duration>,
63    #[serde(default)]
64    #[serde(with = "humantime_serde")]
65    pub max_lifetime: Option<Duration>,
66    pub test_before_acquire: Option<bool>,
67}
68
69// Creates a new DatabaseConfig instance from the configuration
70fn new_database_config(config: &Config) -> Result<DatabaseConfig> {
71    let origin_url = config.get_string("uri").map_err(|e| Error::Common {
72        category: "config".to_string(),
73        message: e.to_string(),
74    })?;
75    let url = origin_url.clone();
76    let parsed = parse_uri::<DatabaseQuery>(&url).map_err(|e| Error::Common {
77        category: "config".to_string(),
78        message: e.to_string(),
79    })?;
80
81    let mut url = parsed.url().map_err(|e| Error::Common {
82        category: "config".to_string(),
83        message: e.to_string(),
84    })?;
85    url.set_query(None);
86
87    let query = &parsed.query;
88    let database_config = DatabaseConfig {
89        origin_url,
90        url: url.to_string(),
91        max_connections: query.max_connections,
92        min_connections: query.min_connections,
93        connect_timeout: query.connect_timeout.unwrap_or(Duration::from_secs(3)),
94        idle_timeout: query.idle_timeout.unwrap_or(Duration::from_secs(60)),
95        max_lifetime: query
96            .max_lifetime
97            .unwrap_or(Duration::from_secs(6 * 60 * 60)),
98        test_before_acquire: query.test_before_acquire.unwrap_or(true),
99        password: parsed.password.map(|v| v.to_string()),
100    };
101    database_config
102        .validate()
103        .map_err(|e| Error::Validate { source: e })?;
104    Ok(database_config)
105}
106
107#[derive(Debug, Snafu)]
108pub enum Error {
109    #[snafu(display("sqlx error: {source}"))]
110    Sqlx { source: sqlx::Error },
111    #[snafu(display("validate error: {source}"))]
112    Validate { source: validator::ValidationErrors },
113    #[snafu(display("category: {category}, error: {message}"))]
114    Common { category: String, message: String },
115}
116
117impl From<Error> for BaseError {
118    fn from(source: Error) -> Self {
119        let err = match source {
120            Error::Sqlx { source } => BaseError::new(source)
121                .with_sub_category("sqlx")
122                .with_exception(true),
123            Error::Validate { source } => BaseError::new(source)
124                .with_sub_category("validate")
125                .with_exception(true),
126            Error::Common { message, .. } => BaseError::new(message).with_exception(true),
127        };
128        err.with_category("sql")
129    }
130}
131
132#[derive(Debug, Default)]
133pub struct PoolStat {
134    connected: AtomicU32,
135    executions: AtomicUsize,
136    idle_for: AtomicU64,
137}
138
139impl PoolStat {
140    pub fn stat(&self) -> (u32, usize, u64) {
141        let connected = self.connected.swap(0, Ordering::Relaxed);
142        let executions = self.executions.swap(0, Ordering::Relaxed);
143        let idle_for = self.idle_for.swap(0, Ordering::Relaxed);
144        (connected, executions, idle_for)
145    }
146}
147
148type Result<T> = std::result::Result<T, Error>;
149
150pub async fn new_mysql_pool(
151    config: &Config,
152    pool_stat: Option<Arc<PoolStat>>,
153) -> Result<MySqlPool> {
154    let database_config = new_database_config(config)?;
155    let password = database_config.password.clone().unwrap_or_default();
156    let url = database_config.url.replace(&password, "***");
157    let category = "sqlx";
158    info!(category, url, "connect to database");
159
160    let mut options = PoolOptions::new()
161        .max_connections(database_config.max_connections)
162        .min_connections(database_config.min_connections)
163        .idle_timeout(database_config.idle_timeout)
164        .max_lifetime(database_config.max_lifetime)
165        .test_before_acquire(database_config.test_before_acquire);
166
167    if let Some(pool_stat) = pool_stat {
168        let after_connect_pool_stat = pool_stat.clone();
169        let before_acquire_pool_stat = pool_stat.clone();
170        options = options
171            .after_connect(move |_conn, _meta| {
172                let stat = after_connect_pool_stat.clone();
173                Box::pin(async move {
174                    stat.connected.fetch_add(1, Ordering::Relaxed);
175                    Ok(())
176                })
177            })
178            .before_acquire(move |_conn, meta| {
179                let stat = before_acquire_pool_stat.clone();
180                Box::pin(async move {
181                    let idle = meta.idle_for.as_secs();
182                    info!(category, age = meta.age.as_secs(), idle, "before acquire");
183                    stat.executions.fetch_add(1, Ordering::Relaxed);
184                    stat.idle_for.fetch_add(idle, Ordering::Relaxed);
185                    Ok(true)
186                })
187            });
188    }
189
190    let pool = options
191        .connect(database_config.url.as_str())
192        .await
193        .map_err(|e| Error::Sqlx { source: e })?;
194    Ok(pool)
195}
tibba_sql/lib.rs

tibba_sql/
lib.rs
