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

tibba_sql/
lib.rs
