chronon_backend_sql_common/
backend.rs

1use std::fmt;
2
3use chronon_core::Result;
4use sqlx::{Executor, Pool, Postgres, Sqlite};
5
6use crate::error_map::map_err;
7use crate::schema;
8
9/// Max Postgres pool connections (`CHRONON_PG_POOL_SIZE`, default 5, cap 200).
10#[must_use]
11pub fn postgres_max_connections() -> u32 {
12    std::env::var("CHRONON_PG_POOL_SIZE")
13        .ok()
14        .and_then(|v| v.parse::<u32>().ok())
15        .unwrap_or(5)
16        .clamp(1, 200)
17}
18
19/// `SQLite` uses `?` placeholders; `PostgreSQL` uses `$1`, `$2`, …
20pub fn bind_sql(dialect: SqlDialect, sql: &str) -> String {
21    match dialect {
22        SqlDialect::Sqlite => sql.to_string(),
23        SqlDialect::Postgres => {
24            let mut out = String::with_capacity(sql.len());
25            let mut n = 1u32;
26            for ch in sql.chars() {
27                if ch == '?' {
28                    out.push('$');
29                    out.push_str(&n.to_string());
30                    n += 1;
31                } else {
32                    out.push(ch);
33                }
34            }
35            out
36        }
37    }
38}
39
40/// SQL dialect for query variants.
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum SqlDialect {
43    /// `PostgreSQL`.
44    Postgres,
45    /// `SQLite`.
46    Sqlite,
47}
48
49/// Connection pool for a SQL backend.
50#[derive(Clone)]
51pub enum SqlPool {
52    /// `SQLite` pool.
53    Sqlite(Pool<Sqlite>),
54    /// `PostgreSQL` pool.
55    Postgres(Pool<Postgres>),
56}
57
58impl fmt::Debug for SqlPool {
59    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
60        match self {
61            Self::Sqlite(_) => f.debug_tuple("SqlPool::Sqlite").finish(),
62            Self::Postgres(_) => f.debug_tuple("SqlPool::Postgres").finish(),
63        }
64    }
65}
66
67/// SQL-backed [`SchedulerStore`](chronon_core::store::SchedulerStore) (`PostgreSQL` or `SQLite`).
68pub struct SqlSchedulerStore {
69    pub(crate) pool: SqlPool,
70    pub(crate) dialect: SqlDialect,
71}
72
73impl SqlSchedulerStore {
74    /// Open a `SQLite` pool, bootstrap schema, and return a store.
75    ///
76    /// # Errors
77    ///
78    /// Returns a storage error if the pool connection or schema bootstrap fails.
79    pub async fn connect_sqlite(url: &str) -> Result<Self> {
80        let pool = sqlx::sqlite::SqlitePoolOptions::new()
81            .max_connections(5)
82            .connect(url)
83            .await
84            .map_err(|e| map_err(&e))?;
85        Self::from_sqlite_pool(pool).await
86    }
87
88    /// Open a `PostgreSQL` pool, bootstrap schema, and return a store.
89    ///
90    /// # Errors
91    ///
92    /// Returns a storage error if the pool connection or schema bootstrap fails.
93    pub async fn connect_postgres(url: &str) -> Result<Self> {
94        let pool = sqlx::postgres::PgPoolOptions::new()
95            .max_connections(postgres_max_connections())
96            .connect(url)
97            .await
98            .map_err(|e| map_err(&e))?;
99        Self::from_postgres_pool(pool).await
100    }
101
102    /// Connect to `PostgreSQL` with an isolated schema for parallel tests.
103    ///
104    /// # Errors
105    ///
106    /// Returns a storage error if schema creation, pool connection, or bootstrap fails.
107    pub async fn connect_postgres_isolated(url: &str, schema: &str) -> Result<Self> {
108        let admin = sqlx::postgres::PgPoolOptions::new()
109            .max_connections(1)
110            .connect(url)
111            .await
112            .map_err(|e| map_err(&e))?;
113        let ddl = format!("CREATE SCHEMA IF NOT EXISTS \"{schema}\"");
114        admin.execute(ddl.as_str()).await.map_err(|e| map_err(&e))?;
115        drop(admin);
116
117        let schema = schema.to_string();
118        let pool = sqlx::postgres::PgPoolOptions::new()
119            .max_connections(postgres_max_connections())
120            .after_connect(move |conn, _meta| {
121                let schema = schema.clone();
122                Box::pin(async move {
123                    let sql = format!("SET search_path TO \"{schema}\"");
124                    sqlx::query(&sql).execute(conn).await?;
125                    Ok(())
126                })
127            })
128            .connect(url)
129            .await
130            .map_err(|e| map_err(&e))?;
131        Self::from_postgres_pool(pool).await
132    }
133
134    /// Attach to an existing isolated schema without re-running DDL bootstrap.
135    ///
136    /// Used when a test process already bootstrapped the schema and worker daemons join.
137    ///
138    /// # Errors
139    ///
140    /// Returns a storage error if the pool cannot be opened.
141    pub async fn attach_postgres_isolated(url: &str, schema: &str) -> Result<Self> {
142        let schema = schema.to_string();
143        let pool = sqlx::postgres::PgPoolOptions::new()
144            .max_connections(postgres_max_connections())
145            .after_connect(move |conn, _meta| {
146                let schema = schema.clone();
147                Box::pin(async move {
148                    let sql = format!("SET search_path TO \"{schema}\"");
149                    sqlx::query(&sql).execute(conn).await?;
150                    Ok(())
151                })
152            })
153            .connect(url)
154            .await
155            .map_err(|e| map_err(&e))?;
156        Ok(Self {
157            pool: SqlPool::Postgres(pool),
158            dialect: SqlDialect::Postgres,
159        })
160    }
161
162    /// Wrap an existing `SQLite` pool (schema bootstrap runs).
163    ///
164    /// # Errors
165    ///
166    /// Returns a storage error if schema bootstrap fails.
167    pub async fn from_sqlite_pool(pool: Pool<Sqlite>) -> Result<Self> {
168        let store = Self {
169            pool: SqlPool::Sqlite(pool),
170            dialect: SqlDialect::Sqlite,
171        };
172        schema::ensure_schema(&store).await?;
173        Ok(store)
174    }
175
176    /// Wrap an existing `PostgreSQL` pool (schema bootstrap runs).
177    ///
178    /// # Errors
179    ///
180    /// Returns a storage error if schema bootstrap fails.
181    pub async fn from_postgres_pool(pool: Pool<Postgres>) -> Result<Self> {
182        let store = Self {
183            pool: SqlPool::Postgres(pool),
184            dialect: SqlDialect::Postgres,
185        };
186        schema::ensure_schema(&store).await?;
187        Ok(store)
188    }
189
190    /// Underlying connection pool.
191    #[must_use]
192    pub const fn pool(&self) -> &SqlPool {
193        &self.pool
194    }
195
196    /// Engine dialect.
197    #[must_use]
198    pub const fn dialect(&self) -> SqlDialect {
199        self.dialect
200    }
201
202    pub(crate) async fn run_ddl(&self, ddl: &str) -> Result<()> {
203        match &self.pool {
204            SqlPool::Sqlite(pool) => {
205                pool.execute(ddl).await.map_err(|e| map_err(&e))?;
206            }
207            SqlPool::Postgres(pool) => {
208                pool.execute(ddl).await.map_err(|e| map_err(&e))?;
209            }
210        }
211        Ok(())
212    }
213
214    /// Drop an isolated PostgreSQL schema (bench cell reset).
215    ///
216    /// # Errors
217    ///
218    /// Returns a storage error when the admin connection or DDL fails.
219    pub async fn drop_postgres_schema(url: &str, schema: &str) -> Result<()> {
220        let admin = sqlx::postgres::PgPoolOptions::new()
221            .max_connections(1)
222            .connect(url)
223            .await
224            .map_err(|e| map_err(&e))?;
225        let ddl = format!("DROP SCHEMA IF EXISTS \"{schema}\" CASCADE");
226        admin.execute(ddl.as_str()).await.map_err(|e| map_err(&e))?;
227        Ok(())
228    }
229}
230
231impl fmt::Debug for SqlSchedulerStore {
232    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
233        f.debug_struct("SqlSchedulerStore")
234            .field("dialect", &self.dialect)
235            .finish_non_exhaustive()
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use super::{bind_sql, SqlDialect};
242
243    #[test]
244    fn bind_sql_sqlite_passthrough() {
245        let sql = "SELECT * FROM t WHERE id = ? AND name = ?";
246        assert_eq!(bind_sql(SqlDialect::Sqlite, sql), sql);
247    }
248
249    #[test]
250    fn bind_sql_postgres_renumbers_placeholders() {
251        let sql = "UPDATE t SET a = ?, b = ? WHERE id = ?";
252        assert_eq!(
253            bind_sql(SqlDialect::Postgres, sql),
254            "UPDATE t SET a = $1, b = $2 WHERE id = $3"
255        );
256    }
257}
chronon_backend_sql_common/backend.rs

chronon_backend_sql_common/
backend.rs
