toasty_driver_sqlite/
lib.rs

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