agentic_core/storage/
pool.rs

1//! Database connection pooling and initialization.
2
3use std::sync::Arc;
4
5use sqlx::any::AnyPoolOptions;
6
7/// Generic database pool type supporting `SQLite`, `PostgreSQL`, and `MySQL`.
8pub type DbPool = sqlx::Pool<sqlx::Any>;
9
10/// Database transaction type for multi-statement operations.
11pub type DbTransaction<'a> = sqlx::Transaction<'a, sqlx::Any>;
12
13/// Convenience type alias for database operation results.
14///
15/// All database queries return `DbResult<T>` which is `Result<T, sqlx::Error>`.
16pub type DbResult<T> = Result<T, sqlx::Error>;
17
18/// Prepares database URL with appropriate parameters.
19///
20/// For `SQLite` connections, adds `?mode=rwc` if not already present.
21/// This enables write mode (`rwc` = read-write-create) for file-based databases.
22///
23/// For other database types (`PostgreSQL`, `MySQL`), returns URL as-is.
24/// Defaults to `sqlite://./agentic_api.db` if no URL is provided.
25fn prepare_db_url(url: Option<&str>) -> String {
26    let url = url.unwrap_or("sqlite://./agentic_api.db");
27    if url.starts_with("sqlite") && !url.contains('?') {
28        format!("{url}?mode=rwc")
29    } else {
30        url.to_string()
31    }
32}
33
34/// Creates a connection pool for the database.
35///
36/// Initializes a connection pool with sensible defaults:
37/// - Max connections: 10 (configurable via [`AnyPoolOptions`])
38/// - Driver auto-detection: supports `SQLite`, `PostgreSQL`, `MySQL`
39/// - `SQLite` file mode: read-write-create for file-based databases
40///
41/// The pool is wrapped in `Arc` for thread-safe sharing across async tasks.
42///
43/// # Arguments
44///
45/// * `db_url` - Optional database connection URL. Defaults to `sqlite://./agentic_api.db` if `None`.
46///   Examples: `sqlite://data.db`, `postgresql://user:pass@host/db`
47///
48/// # Errors
49///
50/// Returns [`sqlx::Error`] if:
51/// - Connection URL is invalid
52/// - Database server is unreachable
53/// - Connection limit is exceeded
54/// - Authentication fails
55///
56pub async fn create_pool(db_url: Option<&str>) -> DbResult<Arc<DbPool>> {
57    // Install default drivers for auto-detection
58    sqlx::any::install_default_drivers();
59
60    // Prepare URL with database-specific parameters
61    let url = prepare_db_url(db_url);
62
63    // SQLite only allows one writer at a time; a single connection in the pool
64    // serializes writes at the pool level (queue).
65    // For other databases, 10 connections is a conservative default.
66    let max_connections = if url.starts_with("sqlite") { 1 } else { 10 };
67    let pool = AnyPoolOptions::new()
68        .max_connections(max_connections)
69        .connect(&url)
70        .await?;
71
72    // Wrap in Arc for thread-safe sharing across async tasks
73    Ok(Arc::new(pool))
74}
75
76/// Creates a connection pool and initializes the database schema.
77///
78/// Combines [`create_pool`] with schema initialization using [`PoolWithSchema`].
79/// Each pool has its own per-pool schema readiness flag.
80/// # Arguments
81///
82/// * `db_url` - Database connection URL
83///
84/// # Errors
85///
86/// Returns error if pool creation or schema initialization fails.
87pub async fn create_pool_with_schema(db_url: Option<&str>) -> DbResult<Arc<DbPool>> {
88    use crate::storage::PoolWithSchema;
89
90    let pool = create_pool(db_url).await?;
91    let pool_with_schema = PoolWithSchema::new(pool);
92    pool_with_schema.ensure_schema_ready().await?;
93
94    Ok(pool_with_schema.pool().clone())
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn test_prepare_sqlite_url_without_params() {
103        let url = "sqlite://test.db";
104        let prepared = prepare_db_url(Some(url));
105        assert_eq!(prepared, "sqlite://test.db?mode=rwc");
106    }
107
108    #[test]
109    fn test_prepare_sqlite_url_with_params() {
110        let url = "sqlite://test.db?cache=shared";
111        let prepared = prepare_db_url(Some(url));
112        assert_eq!(prepared, "sqlite://test.db?cache=shared");
113    }
114
115    #[test]
116    fn test_prepare_postgres_url() {
117        let url = "postgresql://user:pass@localhost/db";
118        let prepared = prepare_db_url(Some(url));
119        assert_eq!(prepared, "postgresql://user:pass@localhost/db");
120    }
121
122    #[test]
123    fn test_prepare_mysql_url() {
124        let url = "mysql://user:pass@localhost/db";
125        let prepared = prepare_db_url(Some(url));
126        assert_eq!(prepared, "mysql://user:pass@localhost/db");
127    }
128
129    #[test]
130    fn test_prepare_default_sqlite_url() {
131        let prepared = prepare_db_url(None);
132        assert_eq!(prepared, "sqlite://./agentic_api.db?mode=rwc");
133    }
134}
agentic_core/storage/pool.rs

agentic_core/storage/
pool.rs
