Skip to main content

boson_backend_sql_common/
backend.rs

1use std::fmt;
2
3use boson_core::Result;
4use sqlx::{Executor, Pool, Postgres, Sqlite};
5
6use crate::enqueue_rate::EnqueueRateLimiter;
7use crate::error_map::map_err;
8use crate::schema;
9
10/// `SQLite` uses `?` placeholders; `PostgreSQL` uses `$1`, `$2`, …
11pub fn bind_sql(dialect: SqlDialect, sql: &str) -> String {
12    match dialect {
13        SqlDialect::Sqlite => sql.to_string(),
14        SqlDialect::Postgres => {
15            let mut out = String::with_capacity(sql.len());
16            let mut n = 1u32;
17            for ch in sql.chars() {
18                if ch == '?' {
19                    out.push('$');
20                    out.push_str(&n.to_string());
21                    n += 1;
22                } else {
23                    out.push(ch);
24                }
25            }
26            out
27        }
28    }
29}
30
31/// SQL dialect for query variants.
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum SqlDialect {
34    /// `PostgreSQL`.
35    Postgres,
36    /// `SQLite`.
37    Sqlite,
38}
39
40/// Connection pool for a SQL backend.
41#[derive(Clone)]
42pub enum SqlPool {
43    /// `SQLite` pool.
44    Sqlite(Pool<Sqlite>),
45    /// `PostgreSQL` pool.
46    Postgres(Pool<Postgres>),
47}
48
49impl fmt::Debug for SqlPool {
50    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
51        match self {
52            Self::Sqlite(_) => f.debug_tuple("SqlPool::Sqlite").finish(),
53            Self::Postgres(_) => f.debug_tuple("SqlPool::Postgres").finish(),
54        }
55    }
56}
57
58/// SQL-backed queue backend (`PostgreSQL` or `SQLite`).
59pub struct SqlQueueBackend {
60    pub(crate) pool: SqlPool,
61    pub(crate) dialect: SqlDialect,
62    pub(crate) enqueue_rate: EnqueueRateLimiter,
63}
64
65impl SqlQueueBackend {
66    /// Open a `SQLite` pool, bootstrap schema, and return a backend.
67    ///
68    /// # Errors
69    ///
70    /// Returns a backend error if the pool connection or schema bootstrap fails.
71    pub async fn connect_sqlite(url: &str) -> Result<Self> {
72        let pool = sqlx::sqlite::SqlitePoolOptions::new()
73            .max_connections(5)
74            .connect(url)
75            .await
76            .map_err(|e| map_err(&e))?;
77        Self::from_sqlite_pool(pool).await
78    }
79
80    /// Open a `PostgreSQL` pool, bootstrap schema, and return a backend.
81    ///
82    /// # Errors
83    ///
84    /// Returns a backend error if the pool connection or schema bootstrap fails.
85    pub async fn connect_postgres(url: &str) -> Result<Self> {
86        let pool = sqlx::postgres::PgPoolOptions::new()
87            .max_connections(5)
88            .connect(url)
89            .await
90            .map_err(|e| map_err(&e))?;
91        Self::from_postgres_pool(pool).await
92    }
93
94    /// Connect to `PostgreSQL` with an isolated schema for parallel tests.
95    ///
96    /// # Errors
97    ///
98    /// Returns a backend error if schema creation, pool connection, or schema bootstrap fails.
99    pub async fn connect_postgres_isolated(url: &str, schema: &str) -> Result<Self> {
100        let admin = sqlx::postgres::PgPoolOptions::new()
101            .max_connections(1)
102            .connect(url)
103            .await
104            .map_err(|e| map_err(&e))?;
105        let ddl = format!("CREATE SCHEMA IF NOT EXISTS \"{schema}\"");
106        admin.execute(ddl.as_str()).await.map_err(|e| map_err(&e))?;
107        drop(admin);
108
109        let schema = schema.to_string();
110        let pool = sqlx::postgres::PgPoolOptions::new()
111            .max_connections(5)
112            .after_connect(move |conn, _meta| {
113                let schema = schema.clone();
114                Box::pin(async move {
115                    let sql = format!("SET search_path TO \"{schema}\"");
116                    sqlx::query(&sql).execute(conn).await?;
117                    Ok(())
118                })
119            })
120            .connect(url)
121            .await
122            .map_err(|e| map_err(&e))?;
123        Self::from_postgres_pool(pool).await
124    }
125
126    /// Wrap an existing `SQLite` pool (schema bootstrap runs).
127    ///
128    /// # Errors
129    ///
130    /// Returns a backend error if schema bootstrap fails.
131    pub async fn from_sqlite_pool(pool: Pool<Sqlite>) -> Result<Self> {
132        let backend = Self {
133            pool: SqlPool::Sqlite(pool),
134            dialect: SqlDialect::Sqlite,
135            enqueue_rate: EnqueueRateLimiter::new(),
136        };
137        schema::ensure_schema(&backend).await?;
138        Ok(backend)
139    }
140
141    /// Wrap an existing `PostgreSQL` pool (schema bootstrap runs).
142    ///
143    /// # Errors
144    ///
145    /// Returns a backend error if schema bootstrap fails.
146    pub async fn from_postgres_pool(pool: Pool<Postgres>) -> Result<Self> {
147        let backend = Self {
148            pool: SqlPool::Postgres(pool),
149            dialect: SqlDialect::Postgres,
150            enqueue_rate: EnqueueRateLimiter::new(),
151        };
152        schema::ensure_schema(&backend).await?;
153        Ok(backend)
154    }
155
156    /// Underlying connection pool.
157    #[must_use]
158    pub const fn pool(&self) -> &SqlPool {
159        &self.pool
160    }
161
162    /// Engine dialect.
163    #[must_use]
164    pub const fn dialect(&self) -> SqlDialect {
165        self.dialect
166    }
167
168    pub(crate) async fn run_ddl(&self, ddl: &str) -> Result<()> {
169        match &self.pool {
170            SqlPool::Sqlite(pool) => {
171                pool.execute(ddl).await.map_err(|e| map_err(&e))?;
172            }
173            SqlPool::Postgres(pool) => {
174                pool.execute(ddl).await.map_err(|e| map_err(&e))?;
175            }
176        }
177        Ok(())
178    }
179}
180
181impl fmt::Debug for SqlQueueBackend {
182    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
183        f.debug_struct("SqlQueueBackend")
184            .field("dialect", &self.dialect)
185            .finish_non_exhaustive()
186    }
187}