photoncloud-monorepo/creditservice/crates/creditservice-api/src/storage.rs

//! 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<Option<Wallet>>;
    async fn create_wallet(&self, wallet: Wallet) -> Result<Wallet>;
    async fn update_wallet(&self, wallet: Wallet) -> Result<Wallet>;
    async fn delete_wallet(&self, project_id: &str) -> Result<bool>;

    // Transaction operations
    async fn add_transaction(&self, transaction: Transaction) -> Result<Transaction>;
    async fn get_transactions(
        &self,
        project_id: &str,
        limit: usize,
        offset: usize,
    ) -> Result<Vec<Transaction>>;

    // Reservation operations
    async fn get_reservation(&self, id: &str) -> Result<Option<Reservation>>;
    async fn create_reservation(&self, reservation: Reservation) -> Result<Reservation>;
    async fn update_reservation(&self, reservation: Reservation) -> Result<Reservation>;
    async fn delete_reservation(&self, id: &str) -> Result<bool>;
    async fn get_pending_reservations(&self, project_id: &str) -> Result<Vec<Reservation>>;

    // Quota operations
    async fn get_quota(
        &self,
        project_id: &str,
        resource_type: ResourceType,
    ) -> Result<Option<Quota>>;
    async fn set_quota(&self, quota: Quota) -> Result<Quota>;
    async fn list_quotas(&self, project_id: &str) -> Result<Vec<Quota>>;
}

/// In-memory storage implementation (for testing and development)
#[derive(Debug, Default)]
pub struct InMemoryStorage {
    wallets: RwLock<HashMap<String, Wallet>>,
    transactions: RwLock<HashMap<String, Vec<Transaction>>>,
    reservations: RwLock<HashMap<String, Reservation>>,
    quotas: RwLock<HashMap<(String, ResourceType), Quota>>,
}

impl InMemoryStorage {
    /// Create a new in-memory storage
    pub fn new() -> Arc<Self> {
        Arc::new(Self::default())
    }
}

#[async_trait]
impl CreditStorage for InMemoryStorage {
    async fn get_wallet(&self, project_id: &str) -> Result<Option<Wallet>> {
        let wallets = self.wallets.read().await;
        Ok(wallets.get(project_id).cloned())
    }

    async fn create_wallet(&self, wallet: Wallet) -> Result<Wallet> {
        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<Wallet> {
        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<bool> {
        let mut wallets = self.wallets.write().await;
        Ok(wallets.remove(project_id).is_some())
    }

    async fn add_transaction(&self, transaction: Transaction) -> Result<Transaction> {
        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<Vec<Transaction>> {
        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<Option<Reservation>> {
        let reservations = self.reservations.read().await;
        Ok(reservations.get(id).cloned())
    }

    async fn create_reservation(&self, reservation: Reservation) -> Result<Reservation> {
        let mut reservations = self.reservations.write().await;
        reservations.insert(reservation.id.clone(), reservation.clone());
        Ok(reservation)
    }

    async fn update_reservation(&self, reservation: Reservation) -> Result<Reservation> {
        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<bool> {
        let mut reservations = self.reservations.write().await;
        Ok(reservations.remove(id).is_some())
    }

    async fn get_pending_reservations(&self, project_id: &str) -> Result<Vec<Reservation>> {
        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<Option<Quota>> {
        let quotas = self.quotas.read().await;
        Ok(quotas
            .get(&(project_id.to_string(), resource_type))
            .cloned())
    }

    async fn set_quota(&self, quota: Quota) -> Result<Quota> {
        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<Vec<Quota>> {
        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());
    }
}