use crate::{Database, error::DbxResult};
use std::collections::HashMap;
#[derive(Debug, Clone)]
pub enum Operation {
Insert {
table: String,
key: Vec<u8>,
value: Vec<u8>,
},
Delete {
table: String,
key: Vec<u8>,
},
}
#[derive(Debug)]
pub struct TransactionBuffer {
operations: Vec<Operation>,
}
impl TransactionBuffer {
pub fn new() -> Self {
Self {
operations: Vec::new(),
}
}
pub fn insert(&mut self, table: String, key: Vec<u8>, value: Vec<u8>) {
self.operations.push(Operation::Insert { table, key, value });
}
pub fn delete(&mut self, table: String, key: Vec<u8>) {
self.operations.push(Operation::Delete { table, key });
}
pub fn commit(&mut self, db: &Database) -> DbxResult<()> {
let mut insert_batches: HashMap<String, Vec<(Vec<u8>, Vec<u8>)>> = HashMap::new();
let mut delete_ops: Vec<(String, Vec<u8>)> = Vec::new();
for op in self.operations.drain(..) {
match op {
Operation::Insert { table, key, value } => {
insert_batches.entry(table).or_default().push((key, value));
}
Operation::Delete { table, key } => {
delete_ops.push((table, key));
}
}
}
for (table, rows) in insert_batches {
db.insert_batch(&table, rows)?;
}
for (table, key) in delete_ops {
db.delete(&table, &key)?;
}
Ok(())
}
pub fn rollback(&mut self) {
self.operations.clear();
}
pub fn is_empty(&self) -> bool {
self.operations.is_empty()
}
pub fn len(&self) -> usize {
self.operations.len()
}
}
impl Default for TransactionBuffer {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new_buffer_is_empty() {
let buffer = TransactionBuffer::new();
assert!(buffer.is_empty());
assert_eq!(buffer.len(), 0);
}
#[test]
fn test_insert_operation() {
let mut buffer = TransactionBuffer::new();
buffer.insert("users".to_string(), b"key1".to_vec(), b"value1".to_vec());
assert_eq!(buffer.len(), 1);
assert!(!buffer.is_empty());
}
#[test]
fn test_delete_operation() {
let mut buffer = TransactionBuffer::new();
buffer.delete("users".to_string(), b"key1".to_vec());
assert_eq!(buffer.len(), 1);
assert!(!buffer.is_empty());
}
#[test]
fn test_rollback() {
let mut buffer = TransactionBuffer::new();
buffer.insert("users".to_string(), b"key1".to_vec(), b"value1".to_vec());
buffer.delete("users".to_string(), b"key2".to_vec());
assert_eq!(buffer.len(), 2);
buffer.rollback();
assert!(buffer.is_empty());
assert_eq!(buffer.len(), 0);
}
}