use super::*; use crate::db::errors::{DataError, MigrationError}; use crate::db::Database; use byteorder::LittleEndian; use memmem::{Searcher, TwoWaySearcher}; use sled::transaction::TransactionError; use sled::{Batch, IVec}; use zerocopy::byteorder::U32; use zerocopy::AsBytes; //Not to be confused with the super::RoomAndUser delineator. #[derive(PartialEq, Eq, std::hash::Hash)] struct RoomAndUser { room_id: String, username: String, } /// Create a version 0 user variable key. fn v0_variable_key(info: &RoomAndUser, variable_name: &str) -> Vec { let mut key = vec![]; key.extend_from_slice(info.room_id.as_bytes()); key.extend_from_slice(info.username.as_bytes()); key.extend_from_slice(variable_name.as_bytes()); key } fn map_value_to_room_and_user( entry: sled::Result<(IVec, IVec)>, ) -> Result { if let Ok((key, _)) = entry { let keys: Vec> = key .split(|&b| b == 0xff) .map(|b| str::from_utf8(b)) .collect(); if let &[_, Ok(room_id), Ok(username), Ok(_variable)] = keys.as_slice() { Ok(RoomAndUser { room_id: room_id.to_owned(), username: username.to_owned(), }) } else { Err(MigrationError::MigrationFailed( "a key violates utf8 schema".to_string(), )) } } else { Err(MigrationError::MigrationFailed( "encountered unexpected key".to_string(), )) } } pub(in crate::db) fn add_room_user_variable_count(db: &Database) -> Result<(), DataError> { let tree = &db.variables.0; let prefix = variables_space_prefix(""); //Extract a vec of tuples, consisting of room id + username. let results: Vec = tree .scan_prefix(prefix) .map(map_value_to_room_and_user) .collect::, MigrationError>>()?; let counts: HashMap = results .into_iter() .fold(HashMap::new(), |mut count_map, room_and_user| { let count = count_map.entry(room_and_user).or_insert(0); *count += 1; count_map }); //Start a transaction on the variables tree. let tx_result: Result<_, TransactionError> = db.variables.0.transaction(|tx_vars| { let batch = counts.iter().fold(Batch::default(), |mut batch, entry| { let (info, count) = entry; //Add variable count according to new schema. let delineator = super::RoomAndUser(&info.room_id, &info.username); let key = variables_space_key(delineator, VARIABLE_COUNT_KEY); let db_value: U32 = U32::new(*count); batch.insert(key, db_value.as_bytes()); //Delete the old variable_count variable if exists. let old_key = v0_variable_key(&info, "variable_count"); batch.remove(old_key); batch }); tx_vars.apply_batch(&batch)?; Ok(()) }); tx_result?; //For some reason, it cannot infer the type Ok(()) } pub(in crate::db) fn delete_v0_schema(db: &Database) -> Result<(), DataError> { let mut vars = db.variables.0.scan_prefix(""); let mut batch = Batch::default(); while let Some(Ok((key, _))) = vars.next() { let key = key.to_vec(); if !key.contains(&0xff) { batch.remove(key); } } db.variables.0.apply_batch(batch)?; Ok(()) } pub(in crate::db) fn delete_variable_count(db: &Database) -> Result<(), DataError> { let prefix = variables_space_prefix(""); let mut vars = db.variables.0.scan_prefix(prefix); let mut batch = Batch::default(); while let Some(Ok((key, _))) = vars.next() { let search = TwoWaySearcher::new(b"variable_count"); let ends_with = { match search.search_in(&key) { Some(index) => key.len() - index == b"variable_count".len(), None => false, } }; if ends_with { batch.remove(key); } } db.variables.0.apply_batch(batch)?; Ok(()) }