reifydb_store_multi/tier/persistent/sqlite/
storage.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::{
5	collections::HashMap,
6	iter::repeat_n,
7	ops::Bound,
8	sync::{
9		Arc,
10		atomic::{AtomicU32, AtomicUsize, Ordering},
11	},
12};
13
14use reifydb_codec::key::encoded::EncodedKey;
15use reifydb_core::{common::CommitVersion, error::diagnostic::internal::internal, interface::store::EntryKind};
16use reifydb_runtime::{
17	shutdown::Shutdown,
18	sync::{
19		map::Map,
20		mutex::{Mutex, MutexGuard},
21	},
22};
23use reifydb_sqlite::{
24	SqliteConfig, SqliteTempPathGuard,
25	connection::{connect, convert_flags, resolve_db_path},
26	pragma,
27};
28use reifydb_value::{Result, error, util::cowvec::CowVec, value::duration::Duration};
29use rusqlite::{
30	Connection, Error::QueryReturnedNoRows, Result as SqliteResult, Row, ToSql, Transaction, TransactionBehavior,
31	params, params_from_iter,
32};
33use tracing::{instrument, warn};
34
35use crate::{
36	MultiVersionScope,
37	tier::{
38		HistoricalCursor, RangeBatch, RangeCursor, RawEntry, TierBackend, TierBatch, TierStorage,
39		VersionedGetResult,
40		persistent::{
41			CheckpointOutcome,
42			sqlite::{
43				entry::current_table_name,
44				query::{
45					build_create_current_sql, build_delete_below_version_sql,
46					build_delete_key_through_sql, build_delete_keys_sql, build_get_current_sql,
47					build_get_many_current_sql, build_range_consistent_sql,
48					build_range_current_sql, build_upsert_current_sql, prefix_upper_bound,
49					version_from_bytes, version_to_bytes,
50				},
51			},
52		},
53	},
54};
55
56const GET_MANY_CHUNK: usize = 900;
57
58const GET_MANY_BUCKETS: [usize; 5] = [1, 8, 64, 512, GET_MANY_CHUNK];
59
60fn bucket_key_count(len: usize) -> usize {
61	for &bucket in GET_MANY_BUCKETS.iter() {
62		if len <= bucket {
63			return bucket;
64		}
65	}
66	GET_MANY_CHUNK
67}
68
69const BUSY_TIMEOUT: Duration = Duration::from_milliseconds_const(200);
70
71#[derive(Clone)]
72pub struct SqlitePersistentStorage {
73	inner: Arc<SqlitePersistentStorageInner>,
74}
75
76struct SqlitePersistentStorageInner {
77	conn: Mutex<Option<Connection>>,
78	readers: ReadPool,
79	checkpoint_threshold_frames: AtomicU32,
80	table_sql: Map<EntryKind, Arc<TableSql>>,
81}
82
83struct TableSql {
84	table_name: String,
85	get_sql: String,
86	upsert_sql: String,
87	create_sql: String,
88}
89
90impl TableSql {
91	fn build(table: EntryKind) -> Self {
92		let table_name = current_table_name(table);
93		let get_sql = build_get_current_sql(&table_name);
94		let upsert_sql = build_upsert_current_sql(&table_name);
95		let create_sql = build_create_current_sql(&table_name);
96		Self {
97			table_name,
98			get_sql,
99			upsert_sql,
100			create_sql,
101		}
102	}
103}
104
105struct ReadPool {
106	conns: Vec<Mutex<Option<Connection>>>,
107	next: AtomicUsize,
108}
109
110impl ReadPool {
111	fn acquire(&self) -> MutexGuard<'_, Option<Connection>> {
112		let n = self.conns.len();
113		let start = self.next.fetch_add(1, Ordering::Relaxed) % n;
114		for i in 0..n {
115			if let Some(guard) = self.conns[(start + i) % n].try_lock() {
116				return guard;
117			}
118		}
119		self.conns[start].lock()
120	}
121
122	fn shutdown(&self) {
123		for slot in &self.conns {
124			drop(slot.lock().take());
125		}
126	}
127}
128
129impl SqlitePersistentStorage {
130	#[instrument(name = "store::multi::persistent::sqlite::new", level = "debug", skip(config), fields(
131		db_path = ?config.path,
132		page_size = config.page_size.as_bytes(),
133		journal_mode = %config.journal_mode.as_str()
134	))]
135	pub fn new(config: SqliteConfig) -> Self {
136		let db_path = resolve_db_path(config.path.clone(), "persistent.db");
137		let flags = convert_flags(&config.flags);
138
139		let conn = connect(&db_path, flags).expect("Failed to connect to persistent database");
140		pragma::apply(&conn, &config).expect("Failed to configure persistent SQLite pragmas");
141		conn.busy_timeout(BUSY_TIMEOUT.to_std()).expect("Failed to set persistent busy timeout");
142
143		let pool_size = config.read_pool_size.max(1) as usize;
144		let mut conns = Vec::with_capacity(pool_size);
145		for _ in 0..pool_size {
146			let reader = connect(&db_path, flags).expect("Failed to open persistent read connection");
147			pragma::apply_read_only(&reader, &config)
148				.expect("Failed to configure persistent read connection");
149			reader.busy_timeout(BUSY_TIMEOUT.to_std()).expect("Failed to set persistent read busy timeout");
150			conns.push(Mutex::new(Some(reader)));
151		}
152
153		Self {
154			inner: Arc::new(SqlitePersistentStorageInner {
155				conn: Mutex::new(Some(conn)),
156				readers: ReadPool {
157					conns,
158					next: AtomicUsize::new(0),
159				},
160				checkpoint_threshold_frames: AtomicU32::new(config.wal_autocheckpoint),
161				table_sql: Map::new(),
162			}),
163		}
164	}
165
166	#[instrument(name = "store::multi::sqlite::conn_acquire", level = "debug", skip(self))]
167	fn lock_conn(&self) -> MutexGuard<'_, Option<Connection>> {
168		self.inner.conn.lock()
169	}
170
171	#[instrument(name = "store::multi::sqlite::checkpoint", level = "debug", skip(self))]
172	pub fn maybe_checkpoint(&self) -> Result<CheckpointOutcome> {
173		let guard = self.lock_conn();
174		let Some(conn) = guard.as_ref() else {
175			return Ok(CheckpointOutcome {
176				log_frames: 0,
177				restarted: false,
178			});
179		};
180
181		let mut log_frames: i64 = 0;
182		conn.pragma(None, "wal_checkpoint", "PASSIVE", |row| {
183			log_frames = row.get(1)?;
184			Ok(())
185		})
186		.map_err(|e| error!(internal(format!("Failed to query persistent WAL size: {}", e))))?;
187
188		let log_frames = log_frames.max(0) as u32;
189		if log_frames <= self.inner.checkpoint_threshold_frames.load(Ordering::Relaxed) {
190			return Ok(CheckpointOutcome {
191				log_frames,
192				restarted: false,
193			});
194		}
195
196		let mut busy: i64 = 1;
197		if let Err(e) = conn.pragma(None, "wal_checkpoint", "RESTART", |row| {
198			busy = row.get(0)?;
199			Ok(())
200		}) {
201			warn!(error = %e, "persistent checkpoint: RESTART failed");
202		}
203
204		Ok(CheckpointOutcome {
205			log_frames,
206			restarted: busy == 0,
207		})
208	}
209
210	pub fn set_checkpoint_threshold(&self, frames: u32) {
211		self.inner.checkpoint_threshold_frames.store(frames, Ordering::Relaxed);
212		let guard = self.lock_conn();
213		if let Some(conn) = guard.as_ref()
214			&& let Err(e) = conn.pragma_update(None, "wal_autocheckpoint", frames)
215		{
216			warn!(error = %e, "failed to update wal_autocheckpoint pragma");
217		}
218	}
219
220	pub fn in_memory() -> (Self, SqliteTempPathGuard) {
221		let (config, guard) = SqliteConfig::in_memory();
222		(Self::new(config), guard)
223	}
224
225	fn table_sql(&self, table: EntryKind) -> Arc<TableSql> {
226		self.inner.table_sql.get_or_insert_with(table, || Arc::new(TableSql::build(table)))
227	}
228
229	pub fn count_current(&self, table: EntryKind) -> Result<u64> {
230		let table_sql = self.table_sql(table);
231		let guard = self.inner.readers.acquire();
232		let Some(conn) = guard.as_ref() else {
233			return Ok(0);
234		};
235		let sql = format!("SELECT COUNT(*) FROM \"{}\"", table_sql.table_name);
236		match conn.query_row(&sql, [], |row| row.get::<_, i64>(0)) {
237			Ok(c) => Ok(c as u64),
238			Err(e) if e.to_string().contains("no such table") => Ok(0),
239			Err(e) => Err(error!(internal(format!("Failed to count persistent current: {}", e)))),
240		}
241	}
242
243	pub fn delete_below_version(
244		&self,
245		table: EntryKind,
246		cutoff_version: CommitVersion,
247		prefix: Option<&[u8]>,
248	) -> Result<Vec<EncodedKey>> {
249		let table_sql = self.table_sql(table);
250		let sql = build_delete_below_version_sql(&table_sql.table_name, prefix.is_some());
251		let cutoff = version_to_bytes(cutoff_version);
252		let guard = self.lock_conn();
253		let Some(conn) = guard.as_ref() else {
254			return Ok(Vec::new());
255		};
256		let mut stmt = match conn.prepare_cached(&sql) {
257			Ok(stmt) => stmt,
258			Err(e) if e.to_string().contains("no such table") => return Ok(Vec::new()),
259			Err(e) => {
260				return Err(error!(internal(format!(
261					"Failed to prepare delete expired for {}: {}",
262					table_sql.table_name, e
263				))));
264			}
265		};
266		let map_key = |row: &Row| row.get::<_, Vec<u8>>(0);
267		let rows = match prefix {
268			Some(prefix) => {
269				let upper = prefix_upper_bound(prefix);
270				stmt.query_map(params![cutoff.as_slice(), prefix, upper.as_slice()], map_key)
271			}
272			None => stmt.query_map(params![cutoff.as_slice()], map_key),
273		};
274		let rows = match rows {
275			Ok(rows) => rows,
276			Err(e) if e.to_string().contains("no such table") => return Ok(Vec::new()),
277			Err(e) => {
278				return Err(error!(internal(format!(
279					"Failed to delete expired persistent rows from {}: {}",
280					table_sql.table_name, e
281				))));
282			}
283		};
284		let mut deleted = Vec::new();
285		for row in rows {
286			match row {
287				Ok(key) => deleted.push(EncodedKey::new(key)),
288				Err(e) => {
289					return Err(error!(internal(format!(
290						"Failed to read deleted key from {}: {}",
291						table_sql.table_name, e
292					))));
293				}
294			}
295		}
296		Ok(deleted)
297	}
298
299	pub fn delete_keys(&self, table: EntryKind, keys: &[EncodedKey]) -> Result<u64> {
300		if keys.is_empty() {
301			return Ok(0);
302		}
303		let table_sql = self.table_sql(table);
304		let guard = self.lock_conn();
305		let Some(conn) = guard.as_ref() else {
306			return Ok(0);
307		};
308		let mut total = 0u64;
309		for chunk in keys.chunks(GET_MANY_CHUNK) {
310			let sql = build_delete_keys_sql(&table_sql.table_name, chunk.len());
311			match conn.execute(&sql, params_from_iter(chunk.iter().map(|k| k.as_slice()))) {
312				Ok(n) => total += n as u64,
313				Err(e) if e.to_string().contains("no such table") => return Ok(total),
314				Err(e) => {
315					return Err(error!(internal(format!(
316						"Failed to delete keys from {}: {}",
317						table_sql.table_name, e
318					))));
319				}
320			}
321		}
322		Ok(total)
323	}
324
325	#[instrument(name = "store::multi::persistent::sqlite::delete_through", level = "debug", skip(self, keys), fields(key_count = keys.len()))]
326	pub fn delete_keys_through(&self, table: EntryKind, keys: &[(EncodedKey, CommitVersion)]) -> Result<u64> {
327		if keys.is_empty() {
328			return Ok(0);
329		}
330		let table_sql = self.table_sql(table);
331		let guard = self.lock_conn();
332		let Some(conn) = guard.as_ref() else {
333			return Ok(0);
334		};
335		let tx = Transaction::new_unchecked(conn, TransactionBehavior::Immediate)
336			.map_err(|e| error!(internal(format!("Failed to start persistent transaction: {}", e))))?;
337		let mut total = 0u64;
338		{
339			let sql = build_delete_key_through_sql(&table_sql.table_name);
340			let mut stmt = match tx.prepare_cached(&sql) {
341				Ok(s) => s,
342				Err(e) if e.to_string().contains("no such table") => return Ok(0),
343				Err(e) => {
344					return Err(error!(internal(format!(
345						"Failed to prepare persistent delete through: {}",
346						e
347					))));
348				}
349			};
350			for (key, version) in keys {
351				let version_bytes = version_to_bytes(*version);
352				let affected = stmt
353					.execute(params![key.as_slice(), version_bytes.as_slice()])
354					.map_err(|e| {
355						error!(internal(format!(
356							"Failed to delete key through version from {}: {}",
357							table_sql.table_name, e
358						)))
359					})?;
360				total += affected as u64;
361			}
362		}
363		tx.commit().map_err(|e| error!(internal(format!("Failed to commit persistent transaction: {}", e))))?;
364		Ok(total)
365	}
366
367	#[instrument(name = "store::multi::persistent::sqlite::set", level = "debug", skip(self, batches), fields(table_count = batches.len(), version = version.0))]
368	pub fn set_collecting_accepted(&self, version: CommitVersion, batches: TierBatch) -> Result<Vec<EncodedKey>> {
369		let mut accepted = Vec::new();
370		if batches.is_empty() {
371			return Ok(accepted);
372		}
373
374		let guard = self.lock_conn();
375		let Some(conn) = guard.as_ref() else {
376			return Ok(accepted);
377		};
378		let tx = Transaction::new_unchecked(conn, TransactionBehavior::Immediate)
379			.map_err(|e| error!(internal(format!("Failed to start persistent transaction: {}", e))))?;
380
381		let new_version_bytes = version_to_bytes(version);
382
383		for (table, entries) in batches {
384			let table_sql = self.table_sql(table);
385			Self::create_table_if_needed(&tx, &table_sql.create_sql)
386				.map_err(|e| error!(internal(format!("Failed to ensure persistent table: {}", e))))?;
387
388			let mut stmt = tx
389				.prepare_cached(&table_sql.upsert_sql)
390				.map_err(|e| error!(internal(format!("Failed to prepare persistent upsert: {}", e))))?;
391
392			for (key, value) in entries {
393				let value_slice = value.as_ref().map(|v| v.as_slice());
394				let affected = stmt
395					.execute(params![key.as_slice(), new_version_bytes.as_slice(), value_slice])
396					.map_err(|e| {
397						error!(internal(format!("Failed to upsert persistent row: {}", e)))
398					})?;
399				if affected > 0 {
400					accepted.push(key);
401				}
402			}
403		}
404
405		tx.commit().map_err(|e| error!(internal(format!("Failed to commit persistent transaction: {}", e))))?;
406		Ok(accepted)
407	}
408
409	#[instrument(name = "store::multi::persistent::sqlite::persist_sweep", level = "debug", skip(self, batches), fields(batch_count = batches.len()))]
410	pub fn persist_sweep(&self, batches: Vec<(CommitVersion, TierBatch)>) -> Result<Vec<EncodedKey>> {
411		let mut accepted = Vec::new();
412		if batches.iter().all(|(_, batch)| batch.is_empty()) {
413			return Ok(accepted);
414		}
415
416		let guard = self.lock_conn();
417		let Some(conn) = guard.as_ref() else {
418			return Err(error!(internal(
419				"Persistent storage is shut down; refusing to acknowledge a flush sweep whose \
420				 writes would then be dropped from the commit buffer unpersisted"
421					.to_string()
422			)));
423		};
424		let tx = Transaction::new_unchecked(conn, TransactionBehavior::Immediate)
425			.map_err(|e| error!(internal(format!("Failed to start persistent transaction: {}", e))))?;
426
427		for (version, batch) in batches {
428			let new_version_bytes = version_to_bytes(version);
429			for (table, entries) in batch {
430				let table_sql = self.table_sql(table);
431				Self::create_table_if_needed(&tx, &table_sql.create_sql).map_err(|e| {
432					error!(internal(format!("Failed to ensure persistent table: {}", e)))
433				})?;
434
435				let mut stmt = tx.prepare_cached(&table_sql.upsert_sql).map_err(|e| {
436					error!(internal(format!("Failed to prepare persistent upsert: {}", e)))
437				})?;
438
439				for (key, value) in entries {
440					let value_slice = value.as_ref().map(|v| v.as_slice());
441					let affected = stmt
442						.execute(params![
443							key.as_slice(),
444							new_version_bytes.as_slice(),
445							value_slice
446						])
447						.map_err(|e| {
448							error!(internal(format!(
449								"Failed to upsert persistent row: {}",
450								e
451							)))
452						})?;
453					if affected > 0 {
454						accepted.push(key);
455					}
456				}
457			}
458		}
459
460		tx.commit().map_err(|e| error!(internal(format!("Failed to commit persistent transaction: {}", e))))?;
461		Ok(accepted)
462	}
463
464	fn create_table_if_needed(conn: &Connection, create_sql: &str) -> SqliteResult<()> {
465		conn.execute_batch(create_sql)?;
466		Ok(())
467	}
468
469	fn range_chunk(&self, cursor: &mut RangeCursor, req: RangeChunkRequest<'_>) -> Result<RangeBatch> {
470		if cursor.exhausted {
471			return Ok(RangeBatch::empty());
472		}
473
474		let table_sql = self.table_sql(req.table);
475		let guard = self.inner.readers.acquire();
476		let Some(conn) = guard.as_ref() else {
477			cursor.exhausted = true;
478			return Ok(RangeBatch::empty());
479		};
480
481		let sql = build_range_current_sql(
482			&table_sql.table_name,
483			bound_shape(req.start),
484			bound_shape(req.end),
485			cursor.last_key.is_some(),
486			req.descending,
487		);
488
489		let mut stmt = match conn.prepare_cached(&sql) {
490			Ok(s) => s,
491			Err(e) if e.to_string().contains("no such table") => {
492				cursor.exhausted = true;
493				return Ok(RangeBatch::empty());
494			}
495			Err(e) => return Err(error!(internal(format!("Failed to prepare persistent range: {}", e)))),
496		};
497
498		let version_bytes = version_to_bytes(req.scope.read()).to_vec();
499		let limit_i64 = req.batch_size as i64;
500		let mut params: Vec<Box<dyn ToSql>> = Vec::new();
501		match req.start {
502			Bound::Included(s) | Bound::Excluded(s) => params.push(Box::new(s.to_vec())),
503			Bound::Unbounded => {}
504		}
505		match req.end {
506			Bound::Included(e) | Bound::Excluded(e) => params.push(Box::new(e.to_vec())),
507			Bound::Unbounded => {}
508		}
509		if let Some(k) = cursor.last_key.as_deref() {
510			params.push(Box::new(k.to_vec()));
511		}
512		params.push(Box::new(version_bytes));
513		params.push(Box::new(limit_i64));
514
515		let raw: Vec<RawEntry> = match stmt.query_map(params_from_iter(params), |row| {
516			let key: Vec<u8> = row.get(0)?;
517			let version_blob: Vec<u8> = row.get(1)?;
518			let value: Option<Vec<u8>> = row.get(2)?;
519			Ok(RawEntry {
520				key: EncodedKey::new(key),
521				version: version_from_bytes(&version_blob),
522				value: value.map(CowVec::new),
523			})
524		}) {
525			Ok(rows) => rows
526				.collect::<SqliteResult<Vec<_>>>()
527				.map_err(|e| error!(internal(format!("Failed to read persistent row: {}", e))))?,
528			Err(e) if e.to_string().contains("no such table") => {
529				cursor.exhausted = true;
530				return Ok(RangeBatch::empty());
531			}
532			Err(e) => return Err(error!(internal(format!("Failed to scan persistent range: {}", e)))),
533		};
534		let entries: Vec<RawEntry> = raw.into_iter().filter(|e| req.scope.contains(e.version)).collect();
535
536		if entries.len() < req.batch_size {
537			cursor.exhausted = true;
538		}
539		if let Some(last) = entries.last() {
540			cursor.last_key = Some(last.key.clone());
541		}
542
543		let has_more = !cursor.exhausted;
544		Ok(RangeBatch {
545			entries,
546			has_more,
547		})
548	}
549
550	#[instrument(name = "store::multi::persistent::sqlite::load_consistent", level = "debug", skip_all, fields(table = ?table))]
551	pub fn load_range_consistent(
552		&self,
553		table: EntryKind,
554		start: Bound<&[u8]>,
555		end: Bound<&[u8]>,
556		read: CommitVersion,
557		limit: Option<usize>,
558	) -> Result<Vec<RawEntry>> {
559		let table_sql = self.table_sql(table);
560		let guard = self.inner.readers.acquire();
561		let Some(conn) = guard.as_ref() else {
562			return Ok(Vec::new());
563		};
564
565		let sql = build_range_consistent_sql(&table_sql.table_name, bound_shape(start), bound_shape(end));
566
567		let mut stmt = match conn.prepare_cached(&sql) {
568			Ok(s) => s,
569			Err(e) if e.to_string().contains("no such table") => return Ok(Vec::new()),
570			Err(e) => {
571				return Err(error!(internal(format!(
572					"Failed to prepare persistent consistent range: {}",
573					e
574				))));
575			}
576		};
577
578		let version_bytes = version_to_bytes(read).to_vec();
579		let mut params: Vec<Box<dyn ToSql>> = Vec::new();
580		match start {
581			Bound::Included(s) | Bound::Excluded(s) => params.push(Box::new(s.to_vec())),
582			Bound::Unbounded => {}
583		}
584		match end {
585			Bound::Included(e) | Bound::Excluded(e) => params.push(Box::new(e.to_vec())),
586			Bound::Unbounded => {}
587		}
588		params.push(Box::new(version_bytes));
589
590		let raw: Vec<RawEntry> = match stmt.query_map(params_from_iter(params), |row| {
591			let key: Vec<u8> = row.get(0)?;
592			let version_blob: Vec<u8> = row.get(1)?;
593			let value: Option<Vec<u8>> = row.get(2)?;
594			Ok(RawEntry {
595				key: EncodedKey::new(key),
596				version: version_from_bytes(&version_blob),
597				value: value.map(CowVec::new),
598			})
599		}) {
600			Ok(rows) => {
601				let mut collected = Vec::new();
602				for row in rows {
603					let entry = row.map_err(|e| {
604						error!(internal(format!(
605							"Failed to read persistent consistent row: {}",
606							e
607						)))
608					})?;
609					collected.push(entry);
610					if limit.is_some_and(|l| collected.len() >= l) {
611						break;
612					}
613				}
614				collected
615			}
616			Err(e) if e.to_string().contains("no such table") => return Ok(Vec::new()),
617			Err(e) => {
618				return Err(error!(internal(format!(
619					"Failed to scan persistent consistent range: {}",
620					e
621				))));
622			}
623		};
624
625		Ok(raw)
626	}
627}
628
629fn bound_shape(b: Bound<&[u8]>) -> Bound<()> {
630	match b {
631		Bound::Included(_) => Bound::Included(()),
632		Bound::Excluded(_) => Bound::Excluded(()),
633		Bound::Unbounded => Bound::Unbounded,
634	}
635}
636
637struct RangeChunkRequest<'a> {
638	table: EntryKind,
639	start: Bound<&'a [u8]>,
640	end: Bound<&'a [u8]>,
641	scope: MultiVersionScope,
642	batch_size: usize,
643	descending: bool,
644}
645
646impl SqlitePersistentStorage {
647	#[instrument(name = "store::multi::persistent::sqlite::get::operator", level = "trace", skip(self), fields(key_len = key.len(), version = version.0))]
648	fn get_operator(&self, table: EntryKind, key: &[u8], version: CommitVersion) -> Result<VersionedGetResult> {
649		self.get_impl(table, key, version)
650	}
651
652	#[instrument(name = "store::multi::persistent::sqlite::get::source", level = "trace", skip(self), fields(key_len = key.len(), version = version.0))]
653	fn get_source(&self, table: EntryKind, key: &[u8], version: CommitVersion) -> Result<VersionedGetResult> {
654		self.get_impl(table, key, version)
655	}
656
657	#[instrument(name = "store::multi::persistent::sqlite::get::operator_internal", level = "trace", skip(self), fields(key_len = key.len(), version = version.0))]
658	fn get_operator_internal(
659		&self,
660		table: EntryKind,
661		key: &[u8],
662		version: CommitVersion,
663	) -> Result<VersionedGetResult> {
664		self.get_impl(table, key, version)
665	}
666
667	#[instrument(name = "store::multi::persistent::sqlite::get::multi", level = "trace", skip(self), fields(key_len = key.len(), version = version.0))]
668	fn get_multi(&self, table: EntryKind, key: &[u8], version: CommitVersion) -> Result<VersionedGetResult> {
669		self.get_impl(table, key, version)
670	}
671
672	fn get_impl(&self, table: EntryKind, key: &[u8], version: CommitVersion) -> Result<VersionedGetResult> {
673		let table_sql = self.table_sql(table);
674		let guard = self.inner.readers.acquire();
675		let Some(conn) = guard.as_ref() else {
676			return Ok(VersionedGetResult::NotFound);
677		};
678
679		let result = match conn.prepare_cached(&table_sql.get_sql) {
680			Ok(mut stmt) => stmt.query_row(params![key], |row| {
681				let version_bytes: Vec<u8> = row.get(0)?;
682				let value: Option<Vec<u8>> = row.get(1)?;
683				Ok((version_from_bytes(&version_bytes), value))
684			}),
685			Err(e) if e.to_string().contains("no such table") => Err(QueryReturnedNoRows),
686			Err(e) => return Err(error!(internal(format!("Failed to prepare persistent get: {}", e)))),
687		};
688
689		match result {
690			Ok((stored_version, value)) if stored_version <= version => Ok(match value {
691				Some(v) => VersionedGetResult::Value {
692					value: CowVec::new(v),
693					version: stored_version,
694				},
695				None => VersionedGetResult::Tombstone,
696			}),
697			Ok(_) => Ok(VersionedGetResult::NotFound),
698			Err(QueryReturnedNoRows) => Ok(VersionedGetResult::NotFound),
699			Err(e) if e.to_string().contains("no such table") => Ok(VersionedGetResult::NotFound),
700			Err(e) => Err(error!(internal(format!("Failed to read persistent: {}", e)))),
701		}
702	}
703
704	#[instrument(name = "store::multi::persistent::sqlite::get_many::operator", level = "trace", skip(self, keys), fields(key_count = keys.len(), version = version.0))]
705	fn get_many_operator(
706		&self,
707		table: EntryKind,
708		keys: &[&[u8]],
709		version: CommitVersion,
710	) -> Result<Vec<VersionedGetResult>> {
711		self.get_many_impl(table, keys, version)
712	}
713
714	#[instrument(name = "store::multi::persistent::sqlite::get_many::source", level = "trace", skip(self, keys), fields(key_count = keys.len(), version = version.0))]
715	fn get_many_source(
716		&self,
717		table: EntryKind,
718		keys: &[&[u8]],
719		version: CommitVersion,
720	) -> Result<Vec<VersionedGetResult>> {
721		self.get_many_impl(table, keys, version)
722	}
723
724	#[instrument(name = "store::multi::persistent::sqlite::get_many::operator_internal", level = "trace", skip(self, keys), fields(key_count = keys.len(), version = version.0))]
725	fn get_many_operator_internal(
726		&self,
727		table: EntryKind,
728		keys: &[&[u8]],
729		version: CommitVersion,
730	) -> Result<Vec<VersionedGetResult>> {
731		self.get_many_impl(table, keys, version)
732	}
733
734	#[instrument(name = "store::multi::persistent::sqlite::get_many::multi", level = "trace", skip(self, keys), fields(key_count = keys.len(), version = version.0))]
735	fn get_many_multi(
736		&self,
737		table: EntryKind,
738		keys: &[&[u8]],
739		version: CommitVersion,
740	) -> Result<Vec<VersionedGetResult>> {
741		self.get_many_impl(table, keys, version)
742	}
743
744	fn get_many_impl(
745		&self,
746		table: EntryKind,
747		keys: &[&[u8]],
748		version: CommitVersion,
749	) -> Result<Vec<VersionedGetResult>> {
750		let mut out = vec![VersionedGetResult::NotFound; keys.len()];
751		if keys.is_empty() {
752			return Ok(out);
753		}
754
755		let index: HashMap<&[u8], usize> = keys.iter().enumerate().map(|(i, &k)| (k, i)).collect();
756		let table_sql = self.table_sql(table);
757		let guard = self.inner.readers.acquire();
758		let Some(conn) = guard.as_ref() else {
759			return Ok(out);
760		};
761
762		for chunk in keys.chunks(GET_MANY_CHUNK) {
763			let bucket = bucket_key_count(chunk.len());
764			let sql = build_get_many_current_sql(&table_sql.table_name, bucket);
765			let mut stmt = match conn.prepare_cached(&sql) {
766				Ok(stmt) => stmt,
767				Err(e) if e.to_string().contains("no such table") => return Ok(out),
768				Err(e) => {
769					return Err(error!(internal(format!(
770						"Failed to prepare persistent get_many: {}",
771						e
772					))));
773				}
774			};
775
776			let pad_key = chunk[0];
777			let padded = chunk.iter().copied().chain(repeat_n(pad_key, bucket - chunk.len()));
778			let mut rows = stmt
779				.query(params_from_iter(padded))
780				.map_err(|e| error!(internal(format!("Failed to query persistent get_many: {}", e))))?;
781
782			while let Some(row) = rows.next().map_err(|e| {
783				error!(internal(format!("Failed to read persistent get_many row: {}", e)))
784			})? {
785				let key_ref = row.get_ref(0).map_err(|e| {
786					error!(internal(format!("Failed to read persistent get_many key: {}", e)))
787				})?;
788				let key = key_ref.as_blob().map_err(|e| {
789					error!(internal(format!("Failed to decode persistent get_many key: {}", e)))
790				})?;
791				let Some(&i) = index.get(key) else {
792					continue;
793				};
794				let version_ref = row.get_ref(1).map_err(|e| {
795					error!(internal(format!("Failed to read persistent get_many version: {}", e)))
796				})?;
797				let version_bytes = version_ref.as_blob().map_err(|e| {
798					error!(internal(format!("Failed to decode persistent get_many version: {}", e)))
799				})?;
800				let stored_version = version_from_bytes(version_bytes);
801				if stored_version > version {
802					continue;
803				}
804				let value: Option<Vec<u8>> = row.get(2).map_err(|e| {
805					error!(internal(format!("Failed to read persistent get_many value: {}", e)))
806				})?;
807				out[i] = match value {
808					Some(v) => VersionedGetResult::Value {
809						value: CowVec::new(v),
810						version: stored_version,
811					},
812					None => VersionedGetResult::Tombstone,
813				};
814			}
815		}
816
817		Ok(out)
818	}
819}
820
821impl TierStorage for SqlitePersistentStorage {
822	fn get(&self, table: EntryKind, key: &[u8], version: CommitVersion) -> Result<VersionedGetResult> {
823		match table {
824			EntryKind::Operator(_) => self.get_operator(table, key, version),
825			EntryKind::OperatorInternal(_) => self.get_operator_internal(table, key, version),
826			EntryKind::Source(_) => self.get_source(table, key, version),
827			_ => self.get_multi(table, key, version),
828		}
829	}
830
831	fn get_many(
832		&self,
833		table: EntryKind,
834		keys: &[&[u8]],
835		version: CommitVersion,
836	) -> Result<Vec<VersionedGetResult>> {
837		match table {
838			EntryKind::Operator(_) => self.get_many_operator(table, keys, version),
839			EntryKind::OperatorInternal(_) => self.get_many_operator_internal(table, keys, version),
840			EntryKind::Source(_) => self.get_many_source(table, keys, version),
841			_ => self.get_many_multi(table, keys, version),
842		}
843	}
844
845	fn set(&self, version: CommitVersion, batches: TierBatch) -> Result<()> {
846		self.set_collecting_accepted(version, batches)?;
847		Ok(())
848	}
849
850	#[instrument(name = "store::multi::persistent::sqlite::range", level = "trace", skip(self, cursor, start, end), fields(table = ?table, batch_size = batch_size))]
851	fn range_next(
852		&self,
853		table: EntryKind,
854		cursor: &mut RangeCursor,
855		start: Bound<&[u8]>,
856		end: Bound<&[u8]>,
857		scope: MultiVersionScope,
858		batch_size: usize,
859	) -> Result<RangeBatch> {
860		self.range_chunk(
861			cursor,
862			RangeChunkRequest {
863				table,
864				start,
865				end,
866				scope,
867				batch_size,
868				descending: false,
869			},
870		)
871	}
872
873	#[instrument(name = "store::multi::persistent::sqlite::range_rev", level = "trace", skip(self, cursor, start, end), fields(table = ?table, batch_size = batch_size))]
874	fn range_rev_next(
875		&self,
876		table: EntryKind,
877		cursor: &mut RangeCursor,
878		start: Bound<&[u8]>,
879		end: Bound<&[u8]>,
880		scope: MultiVersionScope,
881		batch_size: usize,
882	) -> Result<RangeBatch> {
883		self.range_chunk(
884			cursor,
885			RangeChunkRequest {
886				table,
887				start,
888				end,
889				scope,
890				batch_size,
891				descending: true,
892			},
893		)
894	}
895
896	fn ensure_table(&self, table: EntryKind) -> Result<()> {
897		let table_sql = self.table_sql(table);
898		let guard = self.inner.conn.lock();
899		let Some(conn) = guard.as_ref() else {
900			return Ok(());
901		};
902		Self::create_table_if_needed(conn, &table_sql.create_sql)
903			.map_err(|e| error!(internal(format!("Failed to ensure persistent table: {}", e))))
904	}
905
906	fn clear_table(&self, table: EntryKind) -> Result<()> {
907		let table_sql = self.table_sql(table);
908		let guard = self.inner.conn.lock();
909		let Some(conn) = guard.as_ref() else {
910			return Ok(());
911		};
912		let result = conn.execute(&format!("DELETE FROM \"{}\"", table_sql.table_name), []);
913		if let Err(e) = result
914			&& !e.to_string().contains("no such table")
915		{
916			return Err(error!(internal(format!(
917				"Failed to clear persistent {}: {}",
918				table_sql.table_name, e
919			))));
920		}
921		Ok(())
922	}
923
924	fn drop(&self, _batches: HashMap<EntryKind, Vec<(EncodedKey, CommitVersion)>>) -> Result<()> {
925		// TODO: change the TierStorage interface so persistent doesn't have to expose
926
927		panic!("SqlitePersistentStorage::drop: persistent tier has no historical chain to drop versions from");
928	}
929
930	#[instrument(name = "store::multi::persistent::sqlite::get_all_versions", level = "trace", skip(self, key), fields(table = ?table, key_len = key.len()))]
931	fn get_all_versions(&self, table: EntryKind, key: &[u8]) -> Result<Vec<(CommitVersion, Option<CowVec<u8>>)>> {
932		// TODO: change the TierStorage interface to remove the
933
934		let table_sql = self.table_sql(table);
935		let guard = self.inner.readers.acquire();
936		let Some(conn) = guard.as_ref() else {
937			return Ok(Vec::new());
938		};
939
940		let result = match conn.prepare_cached(&table_sql.get_sql) {
941			Ok(mut stmt) => stmt.query_row(params![key], |row| {
942				let version_bytes: Vec<u8> = row.get(0)?;
943				let value: Option<Vec<u8>> = row.get(1)?;
944				Ok((version_from_bytes(&version_bytes), value.map(CowVec::new)))
945			}),
946			Err(e) if e.to_string().contains("no such table") => return Ok(Vec::new()),
947			Err(e) => {
948				return Err(error!(internal(format!(
949					"Failed to prepare persistent get_all_versions: {}",
950					e
951				))));
952			}
953		};
954
955		match result {
956			Ok(row) => Ok(vec![row]),
957			Err(QueryReturnedNoRows) => Ok(Vec::new()),
958			Err(e) if e.to_string().contains("no such table") => Ok(Vec::new()),
959			Err(e) => Err(error!(internal(format!("Failed to read persistent versions: {}", e)))),
960		}
961	}
962
963	fn scan_historical_below(
964		&self,
965		_table: EntryKind,
966		_cutoff: CommitVersion,
967		_cursor: &mut HistoricalCursor,
968		_batch_size: usize,
969	) -> Result<Vec<(EncodedKey, CommitVersion)>> {
970		// TODO: change the TierStorage interface so persistent doesn't have to expose
971
972		panic!("SqlitePersistentStorage::scan_historical_below: persistent tier has no historical chain");
973	}
974}
975
976impl TierBackend for SqlitePersistentStorage {}
977
978impl Shutdown for SqlitePersistentStorage {
979	fn shutdown(&self) {
980		if let Some(conn) = self.inner.conn.lock().take() {
981			if let Err(e) = pragma::shutdown(&conn) {
982				warn!(error = %e, "persistent close: pragma shutdown failed");
983			}
984			drop(conn);
985		}
986		self.inner.readers.shutdown();
987	}
988}
989
990#[cfg(test)]
991mod tests {
992	use std::collections::HashMap;
993
994	use reifydb_core::interface::catalog::{id::TableId, shape::ShapeId};
995
996	use super::*;
997
998	fn table() -> EntryKind {
999		EntryKind::Source(ShapeId::Table(TableId(1)))
1000	}
1001
1002	fn key(n: u64) -> EncodedKey {
1003		EncodedKey::new(n.to_be_bytes().to_vec())
1004	}
1005
1006	fn row(payload: &[u8]) -> CowVec<u8> {
1007		CowVec::new(payload.to_vec())
1008	}
1009
1010	fn visible(s: &SqlitePersistentStorage, k: &EncodedKey) -> bool {
1011		s.get(table(), k.as_slice(), CommitVersion(u64::MAX)).unwrap().value().is_some()
1012	}
1013
1014	#[test]
1015	fn delete_below_version_removes_rows_at_or_below_cutoff() {
1016		let (s, _guard) = SqlitePersistentStorage::in_memory();
1017		// Each key written at a distinct commit version (separate set calls).
1018		s.set(CommitVersion(1), HashMap::from([(table(), vec![(key(1), Some(row(b"a")))])])).unwrap();
1019		s.set(CommitVersion(2), HashMap::from([(table(), vec![(key(2), Some(row(b"b")))])])).unwrap();
1020		s.set(CommitVersion(3), HashMap::from([(table(), vec![(key(3), Some(row(b"c")))])])).unwrap();
1021		assert_eq!(s.count_current(table()).unwrap(), 3);
1022
1023		let deleted = s.delete_below_version(table(), CommitVersion(2), None).unwrap();
1024
1025		assert_eq!(deleted.len(), 2, "rows whose version is <= cutoff(2) must be physically deleted");
1026		assert_eq!(
1027			s.count_current(table()).unwrap(),
1028			1,
1029			"deletion must reclaim sqlite rows, not tombstone them"
1030		);
1031		assert!(!visible(&s, &key(1)));
1032		assert!(!visible(&s, &key(2)));
1033		assert!(visible(&s, &key(3)), "a row written after the cutoff version must survive");
1034	}
1035
1036	#[test]
1037	fn create_table_indexes_the_version_column() {
1038		// Phase 2b: after the created_nanos/updated_nanos indices were dropped, the version-anchored TTL
1039		// delete (DELETE WHERE version <= cutoff) needs an index on `version`, or GC full-scans the live
1040		// set on every tick. The two timestamp indices must stay gone.
1041		let (s, _guard) = SqlitePersistentStorage::in_memory();
1042		s.set(CommitVersion(1), HashMap::from([(table(), vec![(key(1), Some(row(b"a")))])])).unwrap();
1043
1044		let table_name = s.table_sql(table()).table_name.clone();
1045		let guard = s.inner.conn.lock();
1046		let conn = guard.as_ref().expect("write connection is present");
1047
1048		let indices: Vec<String> = conn
1049			.prepare("SELECT name FROM sqlite_master WHERE type = 'index' AND tbl_name = ?1")
1050			.unwrap()
1051			.query_map([table_name.as_str()], |r| r.get::<_, String>(0))
1052			.unwrap()
1053			.map(|r| r.unwrap())
1054			.collect();
1055
1056		assert!(
1057			indices.contains(&format!("{table_name}__version")),
1058			"the version column must be indexed so the TTL delete seeks instead of scanning, got {indices:?}"
1059		);
1060		assert!(
1061			!indices.iter().any(|n| n.ends_with("__created_nanos") || n.ends_with("__updated_nanos")),
1062			"the dropped timestamp indices must not be recreated, got {indices:?}"
1063		);
1064	}
1065
1066	#[test]
1067	fn delete_below_version_keeps_rows_written_after_the_cutoff() {
1068		let (s, _guard) = SqlitePersistentStorage::in_memory();
1069		s.set(CommitVersion(2), HashMap::from([(table(), vec![(key(2), Some(row(b"stale")))])])).unwrap();
1070		s.set(CommitVersion(5), HashMap::from([(table(), vec![(key(1), Some(row(b"fresh")))])])).unwrap();
1071
1072		let deleted = s.delete_below_version(table(), CommitVersion(3), None).unwrap();
1073
1074		assert_eq!(deleted.len(), 1, "only the row whose last write is at or below the cutoff is evicted");
1075		assert!(visible(&s, &key(1)), "a row written after the cutoff version must NOT be evicted");
1076		assert!(!visible(&s, &key(2)));
1077	}
1078
1079	#[test]
1080	fn delete_below_version_boundary_is_inclusive() {
1081		let (s, _guard) = SqlitePersistentStorage::in_memory();
1082		s.set(CommitVersion(5), HashMap::from([(table(), vec![(key(1), Some(row(b"v5")))])])).unwrap();
1083
1084		// Cutoff exactly equal to the row's version: the row IS deleted (version <= cutoff).
1085		let deleted = s.delete_below_version(table(), CommitVersion(5), None).unwrap();
1086		assert_eq!(
1087			deleted.len(),
1088			1,
1089			"a row whose version equals the cutoff is evicted (the bound is inclusive)"
1090		);
1091		assert!(!visible(&s, &key(1)));
1092	}
1093
1094	#[test]
1095	fn delete_below_version_on_missing_table_is_noop() {
1096		let (s, _guard) = SqlitePersistentStorage::in_memory();
1097		let deleted = s
1098			.delete_below_version(EntryKind::Source(ShapeId::Table(TableId(999))), CommitVersion(100), None)
1099			.unwrap();
1100		assert_eq!(deleted.len(), 0);
1101	}
1102
1103	#[test]
1104	fn delete_below_version_with_prefix_only_touches_matching_keys() {
1105		let (s, _guard) = SqlitePersistentStorage::in_memory();
1106		// Two "sides" distinguished by a leading prefix byte, both written at v1.
1107		let left = EncodedKey::new(vec![0x01, 0xAA]);
1108		let right = EncodedKey::new(vec![0x02, 0xBB]);
1109		s.set(
1110			CommitVersion(1),
1111			HashMap::from([(
1112				table(),
1113				vec![(left.clone(), Some(row(b"l"))), (right.clone(), Some(row(b"r")))],
1114			)]),
1115		)
1116		.unwrap();
1117
1118		let deleted = s.delete_below_version(table(), CommitVersion(2), Some(&[0x01])).unwrap();
1119
1120		assert_eq!(deleted.len(), 1, "only the 0x01-prefixed (left) row should be deleted");
1121		assert!(!visible(&s, &left));
1122		assert!(visible(&s, &right), "the 0x02-prefixed (right) row must survive a left-only prefix sweep");
1123	}
1124
1125	#[test]
1126	fn delete_below_version_returns_exactly_the_deleted_keys() {
1127		let (s, _guard) = SqlitePersistentStorage::in_memory();
1128		s.set(CommitVersion(1), HashMap::from([(table(), vec![(key(1), Some(row(b"a")))])])).unwrap();
1129		s.set(CommitVersion(2), HashMap::from([(table(), vec![(key(2), Some(row(b"b")))])])).unwrap();
1130		s.set(CommitVersion(3), HashMap::from([(table(), vec![(key(3), Some(row(b"c")))])])).unwrap();
1131
1132		// The surgical GC invalidation depends on delete_below_version returning the exact keys it deleted,
1133		// so the read cache is invalidated per-key instead of cleared wholesale. A wrong/empty key set
1134		// would silently leave stale entries (or over-clear) and this assertion would catch it.
1135		let mut got: Vec<Vec<u8>> = s
1136			.delete_below_version(table(), CommitVersion(2), None)
1137			.unwrap()
1138			.iter()
1139			.map(|k| k.to_vec())
1140			.collect();
1141		got.sort();
1142		let mut want = vec![key(1).to_vec(), key(2).to_vec()];
1143		want.sort();
1144		assert_eq!(
1145			got, want,
1146			"delete_below_version must return every key it physically deleted, and only those"
1147		);
1148		assert!(visible(&s, &key(3)), "the row newer than the cutoff must neither be deleted nor returned");
1149	}
1150}
reifydb_store_multi/tier/persistent/sqlite/storage.rs

reifydb_store_multi/tier/persistent/sqlite/
storage.rs
