Generalise database

Author: Christopher Skene

src/models/middleware/types/jmix_index.rs

@@ -1,4 +1,4 @@
-use once_cell::sync::OnceCell;
+use crate::storage::{DatabaseBackend, DatabaseManager, DatabaseOperation};
 use redb::{Database, ReadableTable, TableDefinition};
 use serde::{Deserialize, Serialize};
 use std::path::{Path, PathBuf};
@@ -19,69 +19,65 @@ const PACKAGES_BY_ID: TableDefinition<&str, &str> = TableDefinition::new("packag
 const PACKAGES_BY_STUDY_UID: TableDefinition<&str, &str> =
     TableDefinition::new("packages_by_study_uid");
 
-/// Global shared database instances mapped by database path
-/// This allows multiple storage locations in tests while sharing instances per path
-static GLOBAL_JMIX_DB: OnceCell<std::sync::Mutex<std::collections::HashMap<PathBuf, Arc<Database>>>> = OnceCell::new();
-
 /// JMIX package index for fast lookups without filesystem walks
-/// Now uses a shared database instance to prevent concurrent access issues
+/// Now uses the generalized database manager to prevent concurrent access issues
 pub struct JmixIndex {
     db: Arc<Database>,
+    db_path: PathBuf,
 }
 
 impl JmixIndex {
-    /// Open or create the index database using the shared global instance
+    /// Open or create the index database using the generalized database manager
     pub fn open(db_path: &Path) -> Result<Self, String> {
-        let db = get_or_create_shared_database(db_path)?;
-        Ok(Self { db })
+        let db_path_buf = db_path.to_path_buf();
+        let db = DatabaseManager::global().get_or_create_database(&db_path_buf)?;
+        
+        // Initialize tables if this is a new database
+        let instance = Self { db: db.clone(), db_path: db_path_buf };
+        instance.initialize_tables(&db)?;
+        
+        Ok(instance)
     }
 
     /// Create a new JmixIndex with the provided shared database (used for testing)
     #[cfg(test)]
-    pub fn with_shared_db(db: Arc<Database>) -> Self {
-        Self { db }
+    pub fn with_shared_db(db: Arc<Database>, db_path: PathBuf) -> Self {
+        Self { db, db_path }
     }
 
     /// Index a new JMIX package
     pub fn index_package(&self, info: &JmixPackageInfo) -> Result<(), String> {
-        let write_txn = self
-            .db
-            .begin_write()
-            .map_err(|e| format!("Failed to begin write transaction: {}", e))?;
-
-        {
+        let info_clone = info.clone();
+        DatabaseOperation::write(&self.db, |write_txn| {
             // Serialize the package info
-            let json = serde_json::to_string(info)
+            let json = serde_json::to_string(&info_clone)
                 .map_err(|e| format!("Failed to serialize package info: {}", e))?;
 
-            // Store by ID
-            let mut table = write_txn
-                .open_table(PACKAGES_BY_ID)
-                .map_err(|e| format!("Failed to open packages_by_id table: {}", e))?;
-            table
-                .insert(info.id.as_str(), json.as_str())
-                .map_err(|e| format!("Failed to insert package by ID: {}", e))?;
-        }
-
-        {
-            // Store by study UID (for fast queries)
-            let json = serde_json::to_string(info)
-                .map_err(|e| format!("Failed to serialize package info: {}", e))?;
-
-            let mut table = write_txn
-                .open_table(PACKAGES_BY_STUDY_UID)
-                .map_err(|e| format!("Failed to open packages_by_study_uid table: {}", e))?;
+            {
+                // Store by ID
+                let mut table = write_txn
+                    .open_table(PACKAGES_BY_ID)
+                    .map_err(|e| format!("Failed to open packages_by_id table: {}", e))?;
+                table
+                    .insert(info_clone.id.as_str(), json.as_str())
+                    .map_err(|e| format!("Failed to insert package by ID: {}", e))?;
+            }
 
-            // Key format: "study_uid:id" to support multiple packages per study
-            let key = format!("{}:{}", info.study_uid, info.id);
-            table
-                .insert(key.as_str(), json.as_str())
-                .map_err(|e| format!("Failed to insert package by study UID: {}", e))?;
-        }
+            {
+                // Store by study UID (for fast queries)
+                let mut table = write_txn
+                    .open_table(PACKAGES_BY_STUDY_UID)
+                    .map_err(|e| format!("Failed to open packages_by_study_uid table: {}", e))?;
+
+                // Key format: "study_uid:id" to support multiple packages per study
+                let key = format!("{}:{}", info_clone.study_uid, info_clone.id);
+                table
+                    .insert(key.as_str(), json.as_str())
+                    .map_err(|e| format!("Failed to insert package by study UID: {}", e))?;
+            }
 
-        write_txn
-            .commit()
-            .map_err(|e| format!("Failed to commit package index: {}", e))?;
+            Ok(())
+        })?;
 
         tracing::debug!(
             "📇 Indexed JMIX package: id={}, study_uid={}",
@@ -93,57 +89,55 @@ impl JmixIndex {
 
     /// Lookup a package by ID
     pub fn get_by_id(&self, id: &str) -> Result<Option<JmixPackageInfo>, String> {
-        let read_txn = self
-            .db
-            .begin_read()
-            .map_err(|e| format!("Failed to begin read transaction: {}", e))?;
-
-        let table = read_txn
-            .open_table(PACKAGES_BY_ID)
-            .map_err(|e| format!("Failed to open packages_by_id table: {}", e))?;
-
-        match table.get(id) {
-            Ok(Some(value)) => {
-                let json = value.value();
-                let info: JmixPackageInfo = serde_json::from_str(json)
-                    .map_err(|e| format!("Failed to deserialize package info: {}", e))?;
-                Ok(Some(info))
+        let id_owned = id.to_string();
+        DatabaseOperation::read(&self.db, |read_txn| {
+            let table = read_txn
+                .open_table(PACKAGES_BY_ID)
+                .map_err(|e| format!("Failed to open packages_by_id table: {}", e))?;
+
+            match table.get(id_owned.as_str()) {
+                Ok(Some(value)) => {
+                    let json = value.value();
+                    let info: JmixPackageInfo = serde_json::from_str(json)
+                        .map_err(|e| format!("Failed to deserialize package info: {}", e))?;
+                    Ok(Some(info))
+                }
+                Ok(None) => Ok(None),
+                Err(e) => Err(format!("Failed to get package by ID: {}", e)),
             }
-            Ok(None) => Ok(None),
-            Err(e) => Err(format!("Failed to get package by ID: {}", e)),
-        }
+        })
     }
 
     /// Query packages by study UID
     pub fn query_by_study_uid(&self, study_uid: &str) -> Result<Vec<JmixPackageInfo>, String> {
-        let read_txn = self
-            .db
-            .begin_read()
-            .map_err(|e| format!("Failed to begin read transaction: {}", e))?;
-
-        let table = read_txn
-            .open_table(PACKAGES_BY_STUDY_UID)
-            .map_err(|e| format!("Failed to open packages_by_study_uid table: {}", e))?;
-
-        let mut results = Vec::new();
-        let prefix = format!("{}:", study_uid);
-
-        // Iterate over all entries and filter by prefix
-        let iter = table
-            .iter()
-            .map_err(|e| format!("Failed to iterate packages_by_study_uid: {}", e))?;
-
-        for entry in iter {
-            let (key, value) = entry.map_err(|e| format!("Failed to read entry: {}", e))?;
-            let key_str = key.value();
-
-            if key_str.starts_with(&prefix) {
-                let json = value.value();
-                let info: JmixPackageInfo = serde_json::from_str(json)
-                    .map_err(|e| format!("Failed to deserialize package info: {}", e))?;
-                results.push(info);
+        let study_uid_owned = study_uid.to_string();
+        let results = DatabaseOperation::read(&self.db, |read_txn| {
+            let table = read_txn
+                .open_table(PACKAGES_BY_STUDY_UID)
+                .map_err(|e| format!("Failed to open packages_by_study_uid table: {}", e))?;
+
+            let mut results = Vec::new();
+            let prefix = format!("{}:", study_uid_owned);
+
+            // Iterate over all entries and filter by prefix
+            let iter = table
+                .iter()
+                .map_err(|e| format!("Failed to iterate packages_by_study_uid: {}", e))?;
+
+            for entry in iter {
+                let (key, value) = entry.map_err(|e| format!("Failed to read entry: {}", e))?;
+                let key_str = key.value();
+
+                if key_str.starts_with(&prefix) {
+                    let json = value.value();
+                    let info: JmixPackageInfo = serde_json::from_str(json)
+                        .map_err(|e| format!("Failed to deserialize package info: {}", e))?;
+                    results.push(info);
+                }
             }
-        }
+
+            Ok(results)
+        })?;
 
         tracing::debug!(
             "📇 Found {} packages for study_uid={}",
@@ -155,35 +149,33 @@ impl JmixIndex {
 
     /// Remove a package from the index
     pub fn remove_package(&self, id: &str, study_uid: &str) -> Result<(), String> {
-        let write_txn = self
-            .db
-            .begin_write()
-            .map_err(|e| format!("Failed to begin write transaction: {}", e))?;
-
-        {
-            // Remove from ID index
-            let mut table = write_txn
-                .open_table(PACKAGES_BY_ID)
-                .map_err(|e| format!("Failed to open packages_by_id table: {}", e))?;
-            table
-                .remove(id)
-                .map_err(|e| format!("Failed to remove package by ID: {}", e))?;
-        }
+        let id_owned = id.to_string();
+        let study_uid_owned = study_uid.to_string();
+        
+        DatabaseOperation::write(&self.db, |write_txn| {
+            {
+                // Remove from ID index
+                let mut table = write_txn
+                    .open_table(PACKAGES_BY_ID)
+                    .map_err(|e| format!("Failed to open packages_by_id table: {}", e))?;
+                table
+                    .remove(id_owned.as_str())
+                    .map_err(|e| format!("Failed to remove package by ID: {}", e))?;
+            }
 
-        {
-            // Remove from study UID index
-            let mut table = write_txn
-                .open_table(PACKAGES_BY_STUDY_UID)
-                .map_err(|e| format!("Failed to open packages_by_study_uid table: {}", e))?;
-            let key = format!("{}:{}", study_uid, id);
-            table
-                .remove(key.as_str())
-                .map_err(|e| format!("Failed to remove package by study UID: {}", e))?;
-        }
+            {
+                // Remove from study UID index
+                let mut table = write_txn
+                    .open_table(PACKAGES_BY_STUDY_UID)
+                    .map_err(|e| format!("Failed to open packages_by_study_uid table: {}", e))?;
+                let key = format!("{}:{}", study_uid_owned, id_owned);
+                table
+                    .remove(key.as_str())
+                    .map_err(|e| format!("Failed to remove package by study UID: {}", e))?;
+            }
 
-        write_txn
-            .commit()
-            .map_err(|e| format!("Failed to commit package removal: {}", e))?;
+            Ok(())
+        })?;
 
         tracing::debug!("📇 Removed JMIX package from index: id={}", id);
         Ok(())
@@ -195,62 +187,16 @@ impl JmixIndex {
     }
 }
 
-/// Get or create the shared database instance for a specific path
-fn get_or_create_shared_database(db_path: &Path) -> Result<Arc<Database>, String> {
-    let db_path_buf = db_path.to_path_buf();
-    
-    // Get or initialize the global database map
-    let db_map = GLOBAL_JMIX_DB.get_or_init(|| {
-        std::sync::Mutex::new(std::collections::HashMap::new())
-    });
-    
-    // Lock the map and check if we already have a database for this path
-    let mut map = db_map.lock().map_err(|e| format!("Failed to lock database map: {}", e))?;
-    
-    if let Some(existing_db) = map.get(&db_path_buf) {
-        // Return existing database instance for this path
-        tracing::debug!("🔄 Reusing existing database instance for: {}", db_path_buf.display());
-        Ok(existing_db.clone())
-    } else {
-        // Create new database instance for this path
-        tracing::debug!("🆕 Creating new database instance for: {}", db_path_buf.display());
-        let db = init_database(&db_path_buf)?;
-        map.insert(db_path_buf.clone(), db.clone());
-        Ok(db)
-    }
-}
-
-/// Initialize a new database instance
-fn init_database(db_path: &Path) -> Result<Arc<Database>, String> {
-    // Ensure parent directory exists
-    if let Some(parent) = db_path.parent() {
-        std::fs::create_dir_all(parent)
-            .map_err(|e| format!("Failed to create index directory: {}", e))?;
+/// Implement DatabaseBackend trait for JmixIndex
+impl DatabaseBackend for JmixIndex {
+    fn database_path(&self) -> PathBuf {
+        self.db_path.clone()
     }
 
-    tracing::info!("🗄️  Initializing shared JMIX index database: {}", db_path.display());
-    
-    let db = Database::create(db_path)
-        .map_err(|e| format!("Failed to open JMIX index database: {}", e))?;
-
-    // Initialize tables
-    let write_txn = db
-        .begin_write()
-        .map_err(|e| format!("Failed to begin write transaction: {}", e))?;
-    {
-        let _ = write_txn
-            .open_table(PACKAGES_BY_ID)
-            .map_err(|e| format!("Failed to open packages_by_id table: {}", e))?;
-        let _ = write_txn
-            .open_table(PACKAGES_BY_STUDY_UID)
-            .map_err(|e| format!("Failed to open packages_by_study_uid table: {}", e))?;
+    fn initialize_tables(&self, db: &Database) -> Result<(), String> {
+        let table_definitions = &[&PACKAGES_BY_ID, &PACKAGES_BY_STUDY_UID];
+        DatabaseManager::global().initialize_tables(db, table_definitions)
     }
-    write_txn
-        .commit()
-        .map_err(|e| format!("Failed to commit table initialization: {}", e))?;
-
-    tracing::info!("✅ JMIX index database initialized successfully");
-    Ok(Arc::new(db))
 }
 
 /// Get or create the global JMIX index
@@ -263,8 +209,10 @@ pub fn get_jmix_index(store_root: &Path) -> Result<JmixIndex, String> {
 /// Create a new database instance directly (for testing)
 #[cfg(test)]
 pub fn create_test_index(db_path: &Path) -> Result<JmixIndex, String> {
-    let db = init_database(db_path)?;
-    Ok(JmixIndex::with_shared_db(db))
+    let db = DatabaseManager::global().get_or_create_database(db_path)?;
+    let instance = JmixIndex::with_shared_db(db.clone(), db_path.to_path_buf());
+    instance.initialize_tables(&db)?;
+    Ok(instance)
 }
 
 /// Helper to get current Unix timestamp