surfpool_core/storage/
sqlite.rs

1use std::{
2    collections::{HashMap, HashSet},
3    sync::{
4        Mutex, OnceLock,
5        atomic::{AtomicU64, Ordering},
6    },
7};
8
9use log::debug;
10use serde::{Deserialize, Serialize};
11use surfpool_db::diesel::{
12    self, RunQueryDsl,
13    connection::SimpleConnection,
14    r2d2::{ConnectionManager, Pool},
15    sql_query,
16    sql_types::Text,
17};
18
19use crate::storage::{
20    Storage, StorageConstructor, StorageError, StorageResult,
21    diesel_common::{
22        CountRecord, KeyRecord, KvRecord, ValueRecord, deserialize_value, serialize_key,
23        serialize_value,
24    },
25};
26
27/// Applies pragmas when each pool connection is created.
28#[derive(Debug)]
29struct SqlitePragmaCustomizer {
30    is_file_based: bool,
31}
32
33impl diesel::r2d2::CustomizeConnection<diesel::SqliteConnection, diesel::r2d2::Error>
34    for SqlitePragmaCustomizer
35{
36    fn on_acquire(&self, conn: &mut diesel::SqliteConnection) -> Result<(), diesel::r2d2::Error> {
37        let pragmas = if self.is_file_based {
38            "PRAGMA synchronous=NORMAL; PRAGMA temp_store=MEMORY; PRAGMA mmap_size=268435456; PRAGMA cache_size=-64000; PRAGMA busy_timeout=5000;"
39        } else {
40            "PRAGMA synchronous=OFF; PRAGMA temp_store=MEMORY; PRAGMA cache_size=-64000; PRAGMA busy_timeout=5000;"
41        };
42        conn.batch_execute(pragmas)
43            .map_err(diesel::r2d2::Error::QueryError)
44    }
45}
46
47/// Track which database files have already been checkpointed during shutdown.
48/// This prevents multiple SqliteStorage instances sharing the same file from
49/// conflicting when each tries to checkpoint and delete WAL files.
50fn checkpointed_databases() -> &'static Mutex<HashSet<String>> {
51    static CHECKPOINTED: OnceLock<Mutex<HashSet<String>>> = OnceLock::new();
52    CHECKPOINTED.get_or_init(|| Mutex::new(HashSet::new()))
53}
54
55/// Shared pools keyed by connection string (file-based DBs only).
56static SHARED_POOLS: OnceLock<
57    Mutex<HashMap<String, Pool<ConnectionManager<diesel::SqliteConnection>>>>,
58> = OnceLock::new();
59
60/// Counter for unique in-memory database names.
61static MEMORY_DB_COUNTER: AtomicU64 = AtomicU64::new(0);
62
63fn get_or_create_shared_pool(
64    connection_string: &str,
65    is_file_based: bool,
66) -> StorageResult<Pool<ConnectionManager<diesel::SqliteConnection>>> {
67    // In-memory DBs get isolated pools
68    if !is_file_based {
69        let manager = ConnectionManager::<diesel::SqliteConnection>::new(connection_string);
70        return Pool::builder()
71            .max_size(10)
72            .connection_customizer(Box::new(SqlitePragmaCustomizer { is_file_based }))
73            .build(manager)
74            .map_err(|e| StorageError::PooledConnectionError(NAME.into(), e));
75    }
76
77    let pools = SHARED_POOLS.get_or_init(|| Mutex::new(HashMap::new()));
78    let mut guard = pools.lock().map_err(|_| StorageError::LockError)?;
79
80    if let Some(pool) = guard.get(connection_string) {
81        debug!("Reusing shared SQLite pool for {}", connection_string);
82        return Ok(pool.clone());
83    }
84
85    debug!("Creating shared SQLite pool for {}", connection_string);
86    let manager = ConnectionManager::<diesel::SqliteConnection>::new(connection_string);
87    let pool = Pool::builder()
88        .max_size(10)
89        .connection_customizer(Box::new(SqlitePragmaCustomizer { is_file_based }))
90        .build(manager)
91        .map_err(|e| StorageError::PooledConnectionError(NAME.into(), e))?;
92
93    // journal_mode=WAL persists to file; wal_autocheckpoint is per-connection
94    {
95        let mut conn = pool.get().map_err(|_| StorageError::LockError)?;
96        conn.batch_execute("PRAGMA journal_mode=WAL; PRAGMA wal_autocheckpoint=1000;")
97            .map_err(|e| StorageError::create_table("pragma_init", NAME, e))?;
98    }
99
100    guard.insert(connection_string.to_string(), pool.clone());
101    Ok(pool)
102}
103
104#[derive(Clone)]
105pub struct SqliteStorage<K, V> {
106    pool: Pool<ConnectionManager<diesel::SqliteConnection>>,
107    _phantom: std::marker::PhantomData<(K, V)>,
108    table_name: String,
109    surfnet_id: String,
110    /// Whether this is a file-based database (not :memory:)
111    /// Used to determine if WAL checkpoint should be performed on drop
112    is_file_based: bool,
113    /// The connection string for creating direct connections during cleanup
114    connection_string: String,
115}
116
117const NAME: &str = "SQLite";
118
119// Checkpoint implementation that doesn't require K, V bounds
120impl<K, V> SqliteStorage<K, V> {
121    /// Checkpoint the WAL and truncate it to consolidate into the main database file,
122    /// then remove the -wal and -shm files.
123    /// Only runs for file-based databases (not :memory:).
124    /// Uses a static set to track which databases have been checkpointed to avoid
125    /// conflicts when multiple SqliteStorage instances share the same database file.
126    fn checkpoint(&self) {
127        if !self.is_file_based {
128            return;
129        }
130
131        // Extract the file path from the connection string
132        // Connection string is like "file:/path/to/db.sqlite?mode=rwc"
133        let db_path = self
134            .connection_string
135            .strip_prefix("file:")
136            .and_then(|s| s.split('?').next())
137            .unwrap_or(&self.connection_string)
138            .to_string();
139
140        // Check if this database has already been checkpointed by another storage instance
141        {
142            let mut checkpointed = checkpointed_databases().lock().unwrap();
143            if checkpointed.contains(&db_path) {
144                debug!(
145                    "Database {} already checkpointed, skipping for table '{}'",
146                    db_path, self.table_name
147                );
148                return;
149            }
150            checkpointed.insert(db_path.clone());
151        }
152
153        debug!(
154            "Checkpointing WAL for database '{}' (table '{}')",
155            db_path, self.table_name
156        );
157
158        // Use pool connection to checkpoint - this flushes WAL to main database
159        if let Ok(mut conn) = self.pool.get() {
160            if let Err(e) = conn.batch_execute("PRAGMA wal_checkpoint(TRUNCATE);") {
161                debug!("WAL checkpoint failed: {}", e);
162                return;
163            }
164        }
165
166        // Remove the -wal and -shm files
167        let wal_path = format!("{}-wal", db_path);
168        let shm_path = format!("{}-shm", db_path);
169
170        if std::path::Path::new(&wal_path).exists() {
171            if let Err(e) = std::fs::remove_file(&wal_path) {
172                debug!("Failed to remove WAL file {}: {}", wal_path, e);
173            } else {
174                debug!("Removed WAL file: {}", wal_path);
175            }
176        }
177
178        if std::path::Path::new(&shm_path).exists() {
179            if let Err(e) = std::fs::remove_file(&shm_path) {
180                debug!("Failed to remove SHM file {}: {}", shm_path, e);
181            } else {
182                debug!("Removed SHM file: {}", shm_path);
183            }
184        }
185    }
186}
187
188impl<K, V> SqliteStorage<K, V>
189where
190    K: Serialize + for<'de> Deserialize<'de>,
191    V: Serialize + for<'de> Deserialize<'de> + Clone,
192{
193    fn ensure_table_exists(&self) -> StorageResult<()> {
194        debug!("Ensuring table '{}' exists", self.table_name);
195        let create_table_sql = format!(
196            "
197            CREATE TABLE IF NOT EXISTS {} (
198                surfnet_id TEXT NOT NULL,
199                key TEXT NOT NULL,
200                value TEXT NOT NULL,
201                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
202                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
203                PRIMARY KEY (surfnet_id, key)
204            )
205        ",
206            self.table_name
207        );
208
209        debug!("Getting connection from pool for table creation");
210        let mut conn = self.pool.get().map_err(|_| StorageError::LockError)?;
211
212        conn.batch_execute(&create_table_sql)
213            .map_err(|e| StorageError::create_table(&self.table_name, NAME, e))?;
214
215        debug!("Successfully ensured table '{}' exists", self.table_name);
216        Ok(())
217    }
218
219    fn load_value_from_db(&self, key_str: &str) -> StorageResult<Option<V>> {
220        debug!("Loading value from DB for key: {}", key_str);
221        let query = sql_query(format!(
222            "SELECT value FROM {} WHERE surfnet_id = ? AND key = ?",
223            self.table_name
224        ))
225        .bind::<Text, _>(&self.surfnet_id)
226        .bind::<Text, _>(key_str);
227
228        trace!("Getting connection from pool for loading value");
229        let mut conn = self.pool.get().map_err(|_| StorageError::LockError)?;
230
231        let records = query
232            .load::<ValueRecord>(&mut *conn)
233            .map_err(|e| StorageError::get(&self.table_name, NAME, key_str, e))?;
234
235        if let Some(record) = records.into_iter().next() {
236            debug!("Found record for key: {}", key_str);
237            let value = deserialize_value(NAME, &self.table_name, &record.value)?;
238            Ok(Some(value))
239        } else {
240            debug!("No record found for key: {}", key_str);
241            Ok(None)
242        }
243    }
244}
245
246impl<K, V> Storage<K, V> for SqliteStorage<K, V>
247where
248    K: Serialize + for<'de> Deserialize<'de> + Clone + Send + Sync + 'static,
249    V: Serialize + for<'de> Deserialize<'de> + Clone + Send + Sync + 'static,
250{
251    fn store(&mut self, key: K, value: V) -> StorageResult<()> {
252        debug!("Storing value in table '{}", self.table_name);
253        let key_str = serialize_key(NAME, &self.table_name, &key)?;
254        let value_str = serialize_value(NAME, &self.table_name, &value)?;
255
256        // Use prepared statement with sql_query for better safety
257        let query = sql_query(format!(
258            "INSERT OR REPLACE INTO {} (surfnet_id, key, value, updated_at) VALUES (?, ?, ?, CURRENT_TIMESTAMP)",
259            self.table_name
260        ))
261        .bind::<Text, _>(&self.surfnet_id)
262        .bind::<Text, _>(&key_str)
263        .bind::<Text, _>(&value_str);
264
265        trace!("Getting connection from pool for store operation");
266        let mut conn = self.pool.get().map_err(|_| StorageError::LockError)?;
267
268        query
269            .execute(&mut *conn)
270            .map_err(|e| StorageError::store(&self.table_name, NAME, &key_str, e))?;
271
272        debug!("Value stored successfully in table '{}'", self.table_name);
273        Ok(())
274    }
275
276    fn get(&self, key: &K) -> StorageResult<Option<V>> {
277        debug!("Getting value from table '{}", self.table_name);
278        let key_str = serialize_key(NAME, &self.table_name, key)?;
279
280        self.load_value_from_db(&key_str)
281    }
282
283    fn take(&mut self, key: &K) -> StorageResult<Option<V>> {
284        debug!("Taking value from table '{}'", self.table_name);
285        let key_str = serialize_key(NAME, &self.table_name, key)?;
286
287        // If not in cache, try to load from database
288        if let Some(value) = self.load_value_from_db(&key_str)? {
289            debug!("Value found, removing from database");
290            // Remove from database
291            let delete_query = sql_query(format!(
292                "DELETE FROM {} WHERE surfnet_id = ? AND key = ?",
293                self.table_name
294            ))
295            .bind::<Text, _>(&self.surfnet_id)
296            .bind::<Text, _>(&key_str);
297
298            trace!("Getting connection from pool for delete operation");
299            let mut conn = self.pool.get().map_err(|_| StorageError::LockError)?;
300
301            delete_query
302                .execute(&mut *conn)
303                .map_err(|e| StorageError::delete(&self.table_name, NAME, &key_str, e))?;
304
305            debug!(
306                "Value taken and removed successfully from table '{}'",
307                self.table_name
308            );
309            Ok(Some(value))
310        } else {
311            debug!("No value found to take from table '{}'", self.table_name);
312            Ok(None)
313        }
314    }
315
316    fn clear(&mut self) -> StorageResult<()> {
317        debug!("Clearing all data from table '{}'", self.table_name);
318        let delete_query = sql_query(format!(
319            "DELETE FROM {} WHERE surfnet_id = ?",
320            self.table_name
321        ))
322        .bind::<Text, _>(&self.surfnet_id);
323
324        trace!("Getting connection from pool for clear operation");
325        let mut conn = self.pool.get().map_err(|_| StorageError::LockError)?;
326
327        delete_query
328            .execute(&mut *conn)
329            .map_err(|e| StorageError::delete(&self.table_name, NAME, "*all*", e))?;
330
331        debug!("Table '{}' cleared successfully", self.table_name);
332        Ok(())
333    }
334
335    fn keys(&self) -> StorageResult<Vec<K>> {
336        debug!("Fetching all keys from table '{}'", self.table_name);
337        let query = sql_query(format!(
338            "SELECT key FROM {} WHERE surfnet_id = ?",
339            self.table_name
340        ))
341        .bind::<Text, _>(&self.surfnet_id);
342
343        trace!("Getting connection from pool for keys operation");
344        let mut conn = self.pool.get().map_err(|_| StorageError::LockError)?;
345
346        let records = query
347            .load::<KeyRecord>(&mut *conn)
348            .map_err(|e| StorageError::get_all_keys(&self.table_name, NAME, e))?;
349
350        let mut keys = Vec::new();
351        for record in records {
352            let key: K = serde_json::from_str(&record.key)
353                .map_err(|e| StorageError::DeserializeValueError(NAME.into(), e))?;
354            keys.push(key);
355        }
356
357        debug!(
358            "Retrieved {} keys from table '{}'",
359            keys.len(),
360            self.table_name
361        );
362        Ok(keys)
363    }
364
365    fn clone_box(&self) -> Box<dyn Storage<K, V>> {
366        Box::new(self.clone())
367    }
368
369    fn shutdown(&self) {
370        self.checkpoint();
371    }
372
373    fn count(&self) -> StorageResult<u64> {
374        debug!("Counting entries in table '{}'", self.table_name);
375        let query = sql_query(format!(
376            "SELECT COUNT(*) as count FROM {} WHERE surfnet_id = ?",
377            self.table_name
378        ))
379        .bind::<Text, _>(&self.surfnet_id);
380
381        trace!("Getting connection from pool for count operation");
382        let mut conn = self.pool.get().map_err(|_| StorageError::LockError)?;
383
384        let records = query
385            .load::<CountRecord>(&mut *conn)
386            .map_err(|e| StorageError::count(&self.table_name, NAME, e))?;
387
388        let count = records.first().map(|r| r.count as u64).unwrap_or(0);
389        debug!("Table '{}' has {} entries", self.table_name, count);
390        Ok(count)
391    }
392
393    fn into_iter(&self) -> StorageResult<Box<dyn Iterator<Item = (K, V)> + '_>> {
394        debug!(
395            "Creating iterator for all key-value pairs in table '{}'",
396            self.table_name
397        );
398        let query = sql_query(format!(
399            "SELECT key, value FROM {} WHERE surfnet_id = ?",
400            self.table_name
401        ))
402        .bind::<Text, _>(&self.surfnet_id);
403
404        trace!("Getting connection from pool for into_iter operation");
405        let mut conn = self.pool.get().map_err(|_| StorageError::LockError)?;
406
407        let records = query
408            .load::<KvRecord>(&mut *conn)
409            .map_err(|e| StorageError::get_all_key_value_pairs(&self.table_name, NAME, e))?;
410
411        let iter = records.into_iter().filter_map(move |record| {
412            let key: K = match serde_json::from_str(&record.key) {
413                Ok(k) => k,
414                Err(e) => {
415                    debug!("Failed to deserialize key: {}", e);
416                    return None;
417                }
418            };
419            let value: V = match serde_json::from_str(&record.value) {
420                Ok(v) => v,
421                Err(e) => {
422                    debug!("Failed to deserialize value: {}", e);
423                    return None;
424                }
425            };
426            Some((key, value))
427        });
428
429        debug!(
430            "Iterator created successfully for table '{}'",
431            self.table_name
432        );
433        Ok(Box::new(iter))
434    }
435}
436
437impl<K, V> StorageConstructor<K, V> for SqliteStorage<K, V>
438where
439    K: Serialize + for<'de> Deserialize<'de> + Clone + Send + Sync + 'static,
440    V: Serialize + for<'de> Deserialize<'de> + Clone + Send + Sync + 'static,
441{
442    fn connect(database_url: &str, table_name: &str, surfnet_id: &str) -> StorageResult<Self> {
443        debug!(
444            "Connecting to SQLite database: {} with table: {} and surfnet_id: {}",
445            database_url, table_name, surfnet_id
446        );
447
448        let connection_string = if database_url == ":memory:" {
449            // Unique name per storage instance; cache=shared so pool connections share it
450            let id = MEMORY_DB_COUNTER.fetch_add(1, Ordering::Relaxed);
451            format!("file:memdb{}?mode=memory&cache=shared", id)
452        } else if database_url.starts_with("file:") {
453            if database_url.contains('?') {
454                format!("{}&mode=rwc", database_url)
455            } else {
456                format!("{}?mode=rwc", database_url)
457            }
458        } else {
459            format!("file:{}?mode=rwc", database_url)
460        };
461
462        let is_file_based = database_url != ":memory:";
463        let pool = get_or_create_shared_pool(&connection_string, is_file_based)?;
464
465        let storage = SqliteStorage {
466            pool,
467            _phantom: std::marker::PhantomData,
468            table_name: table_name.to_string(),
469            surfnet_id: surfnet_id.to_string(),
470            is_file_based,
471            connection_string,
472        };
473
474        storage.ensure_table_exists()?;
475        debug!(
476            "SQLite storage connected successfully for table: {}",
477            table_name
478        );
479        Ok(storage)
480    }
481}
482
483#[cfg(test)]
484mod tests {
485    use super::*;
486
487    #[derive(diesel::QueryableByName, Debug)]
488    struct PragmaInt {
489        #[diesel(sql_type = diesel::sql_types::BigInt)]
490        value: i64,
491    }
492
493    /// SqliteStorage instances pointing to the same file share one connection pool.
494    #[test]
495    fn test_shared_pool_across_multiple_storages() {
496        let temp_file = tempfile::NamedTempFile::new().unwrap();
497        let db_path = temp_file.path().to_str().unwrap();
498
499        let storage1: SqliteStorage<String, String> =
500            SqliteStorage::connect(db_path, "table1", "surfnet1").unwrap();
501        let storage2: SqliteStorage<String, String> =
502            SqliteStorage::connect(db_path, "table2", "surfnet1").unwrap();
503        let storage3: SqliteStorage<String, String> =
504            SqliteStorage::connect(db_path, "table3", "surfnet1").unwrap();
505
506        let _conn1 = storage1.pool.get().unwrap();
507
508        let state1 = storage1.pool.state();
509        let state2 = storage2.pool.state();
510        let state3 = storage3.pool.state();
511
512        assert_eq!(
513            state1.connections, state2.connections,
514            "pools should be shared"
515        );
516        assert_eq!(
517            state2.connections, state3.connections,
518            "pools should be shared"
519        );
520        assert_eq!(
521            state1.idle_connections, state2.idle_connections,
522            "pools should be shared"
523        );
524    }
525
526    /// Pragmas are applied to all pooled connections, not just the first.
527    #[test]
528    fn test_pragmas_applied_to_pooled_connections() {
529        let temp_file = tempfile::NamedTempFile::new().unwrap();
530        let db_path = temp_file.path().to_str().unwrap();
531
532        let storage: SqliteStorage<String, String> =
533            SqliteStorage::connect(db_path, "pragma_test", "surfnet1").unwrap();
534
535        let _conn1 = storage.pool.get().unwrap();
536        let mut conn2 = storage.pool.get().unwrap();
537
538        let cache_result: Vec<PragmaInt> =
539            diesel::sql_query("SELECT cache_size as value FROM pragma_cache_size")
540                .load(&mut *conn2)
541                .unwrap();
542        assert_eq!(cache_result[0].value, -64000, "cache_size should be -64000");
543
544        let timeout_result: Vec<PragmaInt> =
545            diesel::sql_query("SELECT timeout as value FROM pragma_busy_timeout")
546                .load(&mut *conn2)
547                .unwrap();
548        assert_eq!(timeout_result[0].value, 5000, "busy_timeout should be 5000");
549    }
550
551    /// File-specific pragmas (synchronous, temp_store) are applied to all connections.
552    #[test]
553    fn test_pragmas_applied_to_file_database() {
554        let temp_file = tempfile::NamedTempFile::new().unwrap();
555        let db_path = temp_file.path().to_str().unwrap();
556
557        let storage: SqliteStorage<String, String> =
558            SqliteStorage::connect(db_path, "pragma_test", "surfnet1").unwrap();
559
560        let _conn1 = storage.pool.get().unwrap();
561        let mut conn2 = storage.pool.get().unwrap();
562
563        let result: Vec<PragmaInt> =
564            diesel::sql_query("SELECT synchronous as value FROM pragma_synchronous")
565                .load(&mut *conn2)
566                .unwrap();
567        assert_eq!(result[0].value, 1, "synchronous should be NORMAL (1)");
568
569        let temp_result: Vec<PragmaInt> =
570            diesel::sql_query("SELECT temp_store as value FROM pragma_temp_store")
571                .load(&mut *conn2)
572                .unwrap();
573        assert_eq!(temp_result[0].value, 2, "temp_store should be MEMORY (2)");
574    }
575
576    /// In-memory databases are isolated between SqliteStorage instances.
577    #[test]
578    fn test_in_memory_databases_are_isolated() {
579        let mut storage1: SqliteStorage<String, String> =
580            SqliteStorage::connect(":memory:", "test_table", "surfnet1").unwrap();
581        let storage2: SqliteStorage<String, String> =
582            SqliteStorage::connect(":memory:", "test_table", "surfnet1").unwrap();
583
584        storage1
585            .store("key1".to_string(), "value1".to_string())
586            .unwrap();
587
588        let result = storage2.get(&"key1".to_string()).unwrap();
589        assert!(result.is_none(), "in-memory databases should be isolated");
590    }
591}
surfpool_core/storage/sqlite.rs

surfpool_core/storage/
sqlite.rs
