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::{ResultExt, Snafu};
17use sqlx::PgPool;
18use sqlx::postgres::PgPoolOptions;
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/// 该 crate 所有日志事件的 tracing target。
29/// 可通过 `RUST_LOG=tibba:sql=info`（或 `debug`）进行过滤。
30const LOG_TARGET: &str = "tibba:sql";
31
32#[derive(Debug, Snafu)]
33pub enum Error {
34    /// SQLx 数据库操作错误，属于基础设施异常。
35    #[snafu(display("sqlx error: {source}"))]
36    Sqlx {
37        #[snafu(source(from(sqlx::Error, Box::new)))]
38        source: Box<sqlx::Error>,
39    },
40    /// 配置字段校验失败（如连接数范围越界等）。
41    #[snafu(display("validate error: {source}"))]
42    Validate {
43        #[snafu(source(from(validator::ValidationErrors, Box::new)))]
44        source: Box<validator::ValidationErrors>,
45    },
46    /// 读取应用配置失败。
47    #[snafu(display("config error: {source}"))]
48    Config {
49        #[snafu(source(from(tibba_config::Error, Box::new)))]
50        source: Box<tibba_config::Error>,
51    },
52    /// 数据库 URI 解析失败。
53    #[snafu(display("parse uri error: {source}"))]
54    ParseUri {
55        #[snafu(source(from(tibba_util::Error, Box::new)))]
56        source: Box<tibba_util::Error>,
57    },
58}
59
60impl From<Error> for BaseError {
61    fn from(val: Error) -> Self {
62        let err = match val {
63            Error::Sqlx { source } => BaseError::new(source)
64                .with_sub_category("sqlx")
65                .with_exception(true),
66            Error::Validate { source } => BaseError::new(source).with_sub_category("validate"),
67            Error::Config { source } => BaseError::new(source).with_sub_category("config"),
68            Error::ParseUri { source } => BaseError::new(source).with_sub_category("parse_uri"),
69        };
70        err.with_category("sql")
71    }
72}
73
74type Result<T> = std::result::Result<T, Error>;
75
76/// 数据库连接池配置，字段均通过 URI 查询参数解析后填充。
77#[derive(Debug, Clone, Default, Validate)]
78pub struct DatabaseConfig {
79    /// 原始数据库 URI（含密码，仅内部使用）
80    pub origin_url: String,
81    /// 脱敏后的数据库连接 URL（去除查询参数）
82    #[validate(length(min = 10))]
83    pub url: String,
84    /// 连接池最大连接数（2–1000）
85    #[validate(range(min = 2, max = 1000))]
86    pub max_connections: u32,
87    /// 连接池最小保活连接数（0–10）
88    #[validate(range(min = 0, max = 10))]
89    pub min_connections: u32,
90    /// 建立连接的超时时间
91    pub connect_timeout: Duration,
92    /// 连接空闲超时时间，超出后连接将被回收
93    pub idle_timeout: Duration,
94    /// 连接最大存活时间，超出后强制重建
95    pub max_lifetime: Duration,
96    /// 取出连接前是否先执行健康检测
97    pub test_before_acquire: bool,
98    /// 数据库密码（从 URI 中提取，用于日志脱敏）
99    pub password: Option<String>,
100}
101
102fn default_max_connections() -> u32 {
103    10
104}
105fn default_min_connections() -> u32 {
106    2
107}
108fn default_connect_timeout() -> Duration {
109    Duration::from_secs(3)
110}
111fn default_idle_timeout() -> Duration {
112    Duration::from_secs(60)
113}
114fn default_max_lifetime() -> Duration {
115    Duration::from_secs(6 * 60 * 60)
116}
117fn default_test_before_acquire() -> bool {
118    true
119}
120
121/// 从 URI 查询字符串反序列化的连接池参数，未设置时使用各自的默认值。
122#[derive(Deserialize, Debug, Clone)]
123struct DatabaseQuery {
124    #[serde(default = "default_max_connections")]
125    pub max_connections: u32,
126    #[serde(default = "default_min_connections")]
127    pub min_connections: u32,
128    #[serde(default = "default_connect_timeout")]
129    #[serde(with = "humantime_serde")]
130    pub connect_timeout: Duration,
131    #[serde(default = "default_idle_timeout")]
132    #[serde(with = "humantime_serde")]
133    pub idle_timeout: Duration,
134    #[serde(default = "default_max_lifetime")]
135    #[serde(with = "humantime_serde")]
136    pub max_lifetime: Duration,
137    #[serde(default = "default_test_before_acquire")]
138    pub test_before_acquire: bool,
139}
140
141/// 连接池运行时统计，所有计数器在每次调用 `stat()` 时原子性地读取并重置为 0。
142#[derive(Debug, Default)]
143pub struct PoolStat {
144    /// 自上次读取以来新建的连接数
145    connected: AtomicU32,
146    /// 自上次读取以来连接被取出（acquire）的次数
147    executions: AtomicUsize,
148    /// 自上次读取以来所有连接取出前累计的空闲时间（秒）
149    idle_for: AtomicU64,
150}
151
152impl PoolStat {
153    /// 原子性地读取并重置所有计数器，返回 `(新建连接数, 取出次数, 累计空闲秒数)`。
154    pub fn stat(&self) -> (u32, usize, u64) {
155        let connected = self.connected.swap(0, Ordering::Relaxed);
156        let executions = self.executions.swap(0, Ordering::Relaxed);
157        let idle_for = self.idle_for.swap(0, Ordering::Relaxed);
158        (connected, executions, idle_for)
159    }
160}
161
162/// 从应用配置中解析并校验 `DatabaseConfig`。
163/// 密码从 URI 中单独提取，用于后续日志脱敏；URL 去除查询参数后存储。
164fn new_database_config(config: &Config) -> Result<DatabaseConfig> {
165    let origin_url = config.get_string("uri").context(ConfigSnafu)?;
166    // ParsedUri 借用字符串，需提前 clone 再将 origin_url 移入结构体
167    let url_str = origin_url.clone();
168    let parsed = parse_uri::<DatabaseQuery>(&url_str).context(ParseUriSnafu)?;
169
170    let mut url = parsed.url().context(ParseUriSnafu)?;
171    // 去除查询参数，避免将连接池配置混入实际连接 URL
172    url.set_query(None);
173
174    let query = &parsed.query;
175    let database_config = DatabaseConfig {
176        url: url.to_string(),
177        origin_url,
178        max_connections: query.max_connections,
179        min_connections: query.min_connections,
180        connect_timeout: query.connect_timeout,
181        idle_timeout: query.idle_timeout,
182        max_lifetime: query.max_lifetime,
183        test_before_acquire: query.test_before_acquire,
184        password: parsed.password.map(|v| v.to_string()),
185    };
186    database_config.validate().context(ValidateSnafu)?;
187    Ok(database_config)
188}
189
190/// 根据配置创建并连接 PostgreSQL 连接池。
191/// 若提供了 `pool_stat`，则通过 `after_connect` 和 `before_acquire` 钩子
192/// 原子性地记录新建连接数、取出次数和连接空闲时间。
193pub async fn new_pg_pool(config: &Config, pool_stat: Option<Arc<PoolStat>>) -> Result<PgPool> {
194    let database_config = new_database_config(config)?;
195    // 日志中脱敏密码
196    let password = database_config.password.clone().unwrap_or_default();
197    let url = database_config.url.replace(&password, "***");
198    info!(target: LOG_TARGET, url, "connect to database");
199
200    let mut options = PgPoolOptions::new()
201        .max_connections(database_config.max_connections)
202        .min_connections(database_config.min_connections)
203        .idle_timeout(database_config.idle_timeout)
204        .max_lifetime(database_config.max_lifetime)
205        .test_before_acquire(database_config.test_before_acquire);
206
207    if let Some(pool_stat) = pool_stat {
208        let after_connect_pool_stat = pool_stat.clone();
209        let before_acquire_pool_stat = pool_stat.clone();
210        options = options
211            .after_connect(move |_conn, _meta| {
212                let stat = after_connect_pool_stat.clone();
213                Box::pin(async move {
214                    // 新建连接后累加计数并打印日志
215                    let connected = stat.connected.fetch_add(1, Ordering::Relaxed) + 1;
216                    info!(
217                        target: LOG_TARGET,
218                        connected,
219                        "after connect"
220                    );
221                    Ok(())
222                })
223            })
224            .before_acquire(move |_conn, meta| {
225                let stat = before_acquire_pool_stat.clone();
226                Box::pin(async move {
227                    // 取出连接前记录空闲时间和连接年龄，便于监控连接复用情况
228                    let idle = meta.idle_for.as_secs();
229                    info!(
230                        target: LOG_TARGET,
231                        age = meta.age.as_secs(),
232                        idle,
233                        "before acquire"
234                    );
235                    stat.executions.fetch_add(1, Ordering::Relaxed);
236                    stat.idle_for.fetch_add(idle, Ordering::Relaxed);
237                    Ok(true)
238                })
239            });
240    }
241
242    options
243        .connect(database_config.url.as_str())
244        .await
245        .context(SqlxSnafu)
246}
tibba_sql/lib.rs

tibba_sql/
lib.rs
