diff --git a/rustuna_storages/src/cache.rs b/rustuna_storages/src/cache.rs index 80057bcb..d7950972 100644 --- a/rustuna_storages/src/cache.rs +++ b/rustuna_storages/src/cache.rs @@ -147,7 +147,9 @@ impl rustuna_core::storage::Storage for CachedStorage { self.study_caches.insert(study_id, StudyCache::new()); self.unfinished_trials.insert(study_id, vec![]); self.last_finished_trial_number.insert(study_id, -1); - Ok(self.studies.last().unwrap()) + self.studies + .last() + .ok_or_else(|| Error::new(ErrorKind::StorageError)) } fn delete_study(&mut self, study_id: u32) -> Result<()> { @@ -167,7 +169,9 @@ impl rustuna_core::storage::Storage for CachedStorage { let trials = self.trials.entry(study_id).or_default(); let number = trial.number; trials.insert(number, trial); - let trial_ref = trials.get(&number).unwrap(); + let trial_ref = trials + .get(&number) + .ok_or_else(|| Error::new(ErrorKind::TrialNotFound))?; let study_cache = self.study_caches.entry(study_id).or_default(); let mut trials_vec: Vec<_> = trials.values().cloned().collect(); @@ -615,7 +619,10 @@ mod tests { storage.create_new_study("study2", vec![Direction::Minimize])?; storage.delete_study(study1_id)?; - let err = storage.get_study(study1_id).err().unwrap(); + let err = storage + .get_study(study1_id) + .err() + .expect("Expected StudyNotFound error"); assert!(matches!(err.kind, ErrorKind::StudyNotFound)); let study3 = storage.create_new_study("study3", vec![Direction::Minimize])?; @@ -650,7 +657,10 @@ mod tests { let t1_num = storage.create_new_trial(study)?.number; assert_eq!(t0_num, 0); assert_eq!(t1_num, 1); - let trials = storage.trials.get(&study).unwrap(); + let trials = storage + .trials + .get(&study) + .expect("Trials map should contain study"); assert_eq!(trials.len(), 2); Ok(()) } @@ -762,7 +772,10 @@ mod tests { storage.set_study_attrs(study_id, s_attrs, false)?; let study = storage.get_study(study_id)?; assert_eq!( - study.attrs.get(&AttrKey::User("foo".to_string())).unwrap(), + study + .attrs + .get(&AttrKey::User("foo".to_string())) + .expect("User attr 'foo' should exist"), "bar" ); @@ -774,7 +787,7 @@ mod tests { trial .attrs .get(&AttrKey::System("key".to_string())) - .unwrap(), + .expect("System attr 'key' should exist"), "val" ); Ok(()) @@ -883,7 +896,7 @@ mod tests { let err = storage .set_trial_param(study_id, trial1.number, "x", &int_dist, 1.0) .err() - .unwrap(); + .expect("Expected IncompatibleDistribution error"); assert!(matches!(err.kind, ErrorKind::IncompatibleDistribution)); Ok(()) } diff --git a/rustuna_storages/src/sqlite3.rs b/rustuna_storages/src/sqlite3.rs index c0be1657..c0f20e6f 100644 --- a/rustuna_storages/src/sqlite3.rs +++ b/rustuna_storages/src/sqlite3.rs @@ -27,14 +27,20 @@ impl SQLite3Storage { } pub fn create_database(&self) -> Result<()> { - let conn = self.conn.lock().unwrap(); + let conn = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; conn.execute_batch(SCHEMA_SQL) .map_err(|_e| Error::new(ErrorKind::StorageError))?; Ok(()) } fn validate_study_id(&self, study_id: u32) -> Result<()> { - let guard = self.conn.lock().unwrap(); + let guard = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; let study_exists: Option = guard .query_row( "SELECT study_id FROM studies WHERE study_id = ?", @@ -57,7 +63,10 @@ impl CachedStorageBackend for SQLite3Storage { study_name: &str, directions: Vec, ) -> rustuna_core::Result { - let guard = self.conn.lock().unwrap(); + let guard = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; let existing: Option = guard .query_row( @@ -104,7 +113,10 @@ impl CachedStorageBackend for SQLite3Storage { ) -> rustuna_core::Result { self.validate_study_id(study_id)?; - let guard = self.conn.lock().unwrap(); + let guard = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; guard .execute( "INSERT INTO trials (number, study_id, state, datetime_start, datetime_complete) \ @@ -157,7 +169,10 @@ impl CachedStorageBackend for SQLite3Storage { value: f64, ) -> rustuna_core::Result<()> { // Note: Compatibility between distributions across trials is enforced by CachedStorage. - let guard = self.conn.lock().unwrap(); + let guard = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; let trial_id: Option = guard .query_row( "SELECT trial_id FROM trials WHERE study_id = ? AND number = ?", @@ -187,7 +202,10 @@ impl CachedStorageBackend for SQLite3Storage { trial_number: u32, state_values: rustuna_core::trial::TrialStateValues, ) -> rustuna_core::Result<()> { - let guard = self.conn.lock().unwrap(); + let guard = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; let result: Option<(u32, String)> = guard .query_row( "SELECT trial_id, state FROM trials WHERE study_id = ? AND number = ?", @@ -267,7 +285,10 @@ impl CachedStorageBackend for SQLite3Storage { } fn get_studies(&mut self) -> rustuna_core::Result> { - let guard = self.conn.lock().unwrap(); + let guard = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; let mut studies = Vec::new(); let mut stmt = guard @@ -352,7 +373,10 @@ impl CachedStorageBackend for SQLite3Storage { study_id: u32, trial_number: u32, ) -> rustuna_core::Result { - let guard = self.conn.lock().unwrap(); + let guard = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; // Query to trials table . let trial_row: Option<(u32, String, Option, Option)> = guard @@ -520,7 +544,10 @@ impl CachedStorageBackend for SQLite3Storage { } } - let mut guard = self.conn.lock().unwrap(); + let mut guard = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; let tx = guard .transaction() .map_err(|_e| Error::new(ErrorKind::StorageError))?; @@ -612,7 +639,10 @@ impl CachedStorageBackend for SQLite3Storage { } } - let mut guard = self.conn.lock().unwrap(); + let mut guard = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; let trial_id: Option = guard .query_row( "SELECT trial_id FROM trials WHERE study_id = ? AND number = ?", @@ -704,7 +734,10 @@ impl CachedStorageBackend for SQLite3Storage { included_numbers: &[u32], trial_number_greater_than: i32, ) -> rustuna_core::Result> { - let guard = self.conn.lock().unwrap(); + let guard = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; // Build SQL query with filters let mut sql = String::from( @@ -897,7 +930,10 @@ impl CachedStorageBackend for SQLite3Storage { } fn delete_study(&mut self, study_id: u32) -> Result<()> { - let guard = self.conn.lock().unwrap(); + let guard = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; // Delete trial-related records using subquery guard @@ -973,7 +1009,10 @@ impl CachedStorageBackend for SQLite3Storage { impl OptunaCompatibleStorage for SQLite3Storage { fn get_study_id_trial_number_from_trial_id(&mut self, trial_id: u32) -> Result<(u32, u32)> { - let guard = self.conn.lock().unwrap(); + let guard = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; let result: Option<(u32, u32)> = guard .query_row( "SELECT study_id, number FROM trials WHERE trial_id = ?", @@ -990,7 +1029,10 @@ impl OptunaCompatibleStorage for SQLite3Storage { study_id: u32, trial_number: u32, ) -> Result { - let guard = self.conn.lock().unwrap(); + let guard = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; let trial_id: Option = guard .query_row( "SELECT trial_id FROM trials WHERE study_id = ? AND number = ?", @@ -1011,7 +1053,10 @@ impl OptunaCompatibleStorage for SQLite3Storage { return Ok(()); } - let guard = self.conn.lock().unwrap(); + let guard = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; // TODO(c-bata): Check if Optuna enables PRAGMA foreign_keys and if we can skip this check // Explicitly check trial existence and state since the schema might be created by Optuna @@ -1074,7 +1119,10 @@ impl OptunaCompatibleStorage for SQLite3Storage { datetime_start: Option, datetime_complete: Option, ) -> Result<()> { - let guard = self.conn.lock().unwrap(); + let guard = self + .conn + .lock() + .map_err(|_| Error::new(ErrorKind::StorageError))?; let exists: Option = guard .query_row( "SELECT trial_id FROM trials WHERE trial_id = ?", @@ -1324,7 +1372,7 @@ mod tests { let err = storage .create_new_study("dup", vec![Direction::Minimize]) .err() - .unwrap(); + .expect("Expected DuplicatedStudy error"); assert!(matches!(err.kind, ErrorKind::DuplicatedStudy)); Ok(()) } @@ -1467,7 +1515,7 @@ mod tests { let err = storage .set_study_attrs(study_id, overwrite, true) .err() - .unwrap(); + .expect("Expected AttrOverwriteNotAllowed error"); assert!(matches!(err.kind, ErrorKind::AttrOverwriteNotAllowed)); let study = storage.get_study(study_id)?; @@ -1545,7 +1593,7 @@ mod tests { let err = storage .set_trial_attrs(study_id, trial.number, overwrite, true) .err() - .unwrap(); + .expect("Expected AttrOverwriteNotAllowed error"); assert!(matches!(err.kind, ErrorKind::AttrOverwriteNotAllowed)); let trial = storage.get_trial(study_id, trial.number)?; @@ -1756,9 +1804,9 @@ mod tests { let intermediate_json = trial1_result .attrs .get(&AttrKey::System("intermediate_values".to_string())) - .unwrap(); + .expect("intermediate_values should exist"); let intermediate_entries: Vec = - serde_json::from_str(intermediate_json).unwrap(); + serde_json::from_str(intermediate_json).expect("Failed to parse intermediate values"); assert_eq!(intermediate_entries.len(), 2); assert_eq!(intermediate_entries[0].step, 0); assert_eq!(intermediate_entries[0].value, Some(0.3)); @@ -1778,9 +1826,9 @@ mod tests { let intermediate_json = trial3_result .attrs .get(&AttrKey::System("intermediate_values".to_string())) - .unwrap(); + .expect("intermediate_values should exist"); let intermediate_entries: Vec = - serde_json::from_str(intermediate_json).unwrap(); + serde_json::from_str(intermediate_json).expect("Failed to parse intermediate values"); assert_eq!(intermediate_entries.len(), 4); assert_eq!(intermediate_entries[0].step, 0); assert_eq!(intermediate_entries[0].value, Some(0.1)); @@ -1800,9 +1848,9 @@ mod tests { let intermediate_json = trial4_result .attrs .get(&AttrKey::System("intermediate_values".to_string())) - .unwrap(); + .expect("intermediate_values should exist"); let intermediate_entries: Vec = - serde_json::from_str(intermediate_json).unwrap(); + serde_json::from_str(intermediate_json).expect("Failed to parse intermediate values"); assert_eq!(intermediate_entries.len(), 1); assert_eq!(intermediate_entries[0].step, 0); assert_eq!(intermediate_entries[0].value, None); @@ -1817,9 +1865,9 @@ mod tests { let intermediate_json = trial1_updated .attrs .get(&AttrKey::System("intermediate_values".to_string())) - .unwrap(); + .expect("intermediate_values should exist"); let intermediate_entries: Vec = - serde_json::from_str(intermediate_json).unwrap(); + serde_json::from_str(intermediate_json).expect("Failed to parse intermediate values"); assert_eq!(intermediate_entries.len(), 2); assert_eq!(intermediate_entries[0].step, 0); assert_eq!(intermediate_entries[0].value, Some(0.2)); @@ -1835,7 +1883,7 @@ mod tests { let err = storage .set_trial_intermediate_values(non_existent_trial_id, invalid_values) .err() - .unwrap(); + .expect("Expected TrialNotFound error"); assert!(matches!(err.kind, ErrorKind::TrialNotFound)); Ok(())