dbkit/
connection.rs

1use crate::config::{Backend, DbkitConfig};
2use crate::DbkitError;
3use sqlx::migrate::MigrateDatabase;
4use sqlx::{Any, AnyPool, any::AnyPoolOptions};
5use std::time::Duration;
6use tracing::{info, warn};
7
8/// Multi-backend connection pool with optional automatic database creation.
9///
10/// The backend (Postgres, MySQL, or SQLite) is selected from the URL scheme.
11/// If `auto_create_db` is enabled (default) and the target database doesn't
12/// exist, it is created before the pool is opened.
13pub struct ConnectionManager {
14    pool: AnyPool,
15    backend: Backend,
16    db_name: String,
17    connection_string: String,
18    config: DbkitConfig,
19}
20
21impl ConnectionManager {
22    /// Connect using a [`DbkitConfig`].
23    pub async fn connect(config: DbkitConfig) -> Result<Self, DbkitError> {
24        // Registers the Postgres/MySQL/SQLite drivers compiled in via features.
25        sqlx::any::install_default_drivers();
26
27        let backend = Backend::from_url(&config.url)?;
28        let db_name = Self::extract_db_name(&config.url, backend);
29        let connection_string = config.url.clone();
30
31        if config.auto_create_db {
32            Self::ensure_database(&config.url, &db_name).await?;
33        }
34
35        let pool = Self::build_pool(&config)
36            .await
37            .map_err(|e| Self::map_connect_error(e, &db_name))?;
38
39        info!("connected to database '{}' ({:?})", db_name, backend);
40
41        Ok(Self {
42            pool,
43            backend,
44            db_name,
45            connection_string,
46            config,
47        })
48    }
49
50    /// Connect using a connection URL with default settings.
51    ///
52    /// Shorthand for `ConnectionManager::connect(DbkitConfig::from_url(url))`.
53    pub async fn new(url: &str) -> Result<Self, DbkitError> {
54        Self::connect(DbkitConfig::from_url(url)).await
55    }
56
57    /// Get the underlying sqlx connection pool.
58    pub fn pool(&self) -> &AnyPool {
59        &self.pool
60    }
61
62    /// The detected backend for this connection.
63    pub fn backend(&self) -> Backend {
64        self.backend
65    }
66
67    /// Acquire a connection from the pool.
68    pub async fn get_connection(&self) -> Result<sqlx::pool::PoolConnection<Any>, DbkitError> {
69        self.pool
70            .acquire()
71            .await
72            .map_err(|e| DbkitError::Pool(e.to_string()))
73    }
74
75    /// Check if the database is reachable.
76    pub async fn is_connected(&self) -> bool {
77        self.pool.acquire().await.is_ok()
78    }
79
80    /// The database name extracted from the connection URL.
81    pub fn db_name(&self) -> &str {
82        &self.db_name
83    }
84
85    /// The full connection string.
86    pub fn connection_string(&self) -> &str {
87        &self.connection_string
88    }
89
90    /// The config used to create this connection.
91    pub fn config(&self) -> &DbkitConfig {
92        &self.config
93    }
94
95    /// Create a native sqlx [`PgPool`](sqlx::PgPool) from this connection's URL.
96    ///
97    /// The multi-backend `Any` pool can only represent basic scalar types
98    /// (bool/int/float/text/bytes). This native pool restores full Postgres type
99    /// support — `uuid`, `chrono` timestamps, `json`/`jsonb`, arrays, etc. — for
100    /// the queries that need it. Use it alongside [`pool`](Self::pool): `Any` for
101    /// portable work, the native pool for rich-typed Postgres work.
102    ///
103    /// Errors if this connection is not Postgres.
104    ///
105    /// ```ignore
106    /// let pg = conn.pg_native_pool().await?;
107    /// let row = sqlx::query("SELECT id, created_at FROM users WHERE id = $1")
108    ///     .bind(some_uuid)
109    ///     .fetch_one(&pg)
110    ///     .await?;
111    /// let id: sqlx::types::Uuid = row.get("id");
112    /// ```
113    #[cfg(feature = "postgres-native")]
114    pub async fn pg_native_pool(&self) -> Result<sqlx::PgPool, DbkitError> {
115        if self.backend != Backend::Postgres {
116            return Err(DbkitError::UnsupportedBackend(format!(
117                "pg_native_pool requires a Postgres connection, got {:?}",
118                self.backend
119            )));
120        }
121        sqlx::postgres::PgPoolOptions::new()
122            .max_connections(self.config.pool_size as u32)
123            .acquire_timeout(Duration::from_secs(self.config.connect_timeout_secs))
124            .idle_timeout(Duration::from_secs(self.config.idle_timeout_secs))
125            // Don't ping before every acquire — that's a round-trip per query.
126            // Matches the old deadpool `RecyclingMethod::Fast` behavior.
127            .test_before_acquire(false)
128            .connect(&self.connection_string)
129            .await
130            .map_err(|e| DbkitError::Pool(e.to_string()))
131    }
132
133    /// Pool health metrics.
134    pub fn pool_status(&self) -> PoolStatus {
135        PoolStatus {
136            max_size: self.pool.options().get_max_connections() as usize,
137            size: self.pool.size() as usize,
138            idle: self.pool.num_idle(),
139        }
140    }
141
142    async fn build_pool(config: &DbkitConfig) -> Result<AnyPool, sqlx::Error> {
143        AnyPoolOptions::new()
144            .max_connections(config.pool_size as u32)
145            .acquire_timeout(Duration::from_secs(config.connect_timeout_secs))
146            .idle_timeout(Duration::from_secs(config.idle_timeout_secs))
147            // Skip the per-acquire liveness ping (a round-trip per query).
148            .test_before_acquire(false)
149            .connect(&config.url)
150            .await
151    }
152
153    /// Create the database if it does not already exist.
154    ///
155    /// Uses sqlx's backend-agnostic `MigrateDatabase`, which handles
156    /// `CREATE DATABASE` for Postgres/MySQL and file creation for SQLite.
157    async fn ensure_database(url: &str, db_name: &str) -> Result<(), DbkitError> {
158        let exists = Any::database_exists(url).await.map_err(|e| {
159            DbkitError::DatabaseCreation {
160                name: db_name.to_string(),
161                reason: e.to_string(),
162            }
163        })?;
164
165        if !exists {
166            warn!("database '{}' does not exist, creating...", db_name);
167            Any::create_database(url)
168                .await
169                .map_err(|e| DbkitError::DatabaseCreation {
170                    name: db_name.to_string(),
171                    reason: e.to_string(),
172                })?;
173            info!("database '{}' created", db_name);
174        }
175
176        Ok(())
177    }
178
179    /// Map a sqlx connection error onto a more specific [`DbkitError`] where the
180    /// backend exposes a recognizable SQLSTATE.
181    fn map_connect_error(e: sqlx::Error, db_name: &str) -> DbkitError {
182        if let sqlx::Error::Database(ref db) = e {
183            match db.code().as_deref() {
184                // Postgres: invalid_password / invalid_authorization_specification
185                Some("28P01") | Some("28000") => return DbkitError::AuthFailed,
186                // Postgres: too_many_connections
187                Some("53300") => return DbkitError::TooManyConnections,
188                _ => {}
189            }
190        }
191        DbkitError::Connection(format!("could not connect to '{}': {}", db_name, e))
192    }
193
194    fn extract_db_name(url: &str, backend: Backend) -> String {
195        match backend {
196            // sqlite://path/to/file.db  ->  path/to/file.db
197            Backend::Sqlite => url
198                .split_once(':')
199                .map_or("", |(_, rest)| rest)
200                .trim_start_matches("//")
201                .split('?')
202                .next()
203                .unwrap_or("")
204                .to_string(),
205            // ...://host:port/dbname?params  ->  dbname
206            _ => url
207                .rsplit('/')
208                .next()
209                .unwrap_or("")
210                .split('?')
211                .next()
212                .unwrap_or("")
213                .to_string(),
214        }
215    }
216}
217
218/// Snapshot of connection pool health.
219#[derive(Debug, Clone)]
220pub struct PoolStatus {
221    /// Maximum number of connections in the pool.
222    pub max_size: usize,
223    /// Current number of connections (active + idle).
224    pub size: usize,
225    /// Number of idle connections available.
226    pub idle: usize,
227}
228
229impl std::fmt::Display for PoolStatus {
230    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
231        write!(
232            f,
233            "pool: {}/{} connections, {} idle",
234            self.size, self.max_size, self.idle
235        )
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242
243    #[test]
244    fn extract_db_name_server() {
245        assert_eq!(
246            ConnectionManager::extract_db_name("postgres://u:p@host:5432/myapp", Backend::Postgres),
247            "myapp"
248        );
249        assert_eq!(
250            ConnectionManager::extract_db_name(
251                "mysql://host:3306/app?ssl-mode=required",
252                Backend::MySql
253            ),
254            "app"
255        );
256    }
257
258    #[test]
259    fn extract_db_name_sqlite() {
260        assert_eq!(
261            ConnectionManager::extract_db_name("sqlite://data/app.db", Backend::Sqlite),
262            "data/app.db"
263        );
264    }
265}
dbkit/connection.rs

dbkit/
connection.rs
