Skip to main content

toasty_driver_sqlite/
lib.rs

1#![warn(missing_docs)]
2
3//! Toasty driver for [SQLite](https://www.sqlite.org/) using
4//! [`rusqlite`](https://docs.rs/rusqlite).
5//!
6//! Supports both file-backed and in-memory databases.
7//!
8//! # Examples
9//!
10//! ```
11//! use toasty_driver_sqlite::Sqlite;
12//!
13//! // In-memory database
14//! let driver = Sqlite::in_memory();
15//!
16//! // File-backed database
17//! let driver = Sqlite::open("path/to/db.sqlite3");
18//! ```
19
20mod value;
21pub(crate) use value::Value;
22
23use async_trait::async_trait;
24use rusqlite::Connection as RusqliteConnection;
25use std::{
26    borrow::Cow,
27    path::{Path, PathBuf},
28    sync::Arc,
29};
30use toasty_core::{
31    Result, Schema,
32    driver::{
33        Capability, ConnectContext, ConnectionUrl, Driver, ExecResponse, QueryLogConfig,
34        log::QueryLog,
35        operation::{IsolationLevel, Operation, RawSqlRet, Transaction, TypedValue},
36    },
37    schema::{
38        db::{self, Migration, Table},
39        diff,
40    },
41    stmt,
42};
43use toasty_sql::{self as sql};
44
45enum SqlReturn {
46    Count,
47    Infer,
48    Types(Vec<stmt::Type>),
49}
50
51/// A SQLite [`Driver`] that opens connections to a file or in-memory database.
52///
53/// # Examples
54///
55/// ```
56/// use toasty_driver_sqlite::Sqlite;
57///
58/// let driver = Sqlite::in_memory();
59/// ```
60#[derive(Debug)]
61pub enum Sqlite {
62    /// A database stored at a filesystem path.
63    File(PathBuf),
64    /// An ephemeral in-memory database.
65    InMemory,
66}
67
68impl Sqlite {
69    /// Create a new SQLite driver with an arbitrary connection URL
70    pub fn new(url: impl Into<String>) -> Result<Self> {
71        let url_str = url.into();
72        let url = ConnectionUrl::parse(&url_str)?;
73
74        if !url.has_scheme("sqlite") {
75            return Err(toasty_core::Error::invalid_connection_url(format!(
76                "connection URL does not have a `sqlite` scheme; url={url_str}"
77            )));
78        }
79
80        let path = url.file_path()?;
81        if path == Path::new(":memory:") {
82            return Ok(Self::InMemory);
83        }
84
85        Ok(Self::File(path))
86    }
87
88    /// Create an in-memory SQLite database
89    pub fn in_memory() -> Self {
90        Self::InMemory
91    }
92
93    /// Open a SQLite database at the specified file path
94    pub fn open<P: AsRef<Path>>(path: P) -> Self {
95        Self::File(path.as_ref().to_path_buf())
96    }
97}
98
99#[async_trait]
100impl Driver for Sqlite {
101    fn url(&self) -> Cow<'_, str> {
102        match self {
103            Sqlite::InMemory => Cow::Borrowed("sqlite::memory:"),
104            Sqlite::File(path) => Cow::Owned(format!("sqlite:{}", path.display())),
105        }
106    }
107
108    fn capability(&self) -> &'static Capability {
109        &Capability::SQLITE
110    }
111
112    async fn connect(
113        &self,
114        cx: &ConnectContext,
115    ) -> toasty_core::Result<Box<dyn toasty_core::Connection>> {
116        let mut connection = match self {
117            Sqlite::File(path) => Connection::open(path)?,
118            Sqlite::InMemory => Connection::in_memory(),
119        };
120        connection.query_log = cx.query_log;
121        Ok(Box::new(connection))
122    }
123
124    fn max_connections(&self) -> Option<usize> {
125        matches!(self, Self::InMemory).then_some(1)
126    }
127
128    fn generate_migration(&self, schema_diff: &diff::Schema<'_>) -> Migration {
129        let statements = sql::MigrationStatement::from_diff(schema_diff, &Capability::SQLITE);
130
131        let sql_strings: Vec<String> = statements
132            .iter()
133            .map(|stmt| sql::Serializer::sqlite(stmt.schema()).serialize(stmt.statement()))
134            .collect();
135
136        Migration::new_sql_with_breakpoints(&sql_strings)
137    }
138
139    async fn reset_db(&self) -> toasty_core::Result<()> {
140        match self {
141            Sqlite::File(path) => {
142                // Delete the file and recreate it
143                if path.exists() {
144                    std::fs::remove_file(path)
145                        .map_err(toasty_core::Error::driver_operation_failed)?;
146                }
147            }
148            Sqlite::InMemory => {
149                // Nothing to do — each connect() creates a fresh in-memory database
150            }
151        }
152
153        Ok(())
154    }
155}
156
157/// An open connection to a SQLite database.
158#[derive(Debug)]
159pub struct Connection {
160    connection: RusqliteConnection,
161    query_log: QueryLogConfig,
162}
163
164impl Connection {
165    /// Open an in-memory SQLite connection.
166    pub fn in_memory() -> Self {
167        let connection = RusqliteConnection::open_in_memory().unwrap();
168
169        Self {
170            connection,
171            query_log: QueryLogConfig::default(),
172        }
173    }
174
175    /// Open a SQLite connection to a file at `path`.
176    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
177        let connection =
178            RusqliteConnection::open(path).map_err(toasty_core::Error::driver_operation_failed)?;
179        let sqlite = Self {
180            connection,
181            query_log: QueryLogConfig::default(),
182        };
183        Ok(sqlite)
184    }
185
186    fn exec_sql(
187        &mut self,
188        sql_str: &str,
189        typed_params: Vec<TypedValue>,
190        ret: SqlReturn,
191    ) -> Result<ExecResponse> {
192        let mut log = QueryLog::sql(
193            &self.query_log,
194            "sqlite",
195            sql_str,
196            typed_params.iter().map(|tv| &tv.value),
197        );
198        let result = self.exec_sql_inner(sql_str, typed_params, ret, &mut log);
199        log.finish(&result);
200        result
201    }
202
203    fn exec_sql_inner(
204        &mut self,
205        sql_str: &str,
206        typed_params: Vec<TypedValue>,
207        ret: SqlReturn,
208        log: &mut QueryLog<'_>,
209    ) -> Result<ExecResponse> {
210        let mut stmt = self
211            .connection
212            .prepare_cached(sql_str)
213            .map_err(toasty_core::Error::driver_operation_failed)?;
214
215        let params = typed_params
216            .into_iter()
217            .map(|tv| Value::from(tv.value))
218            .collect::<Vec<_>>();
219
220        if matches!(ret, SqlReturn::Count) {
221            let count = stmt
222                .execute(rusqlite::params_from_iter(params.iter()))
223                .map_err(toasty_core::Error::driver_operation_failed)?;
224
225            return Ok(ExecResponse::count(count as _));
226        }
227
228        let mut rows = stmt
229            .query(rusqlite::params_from_iter(params.iter()))
230            .map_err(toasty_core::Error::driver_operation_failed)?;
231
232        let mut values = vec![];
233        let column_count = rows.as_ref().map(|stmt| stmt.column_count()).unwrap_or(0);
234
235        loop {
236            match rows.next() {
237                Ok(Some(row)) => {
238                    let items = match &ret {
239                        SqlReturn::Count => unreachable!(),
240                        SqlReturn::Infer => (0..column_count)
241                            .map(|index| Value::from_sql_infer(row, index).into_inner())
242                            .collect(),
243                        SqlReturn::Types(ret_tys) => ret_tys
244                            .iter()
245                            .enumerate()
246                            .map(|(index, ret_ty)| Value::from_sql(row, index, ret_ty).into_inner())
247                            .collect(),
248                    };
249
250                    values.push(stmt::ValueRecord::from_vec(items).into());
251                }
252                Ok(None) => break,
253                Err(err) => {
254                    return Err(toasty_core::Error::driver_operation_failed(err));
255                }
256            }
257        }
258
259        log.rows(values.len() as u64);
260        Ok(ExecResponse::value_stream(stmt::ValueStream::from_vec(
261            values,
262        )))
263    }
264}
265
266#[async_trait]
267impl toasty_core::driver::Connection for Connection {
268    async fn exec(&mut self, schema: &Arc<Schema>, op: Operation) -> Result<ExecResponse> {
269        tracing::trace!(driver = "sqlite", op = %op.name(), "driver exec");
270
271        let (sql, typed_params, ret_tys) = match op {
272            Operation::Insert(op) => (sql::Statement::from(op.stmt), op.params, op.ret),
273            Operation::QuerySql(op) => (sql::Statement::from(op.stmt), op.params, op.ret),
274            Operation::RawSql(op) => {
275                let ret = match op.ret {
276                    RawSqlRet::None => SqlReturn::Count,
277                    RawSqlRet::Infer => SqlReturn::Infer,
278                    RawSqlRet::Types(types) => SqlReturn::Types(types),
279                };
280                return self.exec_sql(&op.sql, op.params, ret);
281            }
282            Operation::Transaction(mut op) => {
283                if let Transaction::Start { isolation, .. } = &mut op {
284                    if !matches!(isolation, Some(IsolationLevel::Serializable) | None) {
285                        return Err(toasty_core::Error::unsupported_feature(
286                            "SQLite only supports Serializable isolation",
287                        ));
288                    }
289                    *isolation = None;
290                }
291                let sql = sql::Serializer::sqlite(&schema.db).serialize_transaction(&op);
292                self.connection
293                    .execute(&sql, [])
294                    .map_err(toasty_core::Error::driver_operation_failed)?;
295                return Ok(ExecResponse::count(0));
296            }
297            _ => todo!("op={:#?}", op),
298        };
299
300        let ret = match &sql {
301            sql::Statement::Query(stmt) => match &stmt.body {
302                stmt::ExprSet::Select(_) => SqlReturn::Types(ret_tys.unwrap()),
303                _ => todo!(),
304            },
305            sql::Statement::Insert(stmt) => stmt
306                .returning
307                .as_ref()
308                .map(|_| SqlReturn::Types(ret_tys.unwrap()))
309                .unwrap_or(SqlReturn::Count),
310            sql::Statement::Delete(stmt) => stmt
311                .returning
312                .as_ref()
313                .map(|_| SqlReturn::Types(ret_tys.unwrap()))
314                .unwrap_or(SqlReturn::Count),
315            sql::Statement::Update(stmt) => {
316                assert!(stmt.condition.is_none(), "stmt={stmt:#?}");
317                stmt.returning
318                    .as_ref()
319                    .map(|_| SqlReturn::Types(ret_tys.unwrap()))
320                    .unwrap_or(SqlReturn::Count)
321            }
322            _ => SqlReturn::Count,
323        };
324
325        let sql_str = sql::Serializer::sqlite(&schema.db).serialize(&sql);
326        self.exec_sql(&sql_str, typed_params, ret)
327    }
328
329    async fn push_schema(&mut self, schema: &Schema) -> Result<()> {
330        for table in &schema.db.tables {
331            tracing::debug!(table = %table.name, "creating table");
332            self.create_table(&schema.db, table)?;
333        }
334
335        Ok(())
336    }
337
338    async fn applied_migrations(
339        &mut self,
340    ) -> Result<Vec<toasty_core::schema::db::AppliedMigration>> {
341        // Ensure the migrations table exists
342        self.connection
343            .execute(
344                "CREATE TABLE IF NOT EXISTS __toasty_migrations (
345                id INTEGER PRIMARY KEY,
346                name TEXT NOT NULL,
347                applied_at TEXT NOT NULL
348            )",
349                [],
350            )
351            .map_err(toasty_core::Error::driver_operation_failed)?;
352
353        // Query all applied migrations
354        let mut stmt = self
355            .connection
356            .prepare("SELECT id FROM __toasty_migrations ORDER BY applied_at")
357            .map_err(toasty_core::Error::driver_operation_failed)?;
358
359        let rows = stmt
360            .query_map([], |row| {
361                let id: i64 = row.get(0)?;
362                Ok(toasty_core::schema::db::AppliedMigration::new(id as u64))
363            })
364            .map_err(toasty_core::Error::driver_operation_failed)?;
365
366        rows.collect::<rusqlite::Result<Vec<_>>>()
367            .map_err(toasty_core::Error::driver_operation_failed)
368    }
369
370    async fn apply_migration(
371        &mut self,
372        id: u64,
373        name: &str,
374        migration: &toasty_core::schema::db::Migration,
375    ) -> Result<()> {
376        tracing::info!(id = id, name = %name, "applying migration");
377        // Ensure the migrations table exists
378        self.connection
379            .execute(
380                "CREATE TABLE IF NOT EXISTS __toasty_migrations (
381                id INTEGER PRIMARY KEY,
382                name TEXT NOT NULL,
383                applied_at TEXT NOT NULL
384            )",
385                [],
386            )
387            .map_err(toasty_core::Error::driver_operation_failed)?;
388
389        // Start transaction
390        self.connection
391            .execute("BEGIN", [])
392            .map_err(toasty_core::Error::driver_operation_failed)?;
393
394        // Execute each migration statement
395        for statement in migration.statements() {
396            if let Err(e) = self
397                .connection
398                .execute(statement, [])
399                .map_err(toasty_core::Error::driver_operation_failed)
400            {
401                self.connection
402                    .execute("ROLLBACK", [])
403                    .map_err(toasty_core::Error::driver_operation_failed)?;
404                return Err(e);
405            }
406        }
407
408        // Record the migration
409        if let Err(e) = self.connection.execute(
410            "INSERT INTO __toasty_migrations (id, name, applied_at) VALUES (?1, ?2, datetime('now'))",
411            rusqlite::params![id as i64, name],
412        ).map_err(toasty_core::Error::driver_operation_failed) {
413            self.connection.execute("ROLLBACK", []).map_err(toasty_core::Error::driver_operation_failed)?;
414            return Err(e);
415        }
416
417        // Commit transaction
418        self.connection
419            .execute("COMMIT", [])
420            .map_err(toasty_core::Error::driver_operation_failed)?;
421        Ok(())
422    }
423}
424
425impl Connection {
426    fn create_table(&mut self, schema: &db::Schema, table: &Table) -> Result<()> {
427        let serializer = sql::Serializer::sqlite(schema);
428
429        let stmt = serializer.serialize(&sql::Statement::create_table(table, &Capability::SQLITE));
430
431        self.connection
432            .execute(&stmt, [])
433            .map_err(toasty_core::Error::driver_operation_failed)?;
434
435        // Create any indices
436        for index in &table.indices {
437            // The PK has already been created by the table statement
438            if index.primary_key {
439                continue;
440            }
441
442            let stmt = serializer.serialize(&sql::Statement::create_index(index));
443
444            self.connection
445                .execute(&stmt, [])
446                .map_err(toasty_core::Error::driver_operation_failed)?;
447        }
448        Ok(())
449    }
450}