ic_sqlite_vfs/db/
mod.rs

1//! Public SQLite database facade for canister methods.
2//!
3//! Update paths accept only synchronous closures, which prevents holding a DB
4//! transaction across `await`. Query paths open SQLite in read-only/query-only mode.
5
6pub mod connection;
7pub mod migrate;
8pub mod pragmas;
9pub mod row;
10pub mod statement;
11pub mod transaction;
12pub mod value;
13
14use crate::sqlite_vfs::stable_blob;
15use crate::stable::memory::{self, ContextId, DbMemory};
16use crate::stable::meta::Superblock;
17use connection::Connection;
18pub use row::{FromColumn, Row, TextLen};
19pub use stable_blob::ChecksumRefresh;
20use std::cell::RefCell;
21use std::collections::BTreeMap;
22use std::ffi::c_int;
23use std::rc::Rc;
24pub use transaction::UpdateConnection;
25pub use value::{Null, ToSql, Value, NULL};
26
27thread_local! {
28    static READ_CONNECTIONS: RefCell<BTreeMap<ContextId, Rc<Connection>>> = const { RefCell::new(BTreeMap::new()) };
29    static WRITE_CONNECTIONS: RefCell<BTreeMap<ContextId, Rc<Connection>>> = const { RefCell::new(BTreeMap::new()) };
30    static ACTIVE_READ_CONNECTIONS: RefCell<Vec<(ContextId, usize)>> = const { RefCell::new(Vec::new()) };
31}
32
33#[derive(Debug, thiserror::Error)]
34pub enum DbError {
35    #[error("sqlite error {0}: {1}")]
36    Sqlite(c_int, String),
37    #[error("sqlite constraint failed: {0}")]
38    Constraint(String),
39    #[error("query returned no rows")]
40    NotFound,
41    #[error("column {index} has type {actual}, expected {expected}")]
42    TypeMismatch {
43        index: usize,
44        expected: &'static str,
45        actual: &'static str,
46    },
47    #[error("column index {index} out of range for {count} columns")]
48    ColumnOutOfRange { index: usize, count: usize },
49    #[error("stable memory error: {0}")]
50    Stable(#[from] crate::stable::memory::StableMemoryError),
51    #[error("stable memory backend is not initialized; call Db::init(memory) first")]
52    StableMemoryNotInitialized,
53    #[error("stable memory backend is already initialized")]
54    StableMemoryAlreadyInitialized,
55    #[error("cannot mutate database while a query connection is active")]
56    ReadConnectionInUse,
57    #[error("migration version exceeds SQLite INTEGER range: {0}")]
58    MigrationVersionOutOfRange(u64),
59    #[error("duplicate migration version: {0}")]
60    DuplicateMigrationVersion(u64),
61    #[error("migration versions must be strictly increasing: previous={previous}, next={next}")]
62    MigrationVersionOutOfOrder { previous: u64, next: u64 },
63    #[error("SQL contains an interior NUL byte")]
64    InteriorNul,
65    #[error("SQL contains no statement")]
66    EmptySql,
67    #[error("SQL contains trailing text after the first statement")]
68    TrailingSql,
69    #[error("text value too large")]
70    TextTooLarge,
71    #[error("blob value too large")]
72    BlobTooLarge,
73    #[error("too many SQL parameters")]
74    TooManyParameters,
75    #[error("SQL parameter count mismatch: expected {expected}, actual {actual}")]
76    ParameterCountMismatch { expected: usize, actual: usize },
77    #[error("named bind cannot be used with anonymous SQL parameter at index {index}")]
78    AnonymousParameterInNamedBind { index: usize },
79    #[error("SQL parameter not found: {0}")]
80    ParameterNotFound(String),
81}
82
83pub struct Db;
84
85#[derive(Clone, Copy, Debug, Eq, PartialEq)]
86pub struct DbHandle {
87    context: ContextId,
88}
89
90impl Db {
91    pub fn init(memory: DbMemory) -> Result<(), DbError> {
92        let context = memory::init(memory).map_err(|error| match error {
93            crate::stable::memory::StableMemoryError::AlreadyInitialized => {
94                DbError::StableMemoryAlreadyInitialized
95            }
96            crate::stable::memory::StableMemoryError::NotInitialized => {
97                DbError::StableMemoryNotInitialized
98            }
99            error => DbError::Stable(error),
100        })?;
101        clear_read_connection(context);
102        clear_write_connection(context);
103        let handle = DbHandle::from_context(context);
104        let result = handle.initialize();
105        if result.is_err() {
106            clear_read_connection(context);
107            clear_write_connection(context);
108            memory::clear_failed_initialization(context);
109        }
110        result
111    }
112
113    fn default_handle() -> Result<DbHandle, DbError> {
114        memory::default_context()
115            .map(DbHandle::from_context)
116            .ok_or(DbError::StableMemoryNotInitialized)
117    }
118
119    pub fn update<T, F>(f: F) -> Result<T, DbError>
120    where
121        F: FnOnce(&mut UpdateConnection<'_>) -> Result<T, DbError>,
122    {
123        Self::default_handle()?.update(f)
124    }
125
126    pub fn query<T, F>(f: F) -> Result<T, DbError>
127    where
128        F: FnOnce(&Connection) -> Result<T, DbError>,
129    {
130        Self::default_handle()?.query(f)
131    }
132
133    pub fn migrate(migrations: &[migrate::Migration]) -> Result<(), DbError> {
134        Self::default_handle()?.migrate(migrations)
135    }
136
137    pub fn integrity_check() -> Result<String, DbError> {
138        Self::default_handle()?.integrity_check()
139    }
140
141    pub fn db_checksum() -> Result<u64, DbError> {
142        Self::default_handle()?.db_checksum()
143    }
144
145    pub fn refresh_checksum() -> Result<u64, DbError> {
146        Self::default_handle()?.refresh_checksum()
147    }
148
149    pub fn refresh_checksum_chunk(max_bytes: u64) -> Result<ChecksumRefresh, DbError> {
150        Self::default_handle()?.refresh_checksum_chunk(max_bytes)
151    }
152}
153
154impl DbHandle {
155    pub fn init(memory: DbMemory) -> Result<Self, DbError> {
156        let handle = Self::from_context(memory::init_context(memory)?);
157        clear_read_connection(handle.context);
158        clear_write_connection(handle.context);
159        if let Err(error) = handle.initialize() {
160            clear_read_connection(handle.context);
161            clear_write_connection(handle.context);
162            memory::clear_failed_initialization(handle.context);
163            return Err(error);
164        }
165        Ok(handle)
166    }
167
168    fn from_context(context: ContextId) -> Self {
169        Self { context }
170    }
171
172    fn initialize(self) -> Result<(), DbError> {
173        self.with_context(|| {
174            crate::sqlite_vfs::register();
175            Superblock::load()?;
176            stable_blob::ensure_current_layout()?;
177            Ok(())
178        })
179    }
180
181    fn with_context<T>(self, f: impl FnOnce() -> Result<T, DbError>) -> Result<T, DbError> {
182        memory::with_context(self.context, f)
183    }
184
185    pub fn update<T, F>(self, f: F) -> Result<T, DbError>
186    where
187        F: FnOnce(&mut UpdateConnection<'_>) -> Result<T, DbError>,
188    {
189        self.with_context(|| {
190            reject_active_read_connection(self.context)?;
191            clear_read_connection(self.context);
192            let db_size = stable_blob::begin_update()?;
193            let _overlay_guard = OverlayGuard;
194            let connection = write_connection(self.context, db_size)?;
195            let result = transaction::run_immediate(&connection, f);
196            clear_read_connection(self.context);
197            if result.is_err() {
198                clear_write_connection(self.context);
199            }
200            result
201        })
202    }
203
204    pub fn query<T, F>(self, f: F) -> Result<T, DbError>
205    where
206        F: FnOnce(&Connection) -> Result<T, DbError>,
207    {
208        self.with_context(|| with_read_connection(self.context, f))
209    }
210
211    pub fn migrate(self, migrations: &[migrate::Migration]) -> Result<(), DbError> {
212        self.update(|connection| migrate::apply(connection, migrations))?;
213        self.with_context(|| {
214            let target_version = migrations
215                .iter()
216                .map(|migration| migration.version)
217                .max()
218                .unwrap_or(0);
219            let mut block = Superblock::load()?;
220            if block.schema_version < target_version {
221                clear_read_connection(self.context);
222                block.schema_version = target_version;
223                block.store()?;
224            }
225            Ok(())
226        })
227    }
228
229    pub fn integrity_check(self) -> Result<String, DbError> {
230        self.query(|connection| {
231            connection.query_scalar::<String>("PRAGMA integrity_check", crate::params![])
232        })
233    }
234
235    pub fn db_checksum(self) -> Result<u64, DbError> {
236        self.with_context(|| stable_blob::checksum().map_err(DbError::from))
237    }
238
239    pub fn refresh_checksum(self) -> Result<u64, DbError> {
240        self.with_context(|| {
241            reject_active_read_connection(self.context)?;
242            clear_read_connection(self.context);
243            stable_blob::refresh_checksum().map_err(DbError::from)
244        })
245    }
246
247    pub fn refresh_checksum_chunk(self, max_bytes: u64) -> Result<ChecksumRefresh, DbError> {
248        self.with_context(|| {
249            reject_active_read_connection(self.context)?;
250            clear_read_connection(self.context);
251            stable_blob::refresh_checksum_chunk(max_bytes).map_err(DbError::from)
252        })
253    }
254}
255
256fn write_connection(context: ContextId, db_size: u64) -> Result<Rc<Connection>, DbError> {
257    WRITE_CONNECTIONS.with(|slot| {
258        let cached = { slot.borrow().get(&context).cloned() };
259        if let Some(connection) = cached {
260            return Ok(connection);
261        }
262        let connection = if db_size == 0 {
263            connection::open_read_write()?
264        } else {
265            connection::open_read_write_existing()?
266        };
267        let connection = Rc::new(connection);
268        slot.borrow_mut().insert(context, Rc::clone(&connection));
269        Ok(connection)
270    })
271}
272
273fn with_read_connection<T>(
274    context: ContextId,
275    f: impl FnOnce(&Connection) -> Result<T, DbError>,
276) -> Result<T, DbError> {
277    READ_CONNECTIONS.with(|slot| {
278        let cached = { slot.borrow().get(&context).cloned() };
279        let connection = if let Some(connection) = cached {
280            connection
281        } else {
282            let connection = Rc::new(connection::open_read_only()?);
283            slot.borrow_mut().insert(context, Rc::clone(&connection));
284            connection
285        };
286        let _guard = ReadGuard::enter(context);
287        f(&connection)
288    })
289}
290
291fn reject_active_read_connection(context: ContextId) -> Result<(), DbError> {
292    ACTIVE_READ_CONNECTIONS.with(|slot| {
293        let slot = slot.borrow();
294        let active = active_read_index(&slot, context)
295            .map(|index| slot[index].1)
296            .unwrap_or(0);
297        if active == 0 {
298            Ok(())
299        } else {
300            Err(DbError::ReadConnectionInUse)
301        }
302    })
303}
304
305fn clear_read_connection(context: ContextId) {
306    READ_CONNECTIONS.with(|slot| {
307        slot.borrow_mut().remove(&context);
308    });
309}
310
311fn clear_write_connection(context: ContextId) {
312    WRITE_CONNECTIONS.with(|slot| {
313        slot.borrow_mut().remove(&context);
314    });
315}
316
317struct ReadGuard {
318    context: ContextId,
319}
320
321impl ReadGuard {
322    fn enter(context: ContextId) -> Self {
323        ACTIVE_READ_CONNECTIONS.with(|slot| {
324            let mut slot = slot.borrow_mut();
325            if slot.is_empty() {
326                slot.push((context, 1));
327                return;
328            }
329            if let Some(index) = active_read_index(&slot, context) {
330                slot[index].1 += 1;
331            } else {
332                slot.push((context, 1));
333            }
334        });
335        Self { context }
336    }
337}
338
339impl Drop for ReadGuard {
340    fn drop(&mut self) {
341        ACTIVE_READ_CONNECTIONS.with(|slot| {
342            let mut slot = slot.borrow_mut();
343            if slot.len() == 1 && slot[0].0 == self.context {
344                let depth = &mut slot[0].1;
345                *depth = depth.saturating_sub(1);
346                if *depth == 0 {
347                    slot.clear();
348                }
349                return;
350            }
351            let Some(index) = active_read_index(&slot, self.context) else {
352                return;
353            };
354            let depth = &mut slot[index].1;
355            *depth = depth.saturating_sub(1);
356            if *depth == 0 {
357                slot.swap_remove(index);
358            }
359        });
360    }
361}
362
363fn active_read_index(entries: &[(ContextId, usize)], context: ContextId) -> Option<usize> {
364    entries
365        .iter()
366        .position(|(stored_context, _)| *stored_context == context)
367}
368
369struct OverlayGuard;
370
371impl Drop for OverlayGuard {
372    fn drop(&mut self) {
373        stable_blob::rollback_update();
374    }
375}
ic_sqlite_vfs/db/mod.rs

ic_sqlite_vfs/db/
mod.rs
