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_connect_err, 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/// Maximum length for an isolated Postgres schema identifier (Postgres `NAMEDATALEN` - 1).
41const MAX_POSTGRES_SCHEMA_NAME_LEN: usize = 63;
42
43/// Validate a Postgres schema name before interpolating it into DDL / `search_path`.
44///
45/// Accepts only `^[A-Za-z_][A-Za-z0-9_]*$` up to 63 characters. Rejects quote breakouts and
46/// other identifier injection patterns.
47///
48/// # Errors
49///
50/// Returns [`chronon_core::ChrononError::ParamError`] when the name is empty, too long, or
51/// contains disallowed characters.
52pub fn validate_postgres_schema_name(schema: &str) -> Result<()> {
53    if schema.is_empty() || schema.len() > MAX_POSTGRES_SCHEMA_NAME_LEN {
54        return Err(chronon_core::ChrononError::ParamError(format!(
55            "postgres schema name must be 1..{MAX_POSTGRES_SCHEMA_NAME_LEN} characters"
56        )));
57    }
58    let mut chars = schema.chars();
59    let first = chars.next().ok_or_else(|| {
60        chronon_core::ChrononError::ParamError("postgres schema name must not be empty".into())
61    })?;
62    if !(first.is_ascii_alphabetic() || first == '_') {
63        return Err(chronon_core::ChrononError::ParamError(
64            "postgres schema name must start with ASCII letter or underscore".into(),
65        ));
66    }
67    if !chars.all(|c| c.is_ascii_alphanumeric() || c == '_') {
68        return Err(chronon_core::ChrononError::ParamError(
69            "postgres schema name may contain only ASCII letters, digits, and underscores".into(),
70        ));
71    }
72    Ok(())
73}
74
75/// SQL dialect for query variants.
76#[derive(Debug, Clone, Copy, PartialEq, Eq)]
77pub enum SqlDialect {
78    /// `PostgreSQL`.
79    Postgres,
80    /// `SQLite`.
81    Sqlite,
82}
83
84/// Connection pool for a SQL backend.
85#[derive(Clone)]
86pub enum SqlPool {
87    /// `SQLite` pool.
88    Sqlite(Pool<Sqlite>),
89    /// `PostgreSQL` pool.
90    Postgres(Pool<Postgres>),
91}
92
93impl fmt::Debug for SqlPool {
94    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
95        match self {
96            Self::Sqlite(_) => f.debug_tuple("SqlPool::Sqlite").finish(),
97            Self::Postgres(_) => f.debug_tuple("SqlPool::Postgres").finish(),
98        }
99    }
100}
101
102/// SQL-backed [`SchedulerStore`](chronon_core::store::SchedulerStore) (`PostgreSQL` or `SQLite`).
103pub struct SqlSchedulerStore {
104    pub(crate) pool: SqlPool,
105    pub(crate) dialect: SqlDialect,
106}
107
108impl SqlSchedulerStore {
109    /// Open a `SQLite` pool, bootstrap schema, and return a store.
110    ///
111    /// # Errors
112    ///
113    /// Returns a storage error if the pool connection or schema bootstrap fails.
114    pub async fn connect_sqlite(url: &str) -> Result<Self> {
115        let pool = sqlx::sqlite::SqlitePoolOptions::new()
116            .max_connections(5)
117            .connect(url)
118            .await
119            .map_err(|e| map_connect_err("sqlite", url, e))?;
120        Self::from_sqlite_pool(pool).await
121    }
122
123    /// Open a `PostgreSQL` pool, bootstrap schema, and return a store.
124    ///
125    /// # Errors
126    ///
127    /// Returns a storage error if the pool connection or schema bootstrap fails.
128    pub async fn connect_postgres(url: &str) -> Result<Self> {
129        let pool = sqlx::postgres::PgPoolOptions::new()
130            .max_connections(postgres_max_connections())
131            .connect(url)
132            .await
133            .map_err(|e| map_connect_err("postgres", url, e))?;
134        Self::from_postgres_pool(pool).await
135    }
136
137    /// Connect to `PostgreSQL` with an isolated schema for parallel tests.
138    ///
139    /// # Errors
140    ///
141    /// Returns [`chronon_core::ChrononError::ParamError`] when `schema` fails
142    /// [`validate_postgres_schema_name`]. Returns a storage error if schema creation, pool
143    /// connection, or bootstrap fails.
144    pub async fn connect_postgres_isolated(url: &str, schema: &str) -> Result<Self> {
145        validate_postgres_schema_name(schema)?;
146        let admin = sqlx::postgres::PgPoolOptions::new()
147            .max_connections(1)
148            .connect(url)
149            .await
150            .map_err(|e| map_connect_err("postgres", url, e))?;
151        let ddl = format!("CREATE SCHEMA IF NOT EXISTS \"{schema}\"");
152        admin.execute(ddl.as_str()).await.map_err(map_err)?;
153        drop(admin);
154
155        let schema = schema.to_string();
156        let pool = sqlx::postgres::PgPoolOptions::new()
157            .max_connections(postgres_max_connections())
158            .after_connect(move |conn, _meta| {
159                let schema = schema.clone();
160                Box::pin(async move {
161                    let sql = format!("SET search_path TO \"{schema}\"");
162                    sqlx::query(&sql).execute(conn).await?;
163                    Ok(())
164                })
165            })
166            .connect(url)
167            .await
168            .map_err(|e| map_connect_err("postgres", url, e))?;
169        Self::from_postgres_pool(pool).await
170    }
171
172    /// Attach to an existing isolated schema without re-running DDL bootstrap.
173    ///
174    /// Used when a test process already bootstrapped the schema and worker daemons join.
175    ///
176    /// # Errors
177    ///
178    /// Returns [`chronon_core::ChrononError::ParamError`] when `schema` fails
179    /// [`validate_postgres_schema_name`]. Returns a storage error if the pool cannot be opened.
180    pub async fn attach_postgres_isolated(url: &str, schema: &str) -> Result<Self> {
181        validate_postgres_schema_name(schema)?;
182        let schema = schema.to_string();
183        let pool = sqlx::postgres::PgPoolOptions::new()
184            .max_connections(postgres_max_connections())
185            .after_connect(move |conn, _meta| {
186                let schema = schema.clone();
187                Box::pin(async move {
188                    let sql = format!("SET search_path TO \"{schema}\"");
189                    sqlx::query(&sql).execute(conn).await?;
190                    Ok(())
191                })
192            })
193            .connect(url)
194            .await
195            .map_err(|e| map_connect_err("postgres", url, e))?;
196        Ok(Self {
197            pool: SqlPool::Postgres(pool),
198            dialect: SqlDialect::Postgres,
199        })
200    }
201
202    /// Wrap an existing `SQLite` pool (schema bootstrap runs).
203    ///
204    /// # Errors
205    ///
206    /// Returns a storage error if schema bootstrap fails.
207    pub async fn from_sqlite_pool(pool: Pool<Sqlite>) -> Result<Self> {
208        let store = Self {
209            pool: SqlPool::Sqlite(pool),
210            dialect: SqlDialect::Sqlite,
211        };
212        schema::ensure_schema(&store).await?;
213        Ok(store)
214    }
215
216    /// Wrap an existing `PostgreSQL` pool (schema bootstrap runs).
217    ///
218    /// # Errors
219    ///
220    /// Returns a storage error if schema bootstrap fails.
221    pub async fn from_postgres_pool(pool: Pool<Postgres>) -> Result<Self> {
222        let store = Self {
223            pool: SqlPool::Postgres(pool),
224            dialect: SqlDialect::Postgres,
225        };
226        schema::ensure_schema(&store).await?;
227        Ok(store)
228    }
229
230    /// Underlying connection pool.
231    #[must_use]
232    pub const fn pool(&self) -> &SqlPool {
233        &self.pool
234    }
235
236    /// Engine dialect.
237    #[must_use]
238    pub const fn dialect(&self) -> SqlDialect {
239        self.dialect
240    }
241
242    pub(crate) async fn run_ddl(&self, ddl: &str) -> Result<()> {
243        match &self.pool {
244            SqlPool::Sqlite(pool) => {
245                pool.execute(ddl).await.map_err(map_err)?;
246            }
247            SqlPool::Postgres(pool) => {
248                pool.execute(ddl).await.map_err(map_err)?;
249            }
250        }
251        Ok(())
252    }
253
254    /// Drop an isolated PostgreSQL schema (bench cell reset).
255    ///
256    /// # Errors
257    ///
258    /// Returns [`chronon_core::ChrononError::ParamError`] when `schema` fails
259    /// [`validate_postgres_schema_name`]. Returns a storage error when the admin connection or
260    /// DDL fails.
261    pub async fn drop_postgres_schema(url: &str, schema: &str) -> Result<()> {
262        validate_postgres_schema_name(schema)?;
263        let admin = sqlx::postgres::PgPoolOptions::new()
264            .max_connections(1)
265            .connect(url)
266            .await
267            .map_err(|e| map_connect_err("postgres", url, e))?;
268        let ddl = format!("DROP SCHEMA IF EXISTS \"{schema}\" CASCADE");
269        admin.execute(ddl.as_str()).await.map_err(map_err)?;
270        Ok(())
271    }
272}
273
274impl fmt::Debug for SqlSchedulerStore {
275    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
276        f.debug_struct("SqlSchedulerStore")
277            .field("dialect", &self.dialect)
278            .finish_non_exhaustive()
279    }
280}
281
282#[cfg(test)]
283mod tests {
284    use super::{bind_sql, validate_postgres_schema_name, SqlDialect};
285
286    #[test]
287    fn bind_sql_sqlite_passthrough() {
288        let sql = "SELECT * FROM t WHERE id = ? AND name = ?";
289        assert_eq!(bind_sql(SqlDialect::Sqlite, sql), sql);
290    }
291
292    #[test]
293    fn bind_sql_postgres_renumbers_placeholders() {
294        let sql = "UPDATE t SET a = ?, b = ? WHERE id = ?";
295        assert_eq!(
296            bind_sql(SqlDialect::Postgres, sql),
297            "UPDATE t SET a = $1, b = $2 WHERE id = $3"
298        );
299    }
300
301    #[test]
302    fn schema_name_accepts_safe_identifiers() {
303        assert!(validate_postgres_schema_name("bench_cell_1").is_ok());
304        assert!(validate_postgres_schema_name("_tmp").is_ok());
305        assert!(validate_postgres_schema_name("A").is_ok());
306    }
307
308    #[test]
309    fn schema_name_rejects_injection_and_empty() {
310        use chronon_core::ChrononError;
311
312        let too_long = "x".repeat(64);
313        for name in [
314            "",
315            "a\";drop",
316            "evil-name",
317            "1leading",
318            "has space",
319            too_long.as_str(),
320        ] {
321            match validate_postgres_schema_name(name) {
322                Err(ChrononError::ParamError(_)) => {}
323                other => panic!("expected ParamError for {name:?}, got {other:?}"),
324            }
325        }
326    }
327}
chronon_backend_sql_common/backend.rs

chronon_backend_sql_common/
backend.rs
