uqa_storage_sqlite/
connection.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Pooled `SQLite` connections with explicit session and transaction affinity.
8//!
9//! A [`ManagedConnection`] is a logical session. Clones share that session so
10//! catalog, document, inverted-index, and vector stores participate in the
11//! same explicit transaction. [`ManagedConnection::new_session`] creates an
12//! isolated session over the same physical connection pool. Outside explicit
13//! transactions operations check out independent connections, allowing WAL
14//! readers to make real concurrent progress.
15
16use std::path::{Path, PathBuf};
17use std::sync::Arc;
18
19use parking_lot::{Condvar, Mutex, RwLock};
20use rusqlite::{Connection, OpenFlags};
21use uqa_storage::StorageEncryptionKey;
22
23use crate::compressed_vfs::{self, SQLiteCompressedContainerAnchor, SQLiteCompressionOptions};
24
25#[derive(Debug, thiserror::Error)]
26pub enum SQLiteError {
27    #[error(transparent)]
28    Memory(#[from] uqa_core::memory::MemoryError),
29    #[error(transparent)]
30    Cancelled(#[from] uqa_core::QueryCancelled),
31    #[error("text analysis failed: {0}")]
32    Analysis(#[from] uqa_analysis::AnalysisError),
33    #[error("sqlite error: {0}")]
34    SQLite(#[from] rusqlite::Error),
35    #[error("encryption key must not be empty")]
36    EmptyEncryptionKey,
37    #[error("database requires an encryption key")]
38    EncryptionKeyRequired,
39    #[error("database is not encrypted but an encryption key was provided")]
40    NotEncrypted,
41    #[error("auxiliary database requires DELETE journaling, found {0}")]
42    AuxiliaryJournalMode(String),
43    #[error("compressed sqlite container error: {0}")]
44    CompressedContainer(String),
45    #[error("io error: {0}")]
46    Io(#[from] std::io::Error),
47    #[error("catalog migration {version} failed: {source}")]
48    Migration {
49        version: u32,
50        #[source]
51        source: rusqlite::Error,
52    },
53    #[error("invalid persisted catalog schema version `{0}`")]
54    InvalidSchemaVersion(String),
55    #[error("catalog schema version {found} is newer than this engine supports ({supported})")]
56    UnsupportedSchemaVersion { found: u32, supported: u32 },
57    #[error("corrupt document blob for `{table}` doc {doc_id} field `{field}`: {reason}")]
58    CorruptDocumentBlob {
59        table: String,
60        doc_id: u64,
61        field: String,
62        reason: String,
63    },
64    #[error("payload serialization failed: {0}")]
65    Serde(#[from] serde_json::Error),
66    #[error("storage backend error: {0}")]
67    StorageBackend(String),
68    #[error("transaction already active for this sqlite session")]
69    TransactionAlreadyActive,
70    #[error("no active transaction for this sqlite session")]
71    NoActiveTransaction,
72    #[error("sqlite transaction was aborted by an earlier storage error: {0}")]
73    TransactionAborted(String),
74    #[error("sqlite session cleanup failed: {0}")]
75    SessionCleanupFailed(String),
76    #[error("sqlite connection-pool checkout lost its connection")]
77    MissingCheckedOutConnection,
78}
79
80pub type Result<T> = std::result::Result<T, SQLiteError>;
81
82const MIN_POOL_CONNECTIONS: usize = 4;
83const MAX_POOL_CONNECTIONS: usize = 32;
84
85#[derive(Clone)]
86enum ConnectionSpec {
87    File {
88        path: PathBuf,
89        key: Option<StorageEncryptionKey>,
90    },
91    Auxiliary {
92        path: PathBuf,
93        key: Option<StorageEncryptionKey>,
94    },
95    Compressed {
96        path: PathBuf,
97        compression: SQLiteCompressionOptions,
98        key: Option<StorageEncryptionKey>,
99    },
100    Memory,
101}
102
103impl ConnectionSpec {
104    fn open(&self, initialize_database: bool) -> Result<Connection> {
105        match self {
106            Self::File { path, key } | Self::Auxiliary { path, key } => {
107                let conn = Connection::open(path)?;
108                if let Some(key) = key {
109                    ManagedConnection::apply_encryption_key(&conn, key.expose_secret())?;
110                }
111                if matches!(self, Self::Auxiliary { .. }) {
112                    // Registry writers already serialize their transactions.
113                    // Retain the original journal mode instead of racing to
114                    // promote a new registry to WAL during concurrent opens.
115                    let mode: String =
116                        conn.pragma_query_value(None, "journal_mode", |row| row.get(0))?;
117                    if mode != "delete" {
118                        return Err(SQLiteError::AuxiliaryJournalMode(mode));
119                    }
120                    ManagedConnection::configure_rollback_connection(&conn)?;
121                } else {
122                    if initialize_database {
123                        ManagedConnection::enable_wal(&conn)?;
124                    }
125                    ManagedConnection::configure_wal_connection(&conn)?;
126                }
127                Ok(conn)
128            }
129            Self::Compressed {
130                path, compression, ..
131            } => {
132                let conn = Connection::open_with_flags_and_vfs(
133                    path,
134                    OpenFlags::default(),
135                    compressed_vfs::VFS_NAME,
136                )?;
137                if initialize_database {
138                    conn.pragma_update(None, "page_size", compression.page_size)?;
139                    ManagedConnection::enable_compressed_journal(&conn)?;
140                }
141                ManagedConnection::configure_rollback_connection(&conn)?;
142                Ok(conn)
143            }
144            Self::Memory => {
145                let conn = Connection::open_in_memory()?;
146                if initialize_database {
147                    ManagedConnection::enable_wal(&conn)?;
148                }
149                ManagedConnection::configure_wal_connection(&conn)?;
150                Ok(conn)
151            }
152        }
153    }
154}
155
156struct PoolState {
157    idle: Vec<Connection>,
158    open: usize,
159}
160
161struct ConnectionPool {
162    spec: ConnectionSpec,
163    max_connections: usize,
164    state: Mutex<PoolState>,
165    available: Condvar,
166    /// Stable, never-mutating connection used for `PRAGMA data_version`.
167    /// Every logical session over this pool must compare versions on this
168    /// same connection, and encrypted databases must not repeat key
169    /// derivation merely to create a request-local change monitor.
170    data_version_monitor: Mutex<Option<Connection>>,
171}
172
173impl ConnectionPool {
174    fn new(spec: ConnectionSpec, initial: Connection, max_connections: usize) -> Arc<Self> {
175        Arc::new(Self {
176            spec,
177            max_connections: max_connections.max(1),
178            state: Mutex::new(PoolState {
179                idle: vec![initial],
180                open: 1,
181            }),
182            available: Condvar::new(),
183            data_version_monitor: Mutex::new(None),
184        })
185    }
186
187    fn checkout(self: &Arc<Self>) -> Result<PooledConnection> {
188        loop {
189            let mut state = self.state.lock();
190            if let Some(connection) = state.idle.pop() {
191                return Ok(PooledConnection {
192                    pool: Arc::clone(self),
193                    connection: Some(connection),
194                });
195            }
196            if state.open < self.max_connections {
197                state.open += 1;
198                drop(state);
199                return match self.spec.open(false) {
200                    Ok(connection) => Ok(PooledConnection {
201                        pool: Arc::clone(self),
202                        connection: Some(connection),
203                    }),
204                    Err(error) => {
205                        let mut state = self.state.lock();
206                        state.open -= 1;
207                        self.available.notify_one();
208                        Err(error)
209                    }
210                };
211            }
212            self.available.wait(&mut state);
213        }
214    }
215
216    fn checkin(&self, connection: Connection) {
217        self.state.lock().idle.push(connection);
218        self.available.notify_one();
219    }
220
221    fn discard(&self) {
222        let mut state = self.state.lock();
223        state.open -= 1;
224        self.available.notify_one();
225    }
226}
227
228pub(crate) struct PooledConnection {
229    pool: Arc<ConnectionPool>,
230    connection: Option<Connection>,
231}
232
233impl PooledConnection {
234    pub(crate) fn connection(&self) -> Result<&Connection> {
235        self.connection
236            .as_ref()
237            .ok_or(SQLiteError::MissingCheckedOutConnection)
238    }
239
240    pub(crate) fn connection_mut(&mut self) -> Result<&mut Connection> {
241        self.connection
242            .as_mut()
243            .ok_or(SQLiteError::MissingCheckedOutConnection)
244    }
245}
246
247impl Drop for PooledConnection {
248    fn drop(&mut self) {
249        let Some(connection) = self.connection.take() else {
250            return;
251        };
252        let reusable = connection.is_autocommit() || connection.execute_batch("ROLLBACK").is_ok();
253        if reusable {
254            self.pool.checkin(connection);
255        } else {
256            self.pool.discard();
257        }
258    }
259}
260
261struct SessionState {
262    /// Read guards cover ordinary operations. Transaction lifecycle calls take
263    /// the write guard, making BEGIN/COMMIT/ROLLBACK linearizable with respect
264    /// to every operation issued through the same logical session.
265    gate: RwLock<()>,
266    transaction: Mutex<Option<PooledConnection>>,
267    transaction_failure: Mutex<Option<String>>,
268    cleanup_failure: Mutex<Option<String>>,
269}
270
271impl SessionState {
272    fn new() -> Self {
273        Self {
274            gate: RwLock::new(()),
275            transaction: Mutex::new(None),
276            transaction_failure: Mutex::new(None),
277            cleanup_failure: Mutex::new(None),
278        }
279    }
280}
281
282impl Drop for SessionState {
283    fn drop(&mut self) {
284        // `PooledConnection::drop` performs the rollback and discards the
285        // physical connection when rollback itself fails. Taking the pinned
286        // handle here therefore cannot return a broken connection to the pool.
287        self.transaction.get_mut().take();
288    }
289}
290
291/// Logical `SQLite` session backed by a bounded physical connection pool.
292/// Cloning preserves session/transaction affinity; call [`Self::new_session`]
293/// for an independently isolated transaction context.
294#[derive(Clone)]
295pub struct ManagedConnection {
296    pool: Arc<ConnectionPool>,
297    session: Arc<SessionState>,
298}
299
300impl ManagedConnection {
301    fn surface_cleanup_failure(&self) -> Result<()> {
302        if let Some(error) = self.session.cleanup_failure.lock().take() {
303            return Err(SQLiteError::SessionCleanupFailed(error));
304        }
305        Ok(())
306    }
307
308    pub fn open(path: &Path) -> Result<Self> {
309        if path == Path::new(":memory:") {
310            return Self::open_in_memory();
311        }
312        Self::open_with_optional_key(path, None)
313    }
314
315    /// Return the backing database path for a file-backed connection.
316    #[must_use]
317    pub fn database_path(&self) -> Option<&Path> {
318        match &self.pool.spec {
319            ConnectionSpec::File { path, .. }
320            | ConnectionSpec::Auxiliary { path, .. }
321            | ConnectionSpec::Compressed { path, .. } => Some(path),
322            ConnectionSpec::Memory => None,
323        }
324    }
325
326    /// Retain the database credential for encrypted auxiliary storage, including
327    /// compressed containers whose main-file encryption lives in the VFS.
328    #[must_use]
329    pub fn auxiliary_encryption_key(&self) -> Option<StorageEncryptionKey> {
330        match &self.pool.spec {
331            ConnectionSpec::File { key, .. }
332            | ConnectionSpec::Auxiliary { key, .. }
333            | ConnectionSpec::Compressed { key, .. } => key.clone(),
334            ConnectionSpec::Memory => None,
335        }
336    }
337
338    /// Lease an independent physical connection without joining this logical
339    /// session's transaction. The pool retains encryption and connection policy.
340    pub fn lease_connection(&self) -> Result<crate::SQLiteConnectionLease> {
341        self.pool.checkout().map(crate::SQLiteConnectionLease)
342    }
343
344    /// Open database-owned auxiliary storage with its inherited credential and
345    /// original DELETE journaling. Concurrent first opens never change journal
346    /// modes. An incompatible existing mode is rejected without rewriting it.
347    pub fn open_auxiliary(path: &Path, key: Option<StorageEncryptionKey>) -> Result<Self> {
348        if path == Path::new(":memory:") || path.as_os_str().is_empty() {
349            return Err(SQLiteError::StorageBackend(
350                "auxiliary storage requires a database file".into(),
351            ));
352        }
353        Self::from_spec(
354            ConnectionSpec::Auxiliary {
355                path: path.to_path_buf(),
356                key,
357            },
358            default_pool_connections(),
359        )
360    }
361
362    pub fn open_encrypted(path: &Path, key: &str) -> Result<Self> {
363        Self::open_with_optional_key(path, Some(key))
364    }
365
366    pub fn open_compressed(path: &Path, compression: SQLiteCompressionOptions) -> Result<Self> {
367        Self::open_compressed_with_optional_key(path, compression, None, None)
368    }
369
370    pub fn open_compressed_encrypted(
371        path: &Path,
372        key: &str,
373        compression: SQLiteCompressionOptions,
374    ) -> Result<Self> {
375        if key.is_empty() {
376            return Err(SQLiteError::EmptyEncryptionKey);
377        }
378        Self::open_compressed_with_optional_key(path, compression, Some(key), None)
379    }
380
381    /// Open an encrypted compressed database while enforcing an exact trusted
382    /// external state anchor in the VFS main-file open path.
383    pub fn open_compressed_encrypted_with_anchor(
384        path: &Path,
385        key: &str,
386        compression: SQLiteCompressionOptions,
387        trusted_anchor: SQLiteCompressedContainerAnchor,
388    ) -> Result<Self> {
389        if key.is_empty() {
390            return Err(SQLiteError::EmptyEncryptionKey);
391        }
392        Self::open_compressed_with_optional_key(path, compression, Some(key), Some(trusted_anchor))
393    }
394
395    fn open_with_optional_key(path: &Path, key: Option<&str>) -> Result<Self> {
396        let spec = ConnectionSpec::File {
397            path: path.to_path_buf(),
398            key: key.map(StorageEncryptionKey::new),
399        };
400        Self::from_spec(spec, default_pool_connections())
401    }
402
403    fn open_compressed_with_optional_key(
404        path: &Path,
405        compression: SQLiteCompressionOptions,
406        key: Option<&str>,
407        trusted_anchor: Option<SQLiteCompressedContainerAnchor>,
408    ) -> Result<Self> {
409        let compression = compression
410            .validate()
411            .map_err(SQLiteError::CompressedContainer)?;
412        match trusted_anchor {
413            Some(anchor) => compressed_vfs::register_database_with_anchor(
414                path,
415                compression,
416                key.ok_or(SQLiteError::EncryptionKeyRequired)?,
417                anchor,
418            ),
419            None => compressed_vfs::register_database(path, compression, key),
420        }
421        .map_err(SQLiteError::CompressedContainer)?;
422        let spec = ConnectionSpec::Compressed {
423            path: path.to_path_buf(),
424            compression,
425            key: key.map(StorageEncryptionKey::new),
426        };
427        Self::from_spec(spec, default_pool_connections())
428    }
429
430    pub fn open_in_memory() -> Result<Self> {
431        // Independent `:memory:` connections do not share a database. Keep a
432        // one-connection pool for this special target; file-backed databases
433        // use the real multi-connection pool.
434        Self::from_spec(ConnectionSpec::Memory, 1)
435    }
436
437    fn from_spec(spec: ConnectionSpec, max_connections: usize) -> Result<Self> {
438        let initial = spec.open(true)?;
439        Ok(Self {
440            pool: ConnectionPool::new(spec, initial, max_connections),
441            session: Arc::new(SessionState::new()),
442        })
443    }
444
445    fn apply_encryption_key(conn: &Connection, key: &str) -> Result<()> {
446        if key.is_empty() {
447            return Err(SQLiteError::EmptyEncryptionKey);
448        }
449        conn.pragma_update(None, "key", key)?;
450        Ok(())
451    }
452
453    fn enable_wal(conn: &Connection) -> Result<()> {
454        conn.pragma_update(None, "journal_mode", "WAL")?;
455        Ok(())
456    }
457
458    fn configure_wal_connection(conn: &Connection) -> Result<()> {
459        // 5s busy timeout absorbs short contention without surfacing
460        // SQLITE_BUSY; long contention is a real bug.
461        conn.busy_timeout(std::time::Duration::from_secs(5))?;
462        // Synchronous=NORMAL is the recommended pairing with WAL: safe
463        // against power loss, faster than FULL.
464        conn.pragma_update(None, "synchronous", "NORMAL")?;
465        // Foreign keys are off by default; turn on so per-table cleanup
466        // can use ON DELETE CASCADE later.
467        conn.pragma_update(None, "foreign_keys", "ON")?;
468        Ok(())
469    }
470
471    fn enable_compressed_journal(conn: &Connection) -> Result<()> {
472        // The compressed VFS implements the byte-addressed database file.
473        // Rollback journals stay raw because they are short-lived commit
474        // machinery; compressing them only adds autocommit write amplification.
475        // WAL requires shared-memory VFS methods, so compressed databases use
476        // SQLite's rollback journal and keep temp storage in memory.
477        conn.pragma_update(None, "journal_mode", "DELETE")?;
478        Ok(())
479    }
480
481    fn configure_rollback_connection(conn: &Connection) -> Result<()> {
482        conn.busy_timeout(std::time::Duration::from_secs(5))?;
483        conn.pragma_update(None, "synchronous", "FULL")?;
484        conn.pragma_update(None, "temp_store", "MEMORY")?;
485        conn.pragma_update(None, "foreign_keys", "ON")?;
486        Ok(())
487    }
488
489    /// Create an independent logical session over the same database pool.
490    /// Explicit transactions started on either session are isolated and never
491    /// capture operations issued through the other session.
492    #[must_use]
493    pub fn new_session(&self) -> Self {
494        Self {
495            pool: Arc::clone(&self.pool),
496            session: Arc::new(SessionState::new()),
497        }
498    }
499
500    /// Whether this logical session currently owns a pinned transaction
501    /// connection.
502    #[must_use]
503    pub fn in_transaction(&self) -> bool {
504        let _gate = self.session.gate.read();
505        self.session.transaction.lock().is_some()
506    }
507
508    /// Whether the currently pinned transaction has upgraded to a write
509    /// transaction. Callers use this to enforce read-only execution
510    /// boundaries before COMMIT rather than inferring writes from SQL shape.
511    pub fn transaction_has_written(&self) -> Result<bool> {
512        let _gate = self.session.gate.read();
513        let transaction = self.session.transaction.lock();
514        let transaction = transaction
515            .as_ref()
516            .ok_or(SQLiteError::NoActiveTransaction)?;
517        Ok(matches!(
518            transaction.connection()?.transaction_state(Some("main"))?,
519            rusqlite::TransactionState::Write
520        ))
521    }
522
523    /// Whether this session's independent [`Self::data_version`] monitor can
524    /// read without contending with the currently pinned transaction.
525    ///
526    /// A rollback-journal writer's pending lock blocks new readers while waiting for existing readers to finish. A rollback-journal read transaction must therefore also avoid the independent monitor: its own shared lock may be preventing that waiting writer from proceeding. Callers refresh through the pinned connection instead. WAL sessions and sessions without a pinned transaction permit the independent monitor.
527    pub fn data_version_monitor_is_nonblocking(&self) -> Result<bool> {
528        if self.supports_concurrent_pinned_read_and_write() {
529            return Ok(true);
530        }
531        let _gate = self.session.gate.read();
532        Ok(self.session.transaction.lock().is_none())
533    }
534
535    /// Whether one pooled connection may retain a read snapshot while another writes. Ordinary plain and encrypted `SQLite` databases use WAL; compressed containers and auxiliary files use rollback journaling and therefore require a detached engine snapshot before writer promotion.
536    #[must_use]
537    pub fn supports_concurrent_pinned_read_and_write(&self) -> bool {
538        !matches!(
539            &self.pool.spec,
540            ConnectionSpec::Compressed { .. } | ConnectionSpec::Auxiliary { .. }
541        )
542    }
543
544    /// Database change counter observed on one stable connection shared by
545    /// every logical session over this pool. The value changes when another
546    /// `SQLite` connection commits. In-memory databases have no independent
547    /// connections and therefore return `None`.
548    pub fn data_version(&self) -> Result<Option<u64>> {
549        if matches!(&self.pool.spec, ConnectionSpec::Memory) {
550            return Ok(None);
551        }
552        let mut monitor = self.pool.data_version_monitor.lock();
553        if monitor.is_none() {
554            *monitor = Some(self.pool.spec.open(false)?);
555        }
556        let monitor = monitor.as_ref().ok_or_else(|| {
557            SQLiteError::StorageBackend(
558                "data-version monitor was not initialized after opening it".into(),
559            )
560        })?;
561        let version: i64 = monitor.pragma_query_value(None, "data_version", |row| row.get(0))?;
562        let version = u64::try_from(version).map_err(|_| {
563            SQLiteError::StorageBackend(format!(
564                "SQLite returned a negative PRAGMA data_version: {version}"
565            ))
566        })?;
567        Ok(Some(version))
568    }
569
570    /// Establish the database snapshot for the active transaction without
571    /// depending on the caller's first user query. `BEGIN DEFERRED` alone does
572    /// not start a read transaction, so a writer could otherwise commit after
573    /// the engine checks its cache generations but before the first catalog
574    /// read. Reading `sqlite_schema` is database-wide and keeps the operation
575    /// independent of any application table.
576    pub fn pin_transaction_snapshot(&self) -> Result<()> {
577        self.with(|connection| {
578            if connection.is_autocommit() {
579                return Err(SQLiteError::NoActiveTransaction);
580            }
581            let _: i64 =
582                connection.query_row("SELECT COUNT(*) FROM sqlite_schema", [], |row| row.get(0))?;
583            Ok(())
584        })
585    }
586
587    /// Run a closure using this session. Outside a transaction the closure
588    /// checks out a pooled connection; inside a transaction every clone is
589    /// routed to the session's pinned connection.
590    pub fn with<R>(&self, f: impl FnOnce(&Connection) -> Result<R>) -> Result<R> {
591        self.surface_cleanup_failure()?;
592        let _gate = self.session.gate.read();
593        let transaction = self.session.transaction.lock();
594        if let Some(connection) = transaction.as_ref() {
595            if let Some(error) = self.session.transaction_failure.lock().as_ref() {
596                return Err(SQLiteError::TransactionAborted(error.clone()));
597            }
598            let result = f(connection.connection()?);
599            if let Err(error) = &result {
600                let mut failure = self.session.transaction_failure.lock();
601                if failure.is_none() {
602                    *failure = Some(error.to_string());
603                }
604            }
605            return result;
606        }
607        drop(transaction);
608        let connection = self.pool.checkout()?;
609        f(connection.connection()?)
610    }
611
612    pub fn with_mut<R>(&self, f: impl FnOnce(&mut Connection) -> Result<R>) -> Result<R> {
613        self.surface_cleanup_failure()?;
614        let _gate = self.session.gate.read();
615        let mut transaction = self.session.transaction.lock();
616        if let Some(connection) = transaction.as_mut() {
617            if let Some(error) = self.session.transaction_failure.lock().as_ref() {
618                return Err(SQLiteError::TransactionAborted(error.clone()));
619            }
620            let result = f(connection.connection_mut()?);
621            if let Err(error) = &result {
622                let mut failure = self.session.transaction_failure.lock();
623                if failure.is_none() {
624                    *failure = Some(error.to_string());
625                }
626            }
627            return result;
628        }
629        drop(transaction);
630        let mut connection = self.pool.checkout()?;
631        f(connection.connection_mut()?)
632    }
633
634    /// Rewrite the `SQLite` database into its minimum-sized file. `SQLite` requires `VACUUM` to run in autocommit mode, so the session write gate makes the transaction check and maintenance command one atomic session operation.
635    pub fn vacuum(&self) -> Result<()> {
636        self.surface_cleanup_failure()?;
637        let _gate = self.session.gate.write();
638        if self.session.transaction.lock().is_some() {
639            return Err(SQLiteError::TransactionAlreadyActive);
640        }
641        let connection = self.pool.checkout()?;
642        connection.connection()?.execute_batch("VACUUM")?;
643        Ok(())
644    }
645
646    /// Open an explicit (non-deferred) transaction. Subsequent
647    /// auto-commit hosts (catalog writes, FTS index updates, ...) all
648    /// flow through the same connection so the transaction enclosing
649    /// them is honoured. Use [`Self::commit_transaction`] /
650    /// [`Self::rollback_transaction`] / [`Self::savepoint`] etc. for
651    /// the lifecycle.
652    pub fn begin_transaction(&self) -> Result<()> {
653        self.begin_transaction_with("BEGIN IMMEDIATE")
654    }
655
656    /// Open a deferred transaction. Read-only SQL statements use this mode
657    /// so WAL readers do not take the single writer reservation; if a scalar
658    /// routine performs a write, `SQLite` upgrades the same transaction and
659    /// still preserves the statement's atomic boundary.
660    pub fn begin_deferred_transaction(&self) -> Result<()> {
661        self.begin_transaction_with("BEGIN DEFERRED")
662    }
663
664    fn begin_transaction_with(&self, statement: &str) -> Result<()> {
665        self.surface_cleanup_failure()?;
666        let _gate = self.session.gate.write();
667        let mut transaction = self.session.transaction.lock();
668        if transaction.is_some() {
669            return Err(SQLiteError::TransactionAlreadyActive);
670        }
671        let connection = self.pool.checkout()?;
672        connection.connection()?.execute_batch(statement)?;
673        self.session.transaction_failure.lock().take();
674        *transaction = Some(connection);
675        Ok(())
676    }
677
678    pub fn commit_transaction(&self) -> Result<()> {
679        self.surface_cleanup_failure()?;
680        self.finish_transaction("COMMIT")
681    }
682
683    pub fn rollback_transaction(&self) -> Result<()> {
684        self.surface_cleanup_failure()?;
685        self.finish_transaction("ROLLBACK")
686    }
687
688    /// Drop-only transaction cleanup. A rollback error is recorded on the
689    /// logical session and is returned by its next operation instead of being
690    /// mistaken for a successful rollback.
691    pub(crate) fn rollback_transaction_on_drop(&self) {
692        if let Err(error) = self.finish_transaction("ROLLBACK") {
693            let mut failure = self.session.cleanup_failure.lock();
694            if failure.is_none() {
695                *failure = Some(error.to_string());
696            }
697        }
698    }
699
700    fn finish_transaction(&self, statement: &str) -> Result<()> {
701        let _gate = self.session.gate.write();
702        let mut transaction = self.session.transaction.lock();
703        let connection = transaction
704            .as_ref()
705            .ok_or(SQLiteError::NoActiveTransaction)?;
706        if statement == "COMMIT" {
707            // Materialize the failure before entering the branch. Holding the
708            // mutex guard created by an `if let` scrutinee until the end of
709            // the branch would deadlock when cleanup takes the same lock.
710            let transaction_failure = self.session.transaction_failure.lock().clone();
711            if let Some(error) = transaction_failure {
712                if let Err(rollback_error) = connection.connection()?.execute_batch("ROLLBACK") {
713                    transaction.take();
714                    self.session.transaction_failure.lock().take();
715                    return Err(SQLiteError::SQLite(rollback_error));
716                }
717                transaction.take();
718                self.session.transaction_failure.lock().take();
719                return Err(SQLiteError::TransactionAborted(error));
720            }
721        }
722        if let Err(error) = connection.connection()?.execute_batch(statement) {
723            transaction.take();
724            self.session.transaction_failure.lock().take();
725            return Err(SQLiteError::SQLite(error));
726        }
727        transaction.take();
728        self.session.transaction_failure.lock().take();
729        Ok(())
730    }
731
732    fn with_transaction<R>(&self, f: impl FnOnce(&Connection) -> Result<R>) -> Result<R> {
733        let _gate = self.session.gate.read();
734        let transaction = self.session.transaction.lock();
735        let connection = transaction
736            .as_ref()
737            .ok_or(SQLiteError::NoActiveTransaction)?;
738        f(connection.connection()?)
739    }
740
741    pub fn savepoint(&self, name: &str) -> Result<()> {
742        self.surface_cleanup_failure()?;
743        let stmt = format!("SAVEPOINT \"{}\"", name.replace('"', "\"\""));
744        self.with_transaction(|c| {
745            c.execute_batch(&stmt)?;
746            Ok(())
747        })
748    }
749
750    pub fn release_savepoint(&self, name: &str) -> Result<()> {
751        self.surface_cleanup_failure()?;
752        let stmt = format!("RELEASE SAVEPOINT \"{}\"", name.replace('"', "\"\""));
753        self.with_transaction(|c| {
754            c.execute_batch(&stmt)?;
755            Ok(())
756        })
757    }
758
759    pub fn rollback_to_savepoint(&self, name: &str) -> Result<()> {
760        self.surface_cleanup_failure()?;
761        let stmt = format!("ROLLBACK TO SAVEPOINT \"{}\"", name.replace('"', "\"\""));
762        self.with_transaction(|c| {
763            c.execute_batch(&stmt)?;
764            Ok(())
765        })?;
766        self.session.transaction_failure.lock().take();
767        Ok(())
768    }
769}
770
771fn default_pool_connections() -> usize {
772    std::thread::available_parallelism()
773        .map_or(MIN_POOL_CONNECTIONS, |parallelism| parallelism.get() * 2)
774        .clamp(MIN_POOL_CONNECTIONS, MAX_POOL_CONNECTIONS)
775}
776
777#[cfg(test)]
778mod tests;
779
780impl From<SQLiteError> for uqa_storage::StorageBackendError {
781    fn from(source: SQLiteError) -> Self {
782        match source {
783            SQLiteError::Memory(error) => Self::Memory(error),
784            SQLiteError::Cancelled(error) => Self::Cancelled(error),
785            source => Self::backend("SQLite", source),
786        }
787    }
788}
789
790impl From<SQLiteError> for uqa_storage::TransactionError {
791    fn from(source: SQLiteError) -> Self {
792        Self::Storage(source.into())
793    }
794}
795
796impl From<uqa_storage::StorageBackendError> for SQLiteError {
797    fn from(error: uqa_storage::StorageBackendError) -> Self {
798        use uqa_storage::StorageBackendError;
799        match error {
800            StorageBackendError::Memory(error) => Self::Memory(error),
801            StorageBackendError::Cancelled(error) => Self::Cancelled(error),
802            StorageBackendError::Analysis(error) => Self::Analysis(error),
803            StorageBackendError::Serde(error) => Self::Serde(error),
804            StorageBackendError::Backend { backend, source } => match source.downcast::<Self>() {
805                Ok(error) => *error,
806                Err(source) => Self::StorageBackend(format!("{backend} storage failed: {source}")),
807            },
808            StorageBackendError::Other(message) => Self::StorageBackend(message),
809        }
810    }
811}
uqa_storage_sqlite/connection.rs

uqa_storage_sqlite/
connection.rs
