Skip to main content

ag_store/
connection.rs

1//! `SQLite` connection setup for the database layer.
2
3use std::ops::Deref;
4use std::path::Path;
5use std::sync::Arc;
6use std::time::Duration;
7
8use sqlx::SqlitePool;
9use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous};
10
11use super::{AppRepositories, DbError, TimestampSource};
12use crate::timestamp::system_timestamp_source;
13
14/// Maximum number of pooled `SQLite` connections for the on-disk database.
15///
16/// `SQLite` still serializes writes in WAL mode, so the pool stays small and
17/// biased toward a handful of concurrent readers instead of a large number of
18/// queued writer contenders.
19pub(crate) const DB_POOL_MAX_CONNECTIONS: u32 = 4;
20
21/// Thin wrapper around a `SQLite` connection pool.
22#[derive(Clone)]
23pub struct Database {
24    pool: SqlitePool,
25    repositories: AppRepositories,
26}
27
28impl Database {
29    /// Opens the `SQLite` database and runs embedded migrations.
30    ///
31    /// Uses up to `DB_POOL_MAX_CONNECTIONS` pooled connections so UI reads can
32    /// stay responsive without oversizing the `SQLite` pool beyond what WAL
33    /// can use effectively. Applies a short busy timeout so bursty reducer
34    /// writes wait briefly for the single `SQLite` writer instead of failing
35    /// immediately with `SQLITE_BUSY`.
36    ///
37    /// # Errors
38    /// Returns an error if the directory cannot be created, the database
39    /// cannot be opened, or migrations fail.
40    pub async fn open(db_path: &Path) -> Result<Self, DbError> {
41        Self::open_with_timestamp_source(db_path, system_timestamp_source()).await
42    }
43
44    /// Opens the `SQLite` database with an injected persistence timestamp
45    /// source.
46    ///
47    /// # Errors
48    /// Returns an error if the directory cannot be created, the database
49    /// cannot be opened, or migrations fail.
50    pub async fn open_with_timestamp_source(
51        db_path: &Path,
52        timestamp_source: Arc<dyn TimestampSource>,
53    ) -> Result<Self, DbError> {
54        if let Some(parent) = db_path.parent() {
55            tokio::fs::create_dir_all(parent).await?;
56        }
57
58        let options = SqliteConnectOptions::new()
59            .filename(db_path)
60            .create_if_missing(true)
61            .busy_timeout(Duration::from_secs(2))
62            .journal_mode(SqliteJournalMode::Wal)
63            .synchronous(SqliteSynchronous::Normal)
64            .foreign_keys(true);
65
66        let pool = SqlitePoolOptions::new()
67            .max_connections(DB_POOL_MAX_CONNECTIONS)
68            .connect_with(options)
69            .await?;
70
71        sqlx::migrate!("./migrations").run(&pool).await?;
72
73        let repositories =
74            AppRepositories::from_pool_and_timestamp_source(pool.clone(), timestamp_source);
75
76        Ok(Self { pool, repositories })
77    }
78
79    /// Opens an in-memory `SQLite` database and runs migrations.
80    ///
81    /// This is primarily used by tests and any ephemeral workflows that need
82    /// an isolated database instance while keeping the same durability and
83    /// foreign-key settings as the on-disk database. Applies the same short
84    /// busy timeout as the on-disk configuration so tests exercise the same
85    /// writer wait policy.
86    ///
87    /// # Errors
88    /// Returns an error if the database connection or migrations fail.
89    pub async fn open_in_memory() -> Result<Self, DbError> {
90        Self::open_in_memory_with_timestamp_source(system_timestamp_source()).await
91    }
92
93    /// Opens an in-memory `SQLite` database with an injected timestamp source.
94    ///
95    /// # Errors
96    /// Returns an error if the database connection or migrations fail.
97    pub async fn open_in_memory_with_timestamp_source(
98        timestamp_source: Arc<dyn TimestampSource>,
99    ) -> Result<Self, DbError> {
100        let pool = open_in_memory_pool(1).await?;
101
102        let repositories =
103            AppRepositories::from_pool_and_timestamp_source(pool.clone(), timestamp_source);
104
105        Ok(Self { pool, repositories })
106    }
107
108    /// Returns the shared `SQLite` connection pool for lower-level query
109    /// access.
110    pub fn pool(&self) -> &SqlitePool {
111        &self.pool
112    }
113}
114
115impl Deref for Database {
116    type Target = AppRepositories;
117
118    fn deref(&self) -> &Self::Target {
119        &self.repositories
120    }
121}
122
123impl From<Database> for AppRepositories {
124    fn from(database: Database) -> Self {
125        database.repositories
126    }
127}
128
129/// Opens an in-memory `SQLite` pool with migrations applied.
130///
131/// The caller chooses the connection cap so tests and runtime code can share
132/// the same setup logic while keeping their own concurrency requirements.
133/// Applies the same 2-second busy timeout as the on-disk database so pooled
134/// test connections exercise the same writer wait policy.
135///
136/// # Errors
137/// Returns an error if the in-memory database connection or migrations fail.
138pub(crate) async fn open_in_memory_pool(max_connections: u32) -> Result<SqlitePool, DbError> {
139    let options = SqliteConnectOptions::new()
140        .filename(":memory:")
141        .busy_timeout(Duration::from_secs(2))
142        .journal_mode(SqliteJournalMode::Wal)
143        .synchronous(SqliteSynchronous::Normal)
144        .foreign_keys(true);
145
146    let pool = SqlitePoolOptions::new()
147        .max_connections(max_connections)
148        .connect_with(options)
149        .await?;
150
151    sqlx::migrate!("./migrations").run(&pool).await?;
152
153    Ok(pool)
154}
155
156#[cfg(test)]
157#[path = "connection_test.rs"]
158mod tests;