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
51fn default_connect_timeout() -> Duration {
52    Duration::from_secs(3)
53}
54
55fn default_idle_timeout() -> Duration {
56    Duration::from_secs(60)
57}
58
59fn default_max_lifetime() -> Duration {
60    Duration::from_secs(6 * 60 * 60)
61}
62
63fn default_test_before_acquire() -> bool {
64    true
65}
66
67#[derive(Deserialize, Debug, Clone)]
68struct DatabaseQuery {
69    #[serde(default = "default_max_connections")]
70    pub max_connections: u32,
71    #[serde(default = "default_min_connections")]
72    pub min_connections: u32,
73    #[serde(default = "default_connect_timeout")]
74    #[serde(with = "humantime_serde")]
75    pub connect_timeout: Duration,
76    #[serde(default = "default_idle_timeout")]
77    #[serde(with = "humantime_serde")]
78    pub idle_timeout: Duration,
79    #[serde(default = "default_max_lifetime")]
80    #[serde(with = "humantime_serde")]
81    pub max_lifetime: Duration,
82    #[serde(default = "default_test_before_acquire")]
83    pub test_before_acquire: bool,
84}
85
86fn map_err(e: impl ToString) -> Error {
87    Error::Common {
88        category: "config".to_string(),
89        message: e.to_string(),
90    }
91}
92
93// Creates a new DatabaseConfig instance from the configuration
94fn new_database_config(config: &Config) -> Result<DatabaseConfig> {
95    let origin_url = config.get_string("uri").map_err(map_err)?;
96    let url = origin_url.clone();
97    let parsed = parse_uri::<DatabaseQuery>(&url).map_err(map_err)?;
98
99    let mut url = parsed.url().map_err(map_err)?;
100    url.set_query(None);
101
102    let query = &parsed.query;
103    let database_config = DatabaseConfig {
104        origin_url,
105        url: url.to_string(),
106        max_connections: query.max_connections,
107        min_connections: query.min_connections,
108        connect_timeout: query.connect_timeout,
109        idle_timeout: query.idle_timeout,
110        max_lifetime: query.max_lifetime,
111        test_before_acquire: query.test_before_acquire,
112        password: parsed.password.map(|v| v.to_string()),
113    };
114    database_config
115        .validate()
116        .map_err(|e| Error::Validate { source: e })?;
117    Ok(database_config)
118}
119
120#[derive(Debug, Snafu)]
121pub enum Error {
122    #[snafu(display("sqlx error: {source}"))]
123    Sqlx { source: sqlx::Error },
124    #[snafu(display("validate error: {source}"))]
125    Validate { source: validator::ValidationErrors },
126    #[snafu(display("category: {category}, error: {message}"))]
127    Common { category: String, message: String },
128}
129
130impl From<Error> for BaseError {
131    fn from(source: Error) -> Self {
132        let err = match source {
133            Error::Sqlx { source } => BaseError::new(source)
134                .with_sub_category("sqlx")
135                .with_exception(true),
136            Error::Validate { source } => BaseError::new(source)
137                .with_sub_category("validate")
138                .with_exception(true),
139            Error::Common { message, .. } => BaseError::new(message).with_exception(true),
140        };
141        err.with_category("sql")
142    }
143}
144
145#[derive(Debug, Default)]
146pub struct PoolStat {
147    connected: AtomicU32,
148    executions: AtomicUsize,
149    idle_for: AtomicU64,
150}
151
152impl PoolStat {
153    pub fn stat(&self) -> (u32, usize, u64) {
154        let connected = self.connected.swap(0, Ordering::Relaxed);
155        let executions = self.executions.swap(0, Ordering::Relaxed);
156        let idle_for = self.idle_for.swap(0, Ordering::Relaxed);
157        (connected, executions, idle_for)
158    }
159}
160
161type Result<T> = std::result::Result<T, Error>;
162
163pub async fn new_mysql_pool(
164    config: &Config,
165    pool_stat: Option<Arc<PoolStat>>,
166) -> Result<MySqlPool> {
167    let database_config = new_database_config(config)?;
168    let password = database_config.password.clone().unwrap_or_default();
169    let url = database_config.url.replace(&password, "***");
170    let category = "sqlx";
171    info!(category, url, "connect to database");
172
173    let mut options = PoolOptions::new()
174        .max_connections(database_config.max_connections)
175        .min_connections(database_config.min_connections)
176        .idle_timeout(database_config.idle_timeout)
177        .max_lifetime(database_config.max_lifetime)
178        .test_before_acquire(database_config.test_before_acquire);
179
180    if let Some(pool_stat) = pool_stat {
181        let after_connect_pool_stat = pool_stat.clone();
182        let before_acquire_pool_stat = pool_stat.clone();
183        options = options
184            .after_connect(move |_conn, _meta| {
185                let stat = after_connect_pool_stat.clone();
186                Box::pin(async move {
187                    stat.connected.fetch_add(1, Ordering::Relaxed);
188                    Ok(())
189                })
190            })
191            .before_acquire(move |_conn, meta| {
192                let stat = before_acquire_pool_stat.clone();
193                Box::pin(async move {
194                    let idle = meta.idle_for.as_secs();
195                    info!(category, age = meta.age.as_secs(), idle, "before acquire");
196                    stat.executions.fetch_add(1, Ordering::Relaxed);
197                    stat.idle_for.fetch_add(idle, Ordering::Relaxed);
198                    Ok(true)
199                })
200            });
201    }
202
203    let pool = options
204        .connect(database_config.url.as_str())
205        .await
206        .map_err(|e| Error::Sqlx { source: e })?;
207    Ok(pool)
208}
tibba_sql/lib.rs

tibba_sql/
lib.rs
