//! Storage abstraction for CreditService //! //! Provides trait-based storage for wallets, transactions, and reservations. use async_trait::async_trait; use creditservice_types::{Error, Quota, Reservation, ResourceType, Result, Transaction, Wallet}; use std::collections::HashMap; use std::sync::Arc; use tokio::sync::RwLock; /// Storage trait for CreditService data #[async_trait] pub trait CreditStorage: Send + Sync { // Wallet operations async fn get_wallet(&self, project_id: &str) -> Result>; async fn create_wallet(&self, wallet: Wallet) -> Result; async fn update_wallet(&self, wallet: Wallet) -> Result; async fn delete_wallet(&self, project_id: &str) -> Result; // Transaction operations async fn add_transaction(&self, transaction: Transaction) -> Result; async fn get_transactions( &self, project_id: &str, limit: usize, offset: usize, ) -> Result>; // Reservation operations async fn get_reservation(&self, id: &str) -> Result>; async fn create_reservation(&self, reservation: Reservation) -> Result; async fn update_reservation(&self, reservation: Reservation) -> Result; async fn delete_reservation(&self, id: &str) -> Result; async fn get_pending_reservations(&self, project_id: &str) -> Result>; // Quota operations async fn get_quota( &self, project_id: &str, resource_type: ResourceType, ) -> Result>; async fn set_quota(&self, quota: Quota) -> Result; async fn list_quotas(&self, project_id: &str) -> Result>; } /// In-memory storage implementation (for testing and development) #[derive(Debug, Default)] pub struct InMemoryStorage { wallets: RwLock>, transactions: RwLock>>, reservations: RwLock>, quotas: RwLock>, } impl InMemoryStorage { /// Create a new in-memory storage pub fn new() -> Arc { Arc::new(Self::default()) } } #[async_trait] impl CreditStorage for InMemoryStorage { async fn get_wallet(&self, project_id: &str) -> Result> { let wallets = self.wallets.read().await; Ok(wallets.get(project_id).cloned()) } async fn create_wallet(&self, wallet: Wallet) -> Result { let mut wallets = self.wallets.write().await; if wallets.contains_key(&wallet.project_id) { return Err(Error::WalletAlreadyExists(wallet.project_id)); } wallets.insert(wallet.project_id.clone(), wallet.clone()); Ok(wallet) } async fn update_wallet(&self, wallet: Wallet) -> Result { let mut wallets = self.wallets.write().await; if !wallets.contains_key(&wallet.project_id) { return Err(Error::WalletNotFound(wallet.project_id)); } wallets.insert(wallet.project_id.clone(), wallet.clone()); Ok(wallet) } async fn delete_wallet(&self, project_id: &str) -> Result { let mut wallets = self.wallets.write().await; Ok(wallets.remove(project_id).is_some()) } async fn add_transaction(&self, transaction: Transaction) -> Result { let mut transactions = self.transactions.write().await; let project_txns = transactions .entry(transaction.project_id.clone()) .or_insert_with(Vec::new); project_txns.push(transaction.clone()); Ok(transaction) } async fn get_transactions( &self, project_id: &str, limit: usize, offset: usize, ) -> Result> { let transactions = self.transactions.read().await; let project_txns = transactions.get(project_id); match project_txns { Some(txns) => { let result: Vec<_> = txns .iter() .rev() // Most recent first .skip(offset) .take(limit) .cloned() .collect(); Ok(result) } None => Ok(vec![]), } } async fn get_reservation(&self, id: &str) -> Result> { let reservations = self.reservations.read().await; Ok(reservations.get(id).cloned()) } async fn create_reservation(&self, reservation: Reservation) -> Result { let mut reservations = self.reservations.write().await; reservations.insert(reservation.id.clone(), reservation.clone()); Ok(reservation) } async fn update_reservation(&self, reservation: Reservation) -> Result { let mut reservations = self.reservations.write().await; if !reservations.contains_key(&reservation.id) { return Err(Error::ReservationNotFound(reservation.id)); } reservations.insert(reservation.id.clone(), reservation.clone()); Ok(reservation) } async fn delete_reservation(&self, id: &str) -> Result { let mut reservations = self.reservations.write().await; Ok(reservations.remove(id).is_some()) } async fn get_pending_reservations(&self, project_id: &str) -> Result> { let reservations = self.reservations.read().await; let pending: Vec<_> = reservations .values() .filter(|r| { r.project_id == project_id && r.status == creditservice_types::ReservationStatus::Pending }) .cloned() .collect(); Ok(pending) } async fn get_quota( &self, project_id: &str, resource_type: ResourceType, ) -> Result> { let quotas = self.quotas.read().await; Ok(quotas .get(&(project_id.to_string(), resource_type)) .cloned()) } async fn set_quota(&self, quota: Quota) -> Result { let mut quotas = self.quotas.write().await; quotas.insert( (quota.project_id.clone(), quota.resource_type), quota.clone(), ); Ok(quota) } async fn list_quotas(&self, project_id: &str) -> Result> { let quotas = self.quotas.read().await; let result: Vec<_> = quotas .values() .filter(|q| q.project_id == project_id) .cloned() .collect(); Ok(result) } } #[cfg(test)] mod tests { use super::*; #[tokio::test] async fn test_wallet_crud() { let storage = InMemoryStorage::new(); // Create let wallet = Wallet::new("proj-1".into(), "org-1".into(), 10000); let created = storage.create_wallet(wallet.clone()).await.unwrap(); assert_eq!(created.project_id, "proj-1"); // Get let fetched = storage.get_wallet("proj-1").await.unwrap().unwrap(); assert_eq!(fetched.balance, 10000); // Update let mut updated_wallet = fetched.clone(); updated_wallet.balance = 5000; let updated = storage.update_wallet(updated_wallet).await.unwrap(); assert_eq!(updated.balance, 5000); // Delete let deleted = storage.delete_wallet("proj-1").await.unwrap(); assert!(deleted); // Verify deleted let gone = storage.get_wallet("proj-1").await.unwrap(); assert!(gone.is_none()); } }