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            .connect(&self.connection_string)
126            .await
127            .map_err(|e| DbkitError::Pool(e.to_string()))
128    }
129
130    /// Pool health metrics.
131    pub fn pool_status(&self) -> PoolStatus {
132        PoolStatus {
133            max_size: self.pool.options().get_max_connections() as usize,
134            size: self.pool.size() as usize,
135            idle: self.pool.num_idle(),
136        }
137    }
138
139    async fn build_pool(config: &DbkitConfig) -> Result<AnyPool, sqlx::Error> {
140        AnyPoolOptions::new()
141            .max_connections(config.pool_size as u32)
142            .acquire_timeout(Duration::from_secs(config.connect_timeout_secs))
143            .idle_timeout(Duration::from_secs(config.idle_timeout_secs))
144            .connect(&config.url)
145            .await
146    }
147
148    /// Create the database if it does not already exist.
149    ///
150    /// Uses sqlx's backend-agnostic `MigrateDatabase`, which handles
151    /// `CREATE DATABASE` for Postgres/MySQL and file creation for SQLite.
152    async fn ensure_database(url: &str, db_name: &str) -> Result<(), DbkitError> {
153        let exists = Any::database_exists(url).await.map_err(|e| {
154            DbkitError::DatabaseCreation {
155                name: db_name.to_string(),
156                reason: e.to_string(),
157            }
158        })?;
159
160        if !exists {
161            warn!("database '{}' does not exist, creating...", db_name);
162            Any::create_database(url)
163                .await
164                .map_err(|e| DbkitError::DatabaseCreation {
165                    name: db_name.to_string(),
166                    reason: e.to_string(),
167                })?;
168            info!("database '{}' created", db_name);
169        }
170
171        Ok(())
172    }
173
174    /// Map a sqlx connection error onto a more specific [`DbkitError`] where the
175    /// backend exposes a recognizable SQLSTATE.
176    fn map_connect_error(e: sqlx::Error, db_name: &str) -> DbkitError {
177        if let sqlx::Error::Database(ref db) = e {
178            match db.code().as_deref() {
179                // Postgres: invalid_password / invalid_authorization_specification
180                Some("28P01") | Some("28000") => return DbkitError::AuthFailed,
181                // Postgres: too_many_connections
182                Some("53300") => return DbkitError::TooManyConnections,
183                _ => {}
184            }
185        }
186        DbkitError::Connection(format!("could not connect to '{}': {}", db_name, e))
187    }
188
189    fn extract_db_name(url: &str, backend: Backend) -> String {
190        match backend {
191            // sqlite://path/to/file.db  ->  path/to/file.db
192            Backend::Sqlite => url
193                .splitn(2, ':')
194                .nth(1)
195                .unwrap_or("")
196                .trim_start_matches("//")
197                .split('?')
198                .next()
199                .unwrap_or("")
200                .to_string(),
201            // ...://host:port/dbname?params  ->  dbname
202            _ => url
203                .rsplit('/')
204                .next()
205                .unwrap_or("")
206                .split('?')
207                .next()
208                .unwrap_or("")
209                .to_string(),
210        }
211    }
212}
213
214/// Snapshot of connection pool health.
215#[derive(Debug, Clone)]
216pub struct PoolStatus {
217    /// Maximum number of connections in the pool.
218    pub max_size: usize,
219    /// Current number of connections (active + idle).
220    pub size: usize,
221    /// Number of idle connections available.
222    pub idle: usize,
223}
224
225impl std::fmt::Display for PoolStatus {
226    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
227        write!(
228            f,
229            "pool: {}/{} connections, {} idle",
230            self.size, self.max_size, self.idle
231        )
232    }
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238
239    #[test]
240    fn extract_db_name_server() {
241        assert_eq!(
242            ConnectionManager::extract_db_name("postgres://u:p@host:5432/myapp", Backend::Postgres),
243            "myapp"
244        );
245        assert_eq!(
246            ConnectionManager::extract_db_name(
247                "mysql://host:3306/app?ssl-mode=required",
248                Backend::MySql
249            ),
250            "app"
251        );
252    }
253
254    #[test]
255    fn extract_db_name_sqlite() {
256        assert_eq!(
257            ConnectionManager::extract_db_name("sqlite://data/app.db", Backend::Sqlite),
258            "data/app.db"
259        );
260    }
261}
dbkit/connection.rs

dbkit/
connection.rs
