aster_forge_db/
connection.rs

1//! Database connection helpers and SQLite reader/writer pool setup.
2//!
3//! The module wraps SeaORM connection creation with retry behavior, optional metrics callbacks, and
4//! SQLite-specific pooling rules. SQLite writer connections are constrained so transaction
5//! serialization is explicit, while read-only handles can still be split out when the URL supports
6//! it.
7
8use crate::{DbError, Result, retry};
9use sea_orm::{ConnectOptions, ConnectionTrait, Database, DatabaseConnection, SqlxSqliteConnector};
10use std::sync::Arc;
11
12use aster_forge_metrics::{
13    DbMetricBackend, DbQueryKind, DbQueryMetric, NoopDbMetrics, SharedDbMetricsRecorder,
14};
15
16/// Database connection configuration.
17#[derive(Clone, Debug)]
18pub struct DatabaseConfig {
19    /// Database URL understood by SeaORM/sqlx.
20    pub url: String,
21    /// Maximum pool size for non-SQLite connections and SQLite reader pools.
22    pub pool_size: u32,
23    /// Number of retries for connection establishment.
24    pub retry_count: u32,
25}
26
27impl DatabaseConfig {
28    /// Creates a config with conservative default pool and retry settings.
29    pub fn new(url: impl Into<String>) -> Self {
30        Self {
31            url: url.into(),
32            pool_size: 5,
33            retry_count: 3,
34        }
35    }
36}
37
38/// Pair of writer and reader database handles.
39#[derive(Clone)]
40pub struct DbHandles {
41    writer: DatabaseConnection,
42    reader: DatabaseConnection,
43    sqlite_read_write_split: bool,
44}
45
46impl DbHandles {
47    /// Creates writer/reader handles backed by the same connection.
48    pub fn single(db: DatabaseConnection) -> Self {
49        Self {
50            writer: db.clone(),
51            reader: db,
52            sqlite_read_write_split: false,
53        }
54    }
55
56    /// Returns the writer connection.
57    pub fn writer(&self) -> &DatabaseConnection {
58        &self.writer
59    }
60
61    /// Returns the reader connection.
62    pub fn reader(&self) -> &DatabaseConnection {
63        &self.reader
64    }
65
66    /// Returns whether SQLite writer and reader use separate pools.
67    pub fn sqlite_read_write_split(&self) -> bool {
68        self.sqlite_read_write_split
69    }
70
71    /// Closes the underlying database pools.
72    ///
73    /// Split SQLite handles own two independent pools, so both reader and writer must be closed.
74    /// Single-handle configurations clone the same pool into both fields and only close the writer
75    /// once.
76    ///
77    /// Both pools are always asked to close, even when the reader close fails: an early return
78    /// would leak the writer pool (SQLite WAL/file handles) with no way to retry, because `close`
79    /// consumes the handles. The reader error, being the first failure, is the one returned.
80    pub async fn close(self) -> Result<()> {
81        let reader_result = if self.sqlite_read_write_split {
82            Some(self.reader.close().await)
83        } else {
84            None
85        };
86        let writer_result = self.writer.close().await;
87        first_close_error(reader_result, writer_result)
88    }
89}
90
91/// Returns the first close failure after both pools have been asked to close.
92fn first_close_error(
93    reader: Option<std::result::Result<(), sea_orm::DbErr>>,
94    writer: std::result::Result<(), sea_orm::DbErr>,
95) -> Result<()> {
96    if let Some(Err(error)) = reader {
97        return Err(DbError::from(error));
98    }
99    writer.map_err(DbError::from)
100}
101
102/// Connects to the configured database and installs a metrics callback.
103pub async fn connect_with_metrics(
104    cfg: &DatabaseConfig,
105    metrics: SharedDbMetricsRecorder,
106) -> Result<DatabaseConnection> {
107    let retry_config = retry::RetryConfig {
108        max_retries: cfg.retry_count,
109        ..retry::RetryConfig::connection()
110    };
111    retry::with_retry(&retry_config, || {
112        Box::pin(connect_once(cfg, metrics.clone()))
113    })
114    .await
115}
116
117/// Connects to the configured database without metrics.
118pub async fn connect(cfg: &DatabaseConfig) -> Result<DatabaseConnection> {
119    connect_with_metrics(cfg, Arc::new(NoopDbMetrics)).await
120}
121
122/// Creates reader/writer handles for an existing writer connection and metrics recorder.
123pub async fn connect_reader_for_writer_with_metrics(
124    cfg: &DatabaseConfig,
125    writer: DatabaseConnection,
126    metrics: SharedDbMetricsRecorder,
127) -> Result<DbHandles> {
128    let url = normalize_database_url(&cfg.url);
129    if !sqlite_reader_pool_enabled(&url) {
130        return Ok(DbHandles::single(writer));
131    }
132
133    let retry_config = retry::RetryConfig {
134        max_retries: cfg.retry_count,
135        ..retry::RetryConfig::connection()
136    };
137    let reader = retry::with_retry(&retry_config, || {
138        connect_sqlite_reader_once(cfg, &url, metrics.clone())
139    })
140    .await?;
141    Ok(DbHandles {
142        writer,
143        reader,
144        sqlite_read_write_split: true,
145    })
146}
147
148/// Creates reader/writer handles for an existing writer connection without metrics.
149pub async fn connect_reader_for_writer(
150    cfg: &DatabaseConfig,
151    writer: DatabaseConnection,
152) -> Result<DbHandles> {
153    connect_reader_for_writer_with_metrics(cfg, writer, Arc::new(NoopDbMetrics)).await
154}
155
156async fn connect_once(
157    cfg: &DatabaseConfig,
158    metrics: SharedDbMetricsRecorder,
159) -> Result<DatabaseConnection> {
160    let url = normalize_database_url(&cfg.url);
161    let is_sqlite = url.starts_with("sqlite:");
162    // SQLite relies on a single pooled connection so concurrent writers are serialized at
163    // connection acquisition; repo-layer "lock" helpers do not emulate row locks there.
164    let max_connections = if is_sqlite { 1 } else { cfg.pool_size };
165
166    let mut opt = ConnectOptions::new(&url);
167    opt.max_connections(max_connections)
168        .min_connections(1)
169        .sqlx_logging(false)
170        .test_before_acquire(true);
171
172    // SeaORM's generic Database::connect() pre-validates URLs with url::Url::parse(),
173    // which rejects Windows-style SQLite paths containing backslashes. Route SQLite
174    // through sqlx's dedicated connector instead so platform-native paths keep working.
175    let db = if is_sqlite {
176        SqlxSqliteConnector::connect(opt)
177            .await
178            .map_err(DbError::database_connection)?
179    } else {
180        Database::connect(opt)
181            .await
182            .map_err(DbError::database_connection)?
183    };
184
185    let backend = db.get_database_backend();
186    tracing::info!(backend = ?backend, "database connected");
187
188    if is_sqlite {
189        tracing::info!(max_connections, "applying SQLite PRAGMA optimizations");
190        // PRAGMA failures classify through DbError::from: transient lock contention
191        // (SQLITE_BUSY while another process initializes the file) stays retryable,
192        // persistent failures (permissions, corrupt file) fail fast.
193        db.execute_unprepared("PRAGMA journal_mode=WAL;")
194            .await
195            .map_err(DbError::from)?;
196        db.execute_unprepared("PRAGMA busy_timeout=15000;")
197            .await
198            .map_err(DbError::from)?;
199        db.execute_unprepared("PRAGMA synchronous=NORMAL;")
200            .await
201            .map_err(DbError::from)?;
202        db.execute_unprepared("PRAGMA foreign_keys=ON;")
203            .await
204            .map_err(DbError::from)?;
205    }
206
207    let mut db = db;
208    install_db_metrics(&mut db, metrics);
209
210    Ok(db)
211}
212
213async fn connect_sqlite_reader_once(
214    cfg: &DatabaseConfig,
215    normalized_writer_url: &str,
216    metrics: SharedDbMetricsRecorder,
217) -> Result<DatabaseConnection> {
218    let reader_url = sqlite_reader_url(normalized_writer_url);
219    let max_connections = cfg.pool_size.max(1);
220    let mut opt = ConnectOptions::new(&reader_url);
221    opt.max_connections(max_connections)
222        .min_connections(1)
223        .sqlx_logging(false)
224        .test_before_acquire(true)
225        .map_sqlx_sqlite_pool_opts(|pool_options| {
226            pool_options.after_connect(|conn, _meta| {
227                Box::pin(async move {
228                    use sea_orm::sqlx::Executor;
229
230                    conn.execute("PRAGMA busy_timeout=15000;").await?;
231                    conn.execute("PRAGMA synchronous=NORMAL;").await?;
232                    conn.execute("PRAGMA foreign_keys=ON;").await?;
233                    conn.execute("PRAGMA query_only=ON;").await?;
234                    Ok(())
235                })
236            })
237        });
238
239    let mut db = SqlxSqliteConnector::connect(opt)
240        .await
241        .map_err(DbError::database_connection)?;
242    install_db_metrics(&mut db, metrics);
243
244    tracing::info!(
245        max_connections,
246        "SQLite reader pool connected with query_only pragma"
247    );
248    Ok(db)
249}
250
251fn normalize_database_url(database_url: &str) -> String {
252    if database_url == "sqlite::memory:" {
253        return database_url.to_string();
254    }
255
256    if database_url.starts_with("sqlite://") && !database_url.contains('?') {
257        return format!("{database_url}?mode=rwc");
258    }
259
260    database_url.to_string()
261}
262
263fn sqlite_reader_pool_enabled(normalized_url: &str) -> bool {
264    normalized_url.starts_with("sqlite:") && !is_sqlite_memory_url(normalized_url)
265}
266
267fn is_sqlite_memory_url(normalized_url: &str) -> bool {
268    normalized_url == "sqlite::memory:"
269        || normalized_url
270            .split_once('?')
271            .is_some_and(|(_, query)| query.split('&').any(|param| param == "mode=memory"))
272}
273
274fn sqlite_reader_url(normalized_writer_url: &str) -> String {
275    let Some((base, query)) = normalized_writer_url.split_once('?') else {
276        return format!("{normalized_writer_url}?mode=ro");
277    };
278
279    let mut saw_mode = false;
280    let query = query
281        .split('&')
282        .filter(|param| !param.is_empty())
283        .map(|param| {
284            if param.starts_with("mode=") {
285                saw_mode = true;
286                "mode=ro"
287            } else {
288                param
289            }
290        })
291        .collect::<Vec<_>>();
292
293    if saw_mode {
294        format!("{base}?{}", query.join("&"))
295    } else if query.is_empty() {
296        format!("{base}?mode=ro")
297    } else {
298        format!("{base}?mode=ro&{}", query.join("&"))
299    }
300}
301
302fn install_db_metrics(db: &mut DatabaseConnection, metrics: SharedDbMetricsRecorder) {
303    if !metrics.enabled() {
304        return;
305    }
306
307    db.set_metric_callback(move |info| {
308        metrics.record_db_query(&db_query_metric_from_sea_orm(info));
309    });
310}
311
312fn db_metric_backend_from_sea_orm(backend: sea_orm::DbBackend) -> DbMetricBackend {
313    match backend {
314        sea_orm::DbBackend::Sqlite => DbMetricBackend::Sqlite,
315        sea_orm::DbBackend::MySql => DbMetricBackend::MySql,
316        sea_orm::DbBackend::Postgres => DbMetricBackend::Postgres,
317        _ => DbMetricBackend::Other,
318    }
319}
320
321fn db_query_kind_from_sql(sql: &str) -> DbQueryKind {
322    match sql
323        .trim_start()
324        .split_ascii_whitespace()
325        .next()
326        .unwrap_or_default()
327        .to_ascii_uppercase()
328        .as_str()
329    {
330        "SELECT" => DbQueryKind::Select,
331        "INSERT" => DbQueryKind::Insert,
332        "UPDATE" => DbQueryKind::Update,
333        "DELETE" => DbQueryKind::Delete,
334        "WITH" => DbQueryKind::With,
335        "BEGIN" | "COMMIT" | "ROLLBACK" | "SAVEPOINT" | "RELEASE" => DbQueryKind::Transaction,
336        "CREATE" | "ALTER" | "DROP" | "TRUNCATE" => DbQueryKind::Ddl,
337        "PRAGMA" => DbQueryKind::Pragma,
338        _ => DbQueryKind::Other,
339    }
340}
341
342fn db_query_metric_from_sea_orm(info: &sea_orm::metric::Info<'_>) -> DbQueryMetric {
343    DbQueryMetric::new(
344        db_metric_backend_from_sea_orm(info.statement.db_backend),
345        db_query_kind_from_sql(&info.statement.sql),
346        info.failed,
347        info.elapsed,
348    )
349}
350
351#[cfg(test)]
352mod tests {
353    use super::{first_close_error, normalize_database_url};
354    use crate::connection::DatabaseConfig;
355    use aster_forge_metrics::{DbQueryKind, NoopDbMetrics};
356    use aster_forge_test::temp::SqliteTestDatabase;
357    use sea_orm::{ConnectionTrait, DbErr, TransactionTrait};
358    use std::sync::Arc;
359
360    #[test]
361    fn first_close_error_prefers_the_reader_failure() {
362        let error = first_close_error(
363            Some(Err(DbErr::Custom("reader blew up".to_string()))),
364            Err(DbErr::Custom("writer blew up".to_string())),
365        )
366        .expect_err("both pools failing must surface an error");
367        assert!(error.to_string().contains("reader blew up"));
368
369        let error = first_close_error(
370            Some(Err(DbErr::Custom("reader blew up".to_string()))),
371            Ok(()),
372        )
373        .expect_err("reader failure must surface even when the writer closes cleanly");
374        assert!(error.to_string().contains("reader blew up"));
375    }
376
377    #[test]
378    fn first_close_error_surfaces_writer_failure_and_clean_runs() {
379        let error = first_close_error(
380            Some(Ok(())),
381            Err(DbErr::Custom("writer blew up".to_string())),
382        )
383        .expect_err("writer failure must surface when the reader closed cleanly");
384        assert!(error.to_string().contains("writer blew up"));
385
386        // Single-handle configurations close only the writer.
387        let error = first_close_error(None, Err(DbErr::Custom("writer blew up".to_string())))
388            .expect_err("writer failure must surface without a split reader");
389        assert!(error.to_string().contains("writer blew up"));
390
391        assert!(first_close_error(Some(Ok(())), Ok(())).is_ok());
392        assert!(first_close_error(None, Ok(())).is_ok());
393    }
394
395    #[test]
396    fn sqlite_urls_without_query_default_to_rwc_mode() {
397        assert_eq!(
398            normalize_database_url("sqlite:///var/lib/asterdrive/app.db"),
399            "sqlite:///var/lib/asterdrive/app.db?mode=rwc"
400        );
401        assert_eq!(
402            normalize_database_url("sqlite://data/asterdrive.db"),
403            "sqlite://data/asterdrive.db?mode=rwc"
404        );
405    }
406
407    #[test]
408    fn sqlite_memory_and_existing_queries_are_preserved() {
409        assert_eq!(normalize_database_url("sqlite::memory:"), "sqlite::memory:");
410        assert_eq!(
411            normalize_database_url("sqlite:///var/lib/asterdrive/app.db?mode=ro"),
412            "sqlite:///var/lib/asterdrive/app.db?mode=ro"
413        );
414        assert_eq!(
415            normalize_database_url("postgres://user:pass@localhost/asterdrive"),
416            "postgres://user:pass@localhost/asterdrive"
417        );
418    }
419
420    #[test]
421    fn sqlite_reader_pool_skips_memory_databases() {
422        assert!(!super::sqlite_reader_pool_enabled("sqlite::memory:"));
423        assert!(!super::sqlite_reader_pool_enabled(
424            "sqlite://memory-test?mode=memory&cache=shared"
425        ));
426        assert!(super::sqlite_reader_pool_enabled(
427            "sqlite:///var/lib/asterdrive/app.db?mode=rwc"
428        ));
429    }
430
431    #[test]
432    fn sqlite_reader_url_forces_read_only_mode() {
433        assert_eq!(
434            super::sqlite_reader_url("sqlite:///var/lib/asterdrive/app.db?mode=rwc"),
435            "sqlite:///var/lib/asterdrive/app.db?mode=ro"
436        );
437        assert_eq!(
438            super::sqlite_reader_url("sqlite:///var/lib/asterdrive/app.db?cache=shared"),
439            "sqlite:///var/lib/asterdrive/app.db?mode=ro&cache=shared"
440        );
441        assert_eq!(
442            super::sqlite_reader_url("sqlite:///var/lib/asterdrive/app.db?"),
443            "sqlite:///var/lib/asterdrive/app.db?mode=ro"
444        );
445        assert_eq!(
446            super::sqlite_reader_url("sqlite:///var/lib/asterdrive/app.db?&cache=shared"),
447            "sqlite:///var/lib/asterdrive/app.db?mode=ro&cache=shared"
448        );
449        assert_eq!(
450            super::sqlite_reader_url("sqlite:///var/lib/asterdrive/app.db?mode=rwc&"),
451            "sqlite:///var/lib/asterdrive/app.db?mode=ro"
452        );
453    }
454
455    #[test]
456    fn db_query_kind_from_sql_matches_low_cardinality_labels() {
457        assert_eq!(
458            super::db_query_kind_from_sql(" SELECT 1"),
459            DbQueryKind::Select
460        );
461        assert_eq!(
462            super::db_query_kind_from_sql("insert into users values (1)"),
463            DbQueryKind::Insert
464        );
465        assert_eq!(
466            super::db_query_kind_from_sql("UPDATE users SET name = ?"),
467            DbQueryKind::Update
468        );
469        assert_eq!(
470            super::db_query_kind_from_sql("delete from users"),
471            DbQueryKind::Delete
472        );
473        assert_eq!(
474            super::db_query_kind_from_sql("WITH recent AS (SELECT 1) SELECT * FROM recent"),
475            DbQueryKind::With
476        );
477        assert_eq!(
478            super::db_query_kind_from_sql("BEGIN"),
479            DbQueryKind::Transaction
480        );
481        assert_eq!(
482            super::db_query_kind_from_sql("CREATE TABLE example (id integer)"),
483            DbQueryKind::Ddl
484        );
485        assert_eq!(
486            super::db_query_kind_from_sql("PRAGMA foreign_keys=ON"),
487            DbQueryKind::Pragma
488        );
489        assert_eq!(super::db_query_kind_from_sql("VACUUM"), DbQueryKind::Other);
490    }
491
492    #[tokio::test]
493    async fn sqlite_connector_accepts_windows_style_urls() {
494        let url = format!(
495            "sqlite://windows\\sqlite-url-{}?mode=memory&cache=shared",
496            uuid::Uuid::new_v4()
497        );
498        let db = super::connect_with_metrics(
499            &DatabaseConfig {
500                url,
501                pool_size: 10,
502                retry_count: 3,
503            },
504            Arc::new(NoopDbMetrics),
505        )
506        .await
507        .expect("sqlite connection should succeed for Windows-style URL");
508
509        db.execute_unprepared("SELECT 1;")
510            .await
511            .expect("sqlite query should succeed");
512    }
513
514    #[tokio::test]
515    async fn sqlite_memory_handles_use_single_connection() {
516        let cfg = DatabaseConfig {
517            url: "sqlite::memory:".to_string(),
518            pool_size: 4,
519            retry_count: 0,
520        };
521        let writer = super::connect_with_metrics(&cfg, Arc::new(NoopDbMetrics))
522            .await
523            .expect("sqlite memory writer should connect");
524        let handles =
525            super::connect_reader_for_writer_with_metrics(&cfg, writer, Arc::new(NoopDbMetrics))
526                .await
527                .expect("sqlite memory handles should connect");
528
529        assert!(!handles.sqlite_read_write_split());
530        assert_eq!(
531            handles.writer().get_database_backend(),
532            handles.reader().get_database_backend()
533        );
534
535        handles
536            .close()
537            .await
538            .expect("single sqlite handles should close");
539    }
540
541    #[tokio::test]
542    async fn sqlite_reader_pool_is_query_only() {
543        let database = SqliteTestDatabase::new("reader-pool");
544        let cfg = DatabaseConfig {
545            url: database.url().to_string(),
546            pool_size: 4,
547            retry_count: 0,
548        };
549        let writer = super::connect_with_metrics(&cfg, Arc::new(NoopDbMetrics))
550            .await
551            .expect("sqlite writer should connect");
552        let handles =
553            super::connect_reader_for_writer_with_metrics(&cfg, writer, Arc::new(NoopDbMetrics))
554                .await
555                .expect("sqlite handles should connect");
556        assert!(handles.sqlite_read_write_split());
557
558        handles
559            .writer()
560            .execute_unprepared("CREATE TABLE reader_guard (id INTEGER PRIMARY KEY);")
561            .await
562            .expect("writer should create table");
563
564        let write_result = handles
565            .reader()
566            .execute_unprepared("INSERT INTO reader_guard (id) VALUES (1);")
567            .await;
568        assert!(write_result.is_err(), "reader pool must reject writes");
569
570        handles
571            .close()
572            .await
573            .expect("split sqlite handles should close");
574    }
575
576    #[tokio::test]
577    async fn sqlite_reader_pool_reads_while_writer_connection_is_busy() {
578        let database = SqliteTestDatabase::new("reader-writer-split");
579        let cfg = DatabaseConfig {
580            url: database.url().to_string(),
581            pool_size: 4,
582            retry_count: 0,
583        };
584        let writer = super::connect_with_metrics(&cfg, Arc::new(NoopDbMetrics))
585            .await
586            .expect("sqlite writer should connect");
587        let handles =
588            super::connect_reader_for_writer_with_metrics(&cfg, writer, Arc::new(NoopDbMetrics))
589                .await
590                .expect("sqlite handles should connect");
591        assert!(handles.sqlite_read_write_split());
592
593        handles
594            .writer()
595            .execute_unprepared("CREATE TABLE split_guard (id INTEGER PRIMARY KEY, name TEXT);")
596            .await
597            .expect("writer should create table");
598        handles
599            .writer()
600            .execute_unprepared("INSERT INTO split_guard (id, name) VALUES (1, 'ready');")
601            .await
602            .expect("writer should seed row");
603
604        let txn = handles
605            .writer()
606            .begin()
607            .await
608            .expect("writer transaction should begin");
609        txn.execute_unprepared("UPDATE split_guard SET name = 'held' WHERE id = 1;")
610            .await
611            .expect("writer transaction should hold a write lock");
612
613        let read = tokio::time::timeout(std::time::Duration::from_millis(250), async {
614            handles
615                .reader()
616                .query_one_raw(sea_orm::Statement::from_string(
617                    sea_orm::DbBackend::Sqlite,
618                    "SELECT name FROM split_guard WHERE id = 1",
619                ))
620                .await
621        })
622        .await
623        .expect("reader should not wait on the writer pool queue")
624        .expect("reader query should succeed")
625        .expect("reader query should return row");
626        let name: String = read
627            .try_get_by_index(0)
628            .expect("reader query should decode name");
629        assert_eq!(name, "ready");
630
631        txn.rollback()
632            .await
633            .expect("writer transaction should roll back");
634
635        handles
636            .close()
637            .await
638            .expect("split sqlite handles should close after reader query");
639    }
640}
aster_forge_db/connection.rs

aster_forge_db/
connection.rs
