sql_rs/storage/

transaction.rs

use crate::{SqlRsError, Result};
use std::collections::HashMap;

#[derive(Debug, Clone)]
pub enum TransactionOp {
    Put { key: Vec<u8>, value: Vec<u8> },
    Delete { key: Vec<u8> },
}

#[derive(Debug, Clone, PartialEq)]
pub enum TransactionState {
    Active,
    Committed,
    Aborted,
}

pub struct Transaction {
    id: u64,
    state: TransactionState,
    operations: Vec<TransactionOp>,
    undo_log: HashMap<Vec<u8>, Option<Vec<u8>>>,
}

impl Transaction {
    pub fn new(id: u64) -> Self {
        Self {
            id,
            state: TransactionState::Active,
            operations: Vec::new(),
            undo_log: HashMap::new(),
        }
    }

    pub fn id(&self) -> u64 {
        self.id
    }

    pub fn state(&self) -> &TransactionState {
        &self.state
    }

    pub fn is_active(&self) -> bool {
        self.state == TransactionState::Active
    }

    pub fn add_operation(&mut self, op: TransactionOp, old_value: Option<Vec<u8>>) -> Result<()> {
        if !self.is_active() {
            return Err(SqlRsError::InvalidOperation(
                "Transaction is not active".to_string(),
            ));
        }

        match &op {
            TransactionOp::Put { key, .. } | TransactionOp::Delete { key } => {
                if !self.undo_log.contains_key(key) {
                    self.undo_log.insert(key.clone(), old_value);
                }
            }
        }

        self.operations.push(op);
        Ok(())
    }

    pub fn commit(&mut self) -> Result<()> {
        if !self.is_active() {
            return Err(SqlRsError::InvalidOperation(
                "Transaction is not active".to_string(),
            ));
        }

        self.state = TransactionState::Committed;
        self.undo_log.clear();
        Ok(())
    }

    pub fn rollback(&mut self) -> Result<Vec<(Vec<u8>, Option<Vec<u8>>)>> {
        if !self.is_active() {
            return Err(SqlRsError::InvalidOperation(
                "Transaction is not active".to_string(),
            ));
        }

        self.state = TransactionState::Aborted;

        let undo_ops: Vec<(Vec<u8>, Option<Vec<u8>>)> = self
            .undo_log
            .iter()
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();

        self.undo_log.clear();
        self.operations.clear();

        Ok(undo_ops)
    }

    pub fn operations(&self) -> &[TransactionOp] {
        &self.operations
    }
}

pub struct TransactionManager {
    next_id: u64,
    active_transactions: HashMap<u64, Transaction>,
}

impl TransactionManager {
    pub fn new() -> Self {
        Self {
            next_id: 1,
            active_transactions: HashMap::new(),
        }
    }

    pub fn begin(&mut self) -> u64 {
        let id = self.next_id;
        self.next_id += 1;

        let tx = Transaction::new(id);
        self.active_transactions.insert(id, tx);

        id
    }

    pub fn get(&self, id: u64) -> Option<&Transaction> {
        self.active_transactions.get(&id)
    }

    pub fn get_mut(&mut self, id: u64) -> Option<&mut Transaction> {
        self.active_transactions.get_mut(&id)
    }

    pub fn commit(&mut self, id: u64) -> Result<()> {
        let tx = self
            .active_transactions
            .get_mut(&id)
            .ok_or_else(|| SqlRsError::NotFound(format!("Transaction {} not found", id)))?;

        tx.commit()?;
        self.active_transactions.remove(&id);

        Ok(())
    }

    pub fn rollback(&mut self, id: u64) -> Result<Vec<(Vec<u8>, Option<Vec<u8>>)>> {
        let tx = self
            .active_transactions
            .get_mut(&id)
            .ok_or_else(|| SqlRsError::NotFound(format!("Transaction {} not found", id)))?;

        let undo_ops = tx.rollback()?;
        self.active_transactions.remove(&id);

        Ok(undo_ops)
    }

    pub fn active_count(&self) -> usize {
        self.active_transactions.len()
    }
}

impl Default for TransactionManager {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_transaction_lifecycle() {
        let mut tx = Transaction::new(1);
        assert!(tx.is_active());

        tx.add_operation(
            TransactionOp::Put {
                key: b"key1".to_vec(),
                value: b"value1".to_vec(),
            },
            None,
        )
        .unwrap();

        tx.commit().unwrap();
        assert_eq!(tx.state(), &TransactionState::Committed);
    }

    #[test]
    fn test_transaction_rollback() {
        let mut tx = Transaction::new(1);

        tx.add_operation(
            TransactionOp::Put {
                key: b"key1".to_vec(),
                value: b"value1".to_vec(),
            },
            Some(b"old_value".to_vec()),
        )
        .unwrap();

        let undo_ops = tx.rollback().unwrap();
        assert_eq!(undo_ops.len(), 1);
        assert_eq!(tx.state(), &TransactionState::Aborted);
    }

    #[test]
    fn test_transaction_manager() {
        let mut mgr = TransactionManager::new();

        let tx_id = mgr.begin();
        assert_eq!(mgr.active_count(), 1);

        mgr.commit(tx_id).unwrap();
        assert_eq!(mgr.active_count(), 0);
    }

    #[test]
    fn test_multiple_transactions() {
        let mut mgr = TransactionManager::new();

        let tx1 = mgr.begin();
        let tx2 = mgr.begin();

        assert_eq!(mgr.active_count(), 2);

        mgr.commit(tx1).unwrap();
        assert_eq!(mgr.active_count(), 1);

        mgr.rollback(tx2).unwrap();
        assert_eq!(mgr.active_count(), 0);
    }
}