1use std::path::{Path, PathBuf};
17use std::sync::Arc;
18
19use parking_lot::{Condvar, Mutex, RwLock};
20use rusqlite::{Connection, OpenFlags};
21
22use crate::sqlite::compressed_vfs::{
23 self, SQLiteCompressedContainerAnchor, SQLiteCompressionOptions,
24};
25
26#[derive(Debug, thiserror::Error)]
27pub enum SQLiteError {
28 #[error("text analysis failed: {0}")]
29 Analysis(#[from] uqa_analysis::AnalysisError),
30 #[error("sqlite error: {0}")]
31 SQLite(#[from] rusqlite::Error),
32 #[error("encryption key must not be empty")]
33 EmptyEncryptionKey,
34 #[error("database requires an encryption key")]
35 EncryptionKeyRequired,
36 #[error("database is not encrypted but an encryption key was provided")]
37 NotEncrypted,
38 #[error("compressed sqlite container error: {0}")]
39 CompressedContainer(String),
40 #[error("io error: {0}")]
41 Io(#[from] std::io::Error),
42 #[error("catalog migration {version} failed: {source}")]
43 Migration {
44 version: u32,
45 #[source]
46 source: rusqlite::Error,
47 },
48 #[error("invalid persisted catalog schema version `{0}`")]
49 InvalidSchemaVersion(String),
50 #[error("catalog schema version {found} is newer than this engine supports ({supported})")]
51 UnsupportedSchemaVersion { found: u32, supported: u32 },
52 #[error("corrupt document blob for `{table}` doc {doc_id} field `{field}`: {reason}")]
53 CorruptDocumentBlob {
54 table: String,
55 doc_id: u64,
56 field: String,
57 reason: String,
58 },
59 #[error("payload serialization failed: {0}")]
60 Serde(#[from] serde_json::Error),
61 #[error("storage backend error: {0}")]
62 StorageBackend(String),
63 #[error("transaction already active for this sqlite session")]
64 TransactionAlreadyActive,
65 #[error("no active transaction for this sqlite session")]
66 NoActiveTransaction,
67 #[error("sqlite transaction was aborted by an earlier storage error: {0}")]
68 TransactionAborted(String),
69 #[error("sqlite session cleanup failed: {0}")]
70 SessionCleanupFailed(String),
71 #[error("sqlite connection-pool checkout lost its connection")]
72 MissingCheckedOutConnection,
73}
74
75pub type Result<T> = std::result::Result<T, SQLiteError>;
76
77const MIN_POOL_CONNECTIONS: usize = 4;
78const MAX_POOL_CONNECTIONS: usize = 32;
79
80#[derive(Clone)]
81enum ConnectionSpec {
82 File {
83 path: PathBuf,
84 key: Option<Arc<str>>,
85 },
86 Compressed {
87 path: PathBuf,
88 compression: SQLiteCompressionOptions,
89 },
90 Memory,
91}
92
93impl ConnectionSpec {
94 fn open(&self, initialize_database: bool) -> Result<Connection> {
95 match self {
96 Self::File { path, key } => {
97 let conn = Connection::open(path)?;
98 if let Some(key) = key {
99 ManagedConnection::apply_encryption_key(&conn, key)?;
100 }
101 if initialize_database {
102 ManagedConnection::enable_wal(&conn)?;
103 }
104 ManagedConnection::configure_wal_connection(&conn)?;
105 Ok(conn)
106 }
107 Self::Compressed { path, compression } => {
108 let conn = Connection::open_with_flags_and_vfs(
109 path,
110 OpenFlags::default(),
111 compressed_vfs::VFS_NAME,
112 )?;
113 if initialize_database {
114 conn.pragma_update(None, "page_size", compression.page_size)?;
115 ManagedConnection::enable_compressed_journal(&conn)?;
116 }
117 ManagedConnection::configure_compressed_connection(&conn)?;
118 Ok(conn)
119 }
120 Self::Memory => {
121 let conn = Connection::open_in_memory()?;
122 if initialize_database {
123 ManagedConnection::enable_wal(&conn)?;
124 }
125 ManagedConnection::configure_wal_connection(&conn)?;
126 Ok(conn)
127 }
128 }
129 }
130}
131
132struct PoolState {
133 idle: Vec<Connection>,
134 open: usize,
135}
136
137struct ConnectionPool {
138 spec: ConnectionSpec,
139 max_connections: usize,
140 state: Mutex<PoolState>,
141 available: Condvar,
142}
143
144impl ConnectionPool {
145 fn new(spec: ConnectionSpec, initial: Connection, max_connections: usize) -> Arc<Self> {
146 Arc::new(Self {
147 spec,
148 max_connections: max_connections.max(1),
149 state: Mutex::new(PoolState {
150 idle: vec![initial],
151 open: 1,
152 }),
153 available: Condvar::new(),
154 })
155 }
156
157 fn checkout(self: &Arc<Self>) -> Result<PooledConnection> {
158 loop {
159 let mut state = self.state.lock();
160 if let Some(connection) = state.idle.pop() {
161 return Ok(PooledConnection {
162 pool: Arc::clone(self),
163 connection: Some(connection),
164 });
165 }
166 if state.open < self.max_connections {
167 state.open += 1;
168 drop(state);
169 return match self.spec.open(false) {
170 Ok(connection) => Ok(PooledConnection {
171 pool: Arc::clone(self),
172 connection: Some(connection),
173 }),
174 Err(error) => {
175 let mut state = self.state.lock();
176 state.open -= 1;
177 self.available.notify_one();
178 Err(error)
179 }
180 };
181 }
182 self.available.wait(&mut state);
183 }
184 }
185
186 fn checkin(&self, connection: Connection) {
187 self.state.lock().idle.push(connection);
188 self.available.notify_one();
189 }
190
191 fn discard(&self) {
192 let mut state = self.state.lock();
193 state.open -= 1;
194 self.available.notify_one();
195 }
196}
197
198struct PooledConnection {
199 pool: Arc<ConnectionPool>,
200 connection: Option<Connection>,
201}
202
203impl PooledConnection {
204 fn connection(&self) -> Result<&Connection> {
205 self.connection
206 .as_ref()
207 .ok_or(SQLiteError::MissingCheckedOutConnection)
208 }
209
210 fn connection_mut(&mut self) -> Result<&mut Connection> {
211 self.connection
212 .as_mut()
213 .ok_or(SQLiteError::MissingCheckedOutConnection)
214 }
215}
216
217impl Drop for PooledConnection {
218 fn drop(&mut self) {
219 let Some(connection) = self.connection.take() else {
220 return;
221 };
222 let reusable = connection.is_autocommit() || connection.execute_batch("ROLLBACK").is_ok();
223 if reusable {
224 self.pool.checkin(connection);
225 } else {
226 self.pool.discard();
227 }
228 }
229}
230
231struct SessionState {
232 gate: RwLock<()>,
236 transaction: Mutex<Option<PooledConnection>>,
237 transaction_failure: Mutex<Option<String>>,
238 cleanup_failure: Mutex<Option<String>>,
239 data_version_monitor: Mutex<Option<Connection>>,
243}
244
245impl SessionState {
246 fn new() -> Self {
247 Self {
248 gate: RwLock::new(()),
249 transaction: Mutex::new(None),
250 transaction_failure: Mutex::new(None),
251 cleanup_failure: Mutex::new(None),
252 data_version_monitor: Mutex::new(None),
253 }
254 }
255}
256
257impl Drop for SessionState {
258 fn drop(&mut self) {
259 self.transaction.get_mut().take();
263 }
264}
265
266#[derive(Clone)]
270pub struct ManagedConnection {
271 pool: Arc<ConnectionPool>,
272 session: Arc<SessionState>,
273}
274
275impl ManagedConnection {
276 fn surface_cleanup_failure(&self) -> Result<()> {
277 if let Some(error) = self.session.cleanup_failure.lock().take() {
278 return Err(SQLiteError::SessionCleanupFailed(error));
279 }
280 Ok(())
281 }
282
283 pub fn open(path: &Path) -> Result<Self> {
284 if path == Path::new(":memory:") {
285 return Self::open_in_memory();
286 }
287 Self::open_with_optional_key(path, None)
288 }
289
290 #[must_use]
292 pub fn database_path(&self) -> Option<&Path> {
293 match &self.pool.spec {
294 ConnectionSpec::File { path, .. } | ConnectionSpec::Compressed { path, .. } => {
295 Some(path)
296 }
297 ConnectionSpec::Memory => None,
298 }
299 }
300
301 pub fn open_encrypted(path: &Path, key: &str) -> Result<Self> {
302 Self::open_with_optional_key(path, Some(key))
303 }
304
305 pub fn open_compressed(path: &Path, compression: SQLiteCompressionOptions) -> Result<Self> {
306 Self::open_compressed_with_optional_key(path, compression, None, None)
307 }
308
309 pub fn open_compressed_encrypted(
310 path: &Path,
311 key: &str,
312 compression: SQLiteCompressionOptions,
313 ) -> Result<Self> {
314 if key.is_empty() {
315 return Err(SQLiteError::EmptyEncryptionKey);
316 }
317 Self::open_compressed_with_optional_key(path, compression, Some(key), None)
318 }
319
320 pub fn open_compressed_encrypted_with_anchor(
323 path: &Path,
324 key: &str,
325 compression: SQLiteCompressionOptions,
326 trusted_anchor: SQLiteCompressedContainerAnchor,
327 ) -> Result<Self> {
328 if key.is_empty() {
329 return Err(SQLiteError::EmptyEncryptionKey);
330 }
331 Self::open_compressed_with_optional_key(path, compression, Some(key), Some(trusted_anchor))
332 }
333
334 fn open_with_optional_key(path: &Path, key: Option<&str>) -> Result<Self> {
335 let spec = ConnectionSpec::File {
336 path: path.to_path_buf(),
337 key: key.map(Arc::from),
338 };
339 Self::from_spec(spec, default_pool_connections())
340 }
341
342 fn open_compressed_with_optional_key(
343 path: &Path,
344 compression: SQLiteCompressionOptions,
345 key: Option<&str>,
346 trusted_anchor: Option<SQLiteCompressedContainerAnchor>,
347 ) -> Result<Self> {
348 let compression = compression
349 .validate()
350 .map_err(SQLiteError::CompressedContainer)?;
351 match trusted_anchor {
352 Some(anchor) => compressed_vfs::register_database_with_anchor(
353 path,
354 compression,
355 key.ok_or(SQLiteError::EncryptionKeyRequired)?,
356 anchor,
357 ),
358 None => compressed_vfs::register_database(path, compression, key),
359 }
360 .map_err(SQLiteError::CompressedContainer)?;
361 let spec = ConnectionSpec::Compressed {
362 path: path.to_path_buf(),
363 compression,
364 };
365 Self::from_spec(spec, default_pool_connections())
366 }
367
368 pub fn open_in_memory() -> Result<Self> {
369 Self::from_spec(ConnectionSpec::Memory, 1)
373 }
374
375 fn from_spec(spec: ConnectionSpec, max_connections: usize) -> Result<Self> {
376 let initial = spec.open(true)?;
377 Ok(Self {
378 pool: ConnectionPool::new(spec, initial, max_connections),
379 session: Arc::new(SessionState::new()),
380 })
381 }
382
383 fn apply_encryption_key(conn: &Connection, key: &str) -> Result<()> {
384 if key.is_empty() {
385 return Err(SQLiteError::EmptyEncryptionKey);
386 }
387 conn.pragma_update(None, "key", key)?;
388 Ok(())
389 }
390
391 fn enable_wal(conn: &Connection) -> Result<()> {
392 conn.pragma_update(None, "journal_mode", "WAL")?;
393 Ok(())
394 }
395
396 fn configure_wal_connection(conn: &Connection) -> Result<()> {
397 conn.busy_timeout(std::time::Duration::from_secs(5))?;
400 conn.pragma_update(None, "synchronous", "NORMAL")?;
403 conn.pragma_update(None, "foreign_keys", "ON")?;
406 Ok(())
407 }
408
409 fn enable_compressed_journal(conn: &Connection) -> Result<()> {
410 conn.pragma_update(None, "journal_mode", "DELETE")?;
416 Ok(())
417 }
418
419 fn configure_compressed_connection(conn: &Connection) -> Result<()> {
420 conn.busy_timeout(std::time::Duration::from_secs(5))?;
421 conn.pragma_update(None, "synchronous", "FULL")?;
422 conn.pragma_update(None, "temp_store", "MEMORY")?;
423 conn.pragma_update(None, "foreign_keys", "ON")?;
424 Ok(())
425 }
426
427 #[must_use]
431 pub fn new_session(&self) -> Self {
432 Self {
433 pool: Arc::clone(&self.pool),
434 session: Arc::new(SessionState::new()),
435 }
436 }
437
438 #[must_use]
441 pub fn in_transaction(&self) -> bool {
442 let _gate = self.session.gate.read();
443 self.session.transaction.lock().is_some()
444 }
445
446 pub fn transaction_has_written(&self) -> Result<bool> {
450 let _gate = self.session.gate.read();
451 let transaction = self.session.transaction.lock();
452 let transaction = transaction
453 .as_ref()
454 .ok_or(SQLiteError::NoActiveTransaction)?;
455 Ok(matches!(
456 transaction.connection()?.transaction_state(Some("main"))?,
457 rusqlite::TransactionState::Write
458 ))
459 }
460
461 pub fn data_version_monitor_is_nonblocking(&self) -> Result<bool> {
472 if !matches!(&self.pool.spec, ConnectionSpec::Compressed { .. }) {
473 return Ok(true);
474 }
475 let _gate = self.session.gate.read();
476 let transaction = self.session.transaction.lock();
477 let Some(transaction) = transaction.as_ref() else {
478 return Ok(true);
479 };
480 Ok(!matches!(
481 transaction.connection()?.transaction_state(Some("main"))?,
482 rusqlite::TransactionState::Write
483 ))
484 }
485
486 #[must_use]
488 pub fn supports_concurrent_pinned_read_and_write(&self) -> bool {
489 !matches!(&self.pool.spec, ConnectionSpec::Compressed { .. })
490 }
491
492 pub fn data_version(&self) -> Result<Option<u64>> {
497 if matches!(&self.pool.spec, ConnectionSpec::Memory) {
498 return Ok(None);
499 }
500 let mut monitor = self.session.data_version_monitor.lock();
501 if monitor.is_none() {
502 *monitor = Some(self.pool.spec.open(false)?);
503 }
504 let monitor = monitor.as_ref().ok_or_else(|| {
505 SQLiteError::StorageBackend(
506 "data-version monitor was not initialized after opening it".into(),
507 )
508 })?;
509 let version: i64 = monitor.pragma_query_value(None, "data_version", |row| row.get(0))?;
510 let version = u64::try_from(version).map_err(|_| {
511 SQLiteError::StorageBackend(format!(
512 "SQLite returned a negative PRAGMA data_version: {version}"
513 ))
514 })?;
515 Ok(Some(version))
516 }
517
518 pub fn pin_transaction_snapshot(&self) -> Result<()> {
525 self.with(|connection| {
526 if connection.is_autocommit() {
527 return Err(SQLiteError::NoActiveTransaction);
528 }
529 let _: i64 =
530 connection.query_row("SELECT COUNT(*) FROM sqlite_schema", [], |row| row.get(0))?;
531 Ok(())
532 })
533 }
534
535 pub fn with<R>(&self, f: impl FnOnce(&Connection) -> Result<R>) -> Result<R> {
539 self.surface_cleanup_failure()?;
540 let _gate = self.session.gate.read();
541 let transaction = self.session.transaction.lock();
542 if let Some(connection) = transaction.as_ref() {
543 if let Some(error) = self.session.transaction_failure.lock().as_ref() {
544 return Err(SQLiteError::TransactionAborted(error.clone()));
545 }
546 let result = f(connection.connection()?);
547 if let Err(error) = &result {
548 let mut failure = self.session.transaction_failure.lock();
549 if failure.is_none() {
550 *failure = Some(error.to_string());
551 }
552 }
553 return result;
554 }
555 drop(transaction);
556 let connection = self.pool.checkout()?;
557 f(connection.connection()?)
558 }
559
560 pub fn with_mut<R>(&self, f: impl FnOnce(&mut Connection) -> Result<R>) -> Result<R> {
561 self.surface_cleanup_failure()?;
562 let _gate = self.session.gate.read();
563 let mut transaction = self.session.transaction.lock();
564 if let Some(connection) = transaction.as_mut() {
565 if let Some(error) = self.session.transaction_failure.lock().as_ref() {
566 return Err(SQLiteError::TransactionAborted(error.clone()));
567 }
568 let result = f(connection.connection_mut()?);
569 if let Err(error) = &result {
570 let mut failure = self.session.transaction_failure.lock();
571 if failure.is_none() {
572 *failure = Some(error.to_string());
573 }
574 }
575 return result;
576 }
577 drop(transaction);
578 let mut connection = self.pool.checkout()?;
579 f(connection.connection_mut()?)
580 }
581
582 pub fn vacuum(&self) -> Result<()> {
584 self.surface_cleanup_failure()?;
585 let _gate = self.session.gate.write();
586 if self.session.transaction.lock().is_some() {
587 return Err(SQLiteError::TransactionAlreadyActive);
588 }
589 let connection = self.pool.checkout()?;
590 connection.connection()?.execute_batch("VACUUM")?;
591 Ok(())
592 }
593
594 pub fn begin_transaction(&self) -> Result<()> {
601 self.begin_transaction_with("BEGIN IMMEDIATE")
602 }
603
604 pub fn begin_deferred_transaction(&self) -> Result<()> {
609 self.begin_transaction_with("BEGIN DEFERRED")
610 }
611
612 fn begin_transaction_with(&self, statement: &str) -> Result<()> {
613 self.surface_cleanup_failure()?;
614 let _gate = self.session.gate.write();
615 let mut transaction = self.session.transaction.lock();
616 if transaction.is_some() {
617 return Err(SQLiteError::TransactionAlreadyActive);
618 }
619 let connection = self.pool.checkout()?;
620 connection.connection()?.execute_batch(statement)?;
621 self.session.transaction_failure.lock().take();
622 *transaction = Some(connection);
623 Ok(())
624 }
625
626 pub fn commit_transaction(&self) -> Result<()> {
627 self.surface_cleanup_failure()?;
628 self.finish_transaction("COMMIT")
629 }
630
631 pub fn rollback_transaction(&self) -> Result<()> {
632 self.surface_cleanup_failure()?;
633 self.finish_transaction("ROLLBACK")
634 }
635
636 pub(crate) fn rollback_transaction_on_drop(&self) {
640 if let Err(error) = self.finish_transaction("ROLLBACK") {
641 let mut failure = self.session.cleanup_failure.lock();
642 if failure.is_none() {
643 *failure = Some(error.to_string());
644 }
645 }
646 }
647
648 fn finish_transaction(&self, statement: &str) -> Result<()> {
649 let _gate = self.session.gate.write();
650 let mut transaction = self.session.transaction.lock();
651 let connection = transaction
652 .as_ref()
653 .ok_or(SQLiteError::NoActiveTransaction)?;
654 if statement == "COMMIT" {
655 let transaction_failure = self.session.transaction_failure.lock().clone();
659 if let Some(error) = transaction_failure {
660 if let Err(rollback_error) = connection.connection()?.execute_batch("ROLLBACK") {
661 transaction.take();
662 self.session.transaction_failure.lock().take();
663 return Err(SQLiteError::SQLite(rollback_error));
664 }
665 transaction.take();
666 self.session.transaction_failure.lock().take();
667 return Err(SQLiteError::TransactionAborted(error));
668 }
669 }
670 if let Err(error) = connection.connection()?.execute_batch(statement) {
671 transaction.take();
672 self.session.transaction_failure.lock().take();
673 return Err(SQLiteError::SQLite(error));
674 }
675 transaction.take();
676 self.session.transaction_failure.lock().take();
677 Ok(())
678 }
679
680 fn with_transaction<R>(&self, f: impl FnOnce(&Connection) -> Result<R>) -> Result<R> {
681 let _gate = self.session.gate.read();
682 let transaction = self.session.transaction.lock();
683 let connection = transaction
684 .as_ref()
685 .ok_or(SQLiteError::NoActiveTransaction)?;
686 f(connection.connection()?)
687 }
688
689 pub fn savepoint(&self, name: &str) -> Result<()> {
690 self.surface_cleanup_failure()?;
691 let stmt = format!("SAVEPOINT \"{}\"", name.replace('"', "\"\""));
692 self.with_transaction(|c| {
693 c.execute_batch(&stmt)?;
694 Ok(())
695 })
696 }
697
698 pub fn release_savepoint(&self, name: &str) -> Result<()> {
699 self.surface_cleanup_failure()?;
700 let stmt = format!("RELEASE SAVEPOINT \"{}\"", name.replace('"', "\"\""));
701 self.with_transaction(|c| {
702 c.execute_batch(&stmt)?;
703 Ok(())
704 })
705 }
706
707 pub fn rollback_to_savepoint(&self, name: &str) -> Result<()> {
708 self.surface_cleanup_failure()?;
709 let stmt = format!("ROLLBACK TO SAVEPOINT \"{}\"", name.replace('"', "\"\""));
710 self.with_transaction(|c| {
711 c.execute_batch(&stmt)?;
712 Ok(())
713 })?;
714 self.session.transaction_failure.lock().take();
715 Ok(())
716 }
717}
718
719fn default_pool_connections() -> usize {
720 std::thread::available_parallelism()
721 .map_or(MIN_POOL_CONNECTIONS, |parallelism| parallelism.get() * 2)
722 .clamp(MIN_POOL_CONNECTIONS, MAX_POOL_CONNECTIONS)
723}
724
725#[cfg(test)]
726mod tests;