chronon_backend_sql_common/
backend.rs

1use std::fmt;
2use std::str::FromStr;
3use std::time::Duration;
4
5use chronon_core::Result;
6use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode};
7use sqlx::{Executor, Pool, Postgres, Sqlite};
8
9use crate::error_map::{map_connect_err, map_err};
10use crate::schema;
11
12/// Max Postgres pool connections (`CHRONON_PG_POOL_SIZE`, default 5, cap 200).
13#[must_use]
14pub fn postgres_max_connections() -> u32 {
15    std::env::var("CHRONON_PG_POOL_SIZE")
16        .ok()
17        .and_then(|v| v.parse::<u32>().ok())
18        .unwrap_or(5)
19        .clamp(1, 200)
20}
21
22fn sqlite_url_is_file_backed(url: &str) -> bool {
23    let lower = url.to_ascii_lowercase();
24    !lower.contains(":memory:") && !lower.contains("mode=memory")
25}
26
27/// `SQLite` uses `?` placeholders; `PostgreSQL` uses `$1`, `$2`, …
28pub fn bind_sql(dialect: SqlDialect, sql: &str) -> String {
29    match dialect {
30        SqlDialect::Sqlite => sql.to_string(),
31        SqlDialect::Postgres => {
32            let mut out = String::with_capacity(sql.len());
33            let mut n = 1u32;
34            for ch in sql.chars() {
35                if ch == '?' {
36                    out.push('$');
37                    out.push_str(&n.to_string());
38                    n += 1;
39                } else {
40                    out.push(ch);
41                }
42            }
43            out
44        }
45    }
46}
47
48/// Maximum length for an isolated Postgres schema identifier (Postgres `NAMEDATALEN` - 1).
49const MAX_POSTGRES_SCHEMA_NAME_LEN: usize = 63;
50
51/// Validate a Postgres schema name before interpolating it into DDL / `search_path`.
52///
53/// Accepts only `^[A-Za-z_][A-Za-z0-9_]*$` up to 63 characters. Rejects quote breakouts and
54/// other identifier injection patterns.
55///
56/// # Errors
57///
58/// Returns [`chronon_core::ChrononError::ParamError`] when the name is empty, too long, or
59/// contains disallowed characters.
60pub fn validate_postgres_schema_name(schema: &str) -> Result<()> {
61    if schema.is_empty() || schema.len() > MAX_POSTGRES_SCHEMA_NAME_LEN {
62        return Err(chronon_core::ChrononError::ParamError(format!(
63            "postgres schema name must be 1..{MAX_POSTGRES_SCHEMA_NAME_LEN} characters"
64        )));
65    }
66    let mut chars = schema.chars();
67    let first = chars.next().ok_or_else(|| {
68        chronon_core::ChrononError::ParamError("postgres schema name must not be empty".into())
69    })?;
70    if !(first.is_ascii_alphabetic() || first == '_') {
71        return Err(chronon_core::ChrononError::ParamError(
72            "postgres schema name must start with ASCII letter or underscore".into(),
73        ));
74    }
75    if !chars.all(|c| c.is_ascii_alphanumeric() || c == '_') {
76        return Err(chronon_core::ChrononError::ParamError(
77            "postgres schema name may contain only ASCII letters, digits, and underscores".into(),
78        ));
79    }
80    Ok(())
81}
82
83/// SQL dialect for query variants.
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub enum SqlDialect {
86    /// `PostgreSQL`.
87    Postgres,
88    /// `SQLite`.
89    Sqlite,
90}
91
92/// Connection pool for a SQL backend.
93#[derive(Clone)]
94pub enum SqlPool {
95    /// `SQLite` pool.
96    Sqlite(Pool<Sqlite>),
97    /// `PostgreSQL` pool.
98    Postgres(Pool<Postgres>),
99}
100
101impl fmt::Debug for SqlPool {
102    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
103        match self {
104            Self::Sqlite(_) => f.debug_tuple("SqlPool::Sqlite").finish(),
105            Self::Postgres(_) => f.debug_tuple("SqlPool::Postgres").finish(),
106        }
107    }
108}
109
110/// SQL-backed [`SchedulerStore`](chronon_core::store::SchedulerStore) (`PostgreSQL` or `SQLite`).
111pub struct SqlSchedulerStore {
112    pub(crate) pool: SqlPool,
113    pub(crate) dialect: SqlDialect,
114}
115
116impl SqlSchedulerStore {
117    /// Open a `SQLite` pool, bootstrap schema, and return a store.
118    ///
119    /// File-backed URLs use WAL and a 5s busy timeout so concurrent scheduler
120    /// writes wait instead of failing immediately with `SQLITE_BUSY`. In-memory
121    /// URLs keep the default journal; WAL is not valid for `:memory:`.
122    ///
123    /// # Errors
124    ///
125    /// Returns a storage error if the pool connection or schema bootstrap fails.
126    pub async fn connect_sqlite(url: &str) -> Result<Self> {
127        let mut options = SqliteConnectOptions::from_str(url)
128            .map_err(|e| map_connect_err("sqlite", url, e))?
129            .create_if_missing(true)
130            .busy_timeout(Duration::from_secs(5));
131        if sqlite_url_is_file_backed(url) {
132            options = options.journal_mode(SqliteJournalMode::Wal);
133        }
134        let pool = sqlx::sqlite::SqlitePoolOptions::new()
135            .max_connections(5)
136            .connect_with(options)
137            .await
138            .map_err(|e| map_connect_err("sqlite", url, e))?;
139        Self::from_sqlite_pool(pool).await
140    }
141
142    /// Open a `PostgreSQL` pool, bootstrap schema, and return a store.
143    ///
144    /// # Errors
145    ///
146    /// Returns a storage error if the pool connection or schema bootstrap fails.
147    pub async fn connect_postgres(url: &str) -> Result<Self> {
148        let pool = sqlx::postgres::PgPoolOptions::new()
149            .max_connections(postgres_max_connections())
150            .connect(url)
151            .await
152            .map_err(|e| map_connect_err("postgres", url, e))?;
153        Self::from_postgres_pool(pool).await
154    }
155
156    /// Connect to `PostgreSQL` with an isolated schema for parallel tests.
157    ///
158    /// # Errors
159    ///
160    /// Returns [`chronon_core::ChrononError::ParamError`] when `schema` fails
161    /// [`validate_postgres_schema_name`]. Returns a storage error if schema creation, pool
162    /// connection, or bootstrap fails.
163    pub async fn connect_postgres_isolated(url: &str, schema: &str) -> Result<Self> {
164        validate_postgres_schema_name(schema)?;
165        let admin = sqlx::postgres::PgPoolOptions::new()
166            .max_connections(1)
167            .connect(url)
168            .await
169            .map_err(|e| map_connect_err("postgres", url, e))?;
170        let ddl = format!("CREATE SCHEMA IF NOT EXISTS \"{schema}\"");
171        admin.execute(ddl.as_str()).await.map_err(map_err)?;
172        drop(admin);
173
174        let schema = schema.to_string();
175        let pool = sqlx::postgres::PgPoolOptions::new()
176            .max_connections(postgres_max_connections())
177            .after_connect(move |conn, _meta| {
178                let schema = schema.clone();
179                Box::pin(async move {
180                    let sql = format!("SET search_path TO \"{schema}\"");
181                    sqlx::query(&sql).execute(conn).await?;
182                    Ok(())
183                })
184            })
185            .connect(url)
186            .await
187            .map_err(|e| map_connect_err("postgres", url, e))?;
188        Self::from_postgres_pool(pool).await
189    }
190
191    /// Attach to an existing isolated schema without re-running DDL bootstrap.
192    ///
193    /// Used when a test process already bootstrapped the schema and worker daemons join.
194    ///
195    /// # Errors
196    ///
197    /// Returns [`chronon_core::ChrononError::ParamError`] when `schema` fails
198    /// [`validate_postgres_schema_name`]. Returns a storage error if the pool cannot be opened.
199    pub async fn attach_postgres_isolated(url: &str, schema: &str) -> Result<Self> {
200        validate_postgres_schema_name(schema)?;
201        let schema = schema.to_string();
202        let pool = sqlx::postgres::PgPoolOptions::new()
203            .max_connections(postgres_max_connections())
204            .after_connect(move |conn, _meta| {
205                let schema = schema.clone();
206                Box::pin(async move {
207                    let sql = format!("SET search_path TO \"{schema}\"");
208                    sqlx::query(&sql).execute(conn).await?;
209                    Ok(())
210                })
211            })
212            .connect(url)
213            .await
214            .map_err(|e| map_connect_err("postgres", url, e))?;
215        Ok(Self {
216            pool: SqlPool::Postgres(pool),
217            dialect: SqlDialect::Postgres,
218        })
219    }
220
221    /// Wrap an existing `SQLite` pool (schema bootstrap runs).
222    ///
223    /// # Errors
224    ///
225    /// Returns a storage error if schema bootstrap fails.
226    pub async fn from_sqlite_pool(pool: Pool<Sqlite>) -> Result<Self> {
227        let store = Self {
228            pool: SqlPool::Sqlite(pool),
229            dialect: SqlDialect::Sqlite,
230        };
231        schema::ensure_schema(&store).await?;
232        Ok(store)
233    }
234
235    /// Wrap an existing `PostgreSQL` pool (schema bootstrap runs).
236    ///
237    /// # Errors
238    ///
239    /// Returns a storage error if schema bootstrap fails.
240    pub async fn from_postgres_pool(pool: Pool<Postgres>) -> Result<Self> {
241        let store = Self {
242            pool: SqlPool::Postgres(pool),
243            dialect: SqlDialect::Postgres,
244        };
245        schema::ensure_schema(&store).await?;
246        Ok(store)
247    }
248
249    /// Underlying connection pool.
250    #[must_use]
251    pub const fn pool(&self) -> &SqlPool {
252        &self.pool
253    }
254
255    /// Engine dialect.
256    #[must_use]
257    pub const fn dialect(&self) -> SqlDialect {
258        self.dialect
259    }
260
261    pub(crate) async fn run_ddl(&self, ddl: &str) -> Result<()> {
262        match &self.pool {
263            SqlPool::Sqlite(pool) => {
264                pool.execute(ddl).await.map_err(map_err)?;
265            }
266            SqlPool::Postgres(pool) => {
267                pool.execute(ddl).await.map_err(map_err)?;
268            }
269        }
270        Ok(())
271    }
272
273    /// Drop an isolated PostgreSQL schema (bench cell reset).
274    ///
275    /// # Errors
276    ///
277    /// Returns [`chronon_core::ChrononError::ParamError`] when `schema` fails
278    /// [`validate_postgres_schema_name`]. Returns a storage error when the admin connection or
279    /// DDL fails.
280    pub async fn drop_postgres_schema(url: &str, schema: &str) -> Result<()> {
281        validate_postgres_schema_name(schema)?;
282        let admin = sqlx::postgres::PgPoolOptions::new()
283            .max_connections(1)
284            .connect(url)
285            .await
286            .map_err(|e| map_connect_err("postgres", url, e))?;
287        let ddl = format!("DROP SCHEMA IF EXISTS \"{schema}\" CASCADE");
288        admin.execute(ddl.as_str()).await.map_err(map_err)?;
289        Ok(())
290    }
291}
292
293impl fmt::Debug for SqlSchedulerStore {
294    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
295        f.debug_struct("SqlSchedulerStore")
296            .field("dialect", &self.dialect)
297            .finish_non_exhaustive()
298    }
299}
300
301#[cfg(test)]
302mod tests {
303    use super::{bind_sql, sqlite_url_is_file_backed, validate_postgres_schema_name, SqlDialect};
304
305    #[test]
306    fn bind_sql_sqlite_passthrough() {
307        let sql = "SELECT * FROM t WHERE id = ? AND name = ?";
308        assert_eq!(bind_sql(SqlDialect::Sqlite, sql), sql);
309    }
310
311    #[test]
312    fn bind_sql_postgres_renumbers_placeholders() {
313        let sql = "UPDATE t SET a = ?, b = ? WHERE id = ?";
314        assert_eq!(
315            bind_sql(SqlDialect::Postgres, sql),
316            "UPDATE t SET a = $1, b = $2 WHERE id = $3"
317        );
318    }
319
320    #[test]
321    fn schema_name_accepts_safe_identifiers() {
322        assert!(validate_postgres_schema_name("bench_cell_1").is_ok());
323        assert!(validate_postgres_schema_name("_tmp").is_ok());
324        assert!(validate_postgres_schema_name("A").is_ok());
325    }
326
327    #[test]
328    fn schema_name_rejects_injection_and_empty() {
329        use chronon_core::ChrononError;
330
331        let too_long = "x".repeat(64);
332        for name in [
333            "",
334            "a\";drop",
335            "evil-name",
336            "1leading",
337            "has space",
338            too_long.as_str(),
339        ] {
340            match validate_postgres_schema_name(name) {
341                Err(ChrononError::ParamError(_)) => {}
342                other => panic!("expected ParamError for {name:?}, got {other:?}"),
343            }
344        }
345    }
346
347    #[test]
348    fn sqlite_url_file_backed_detection() {
349        assert!(sqlite_url_is_file_backed("sqlite:///tmp/chronon.db"));
350        assert!(sqlite_url_is_file_backed("sqlite://./bench.db?mode=rwc"));
351        assert!(!sqlite_url_is_file_backed("sqlite://:memory:"));
352        assert!(!sqlite_url_is_file_backed(
353            "sqlite://file:mem?mode=memory&cache=shared"
354        ));
355    }
356}
chronon_backend_sql_common/backend.rs

chronon_backend_sql_common/
backend.rs
