Merge pull request #250 from Jagadeeshftw/main

feat(api/vault): Add endpoints for vault balance retrieval and update

web_app/api/vault.py

@@ -3,41 +3,88 @@
 """
 
 import logging
-from fastapi import APIRouter, HTTPException, Depends
+
+from fastapi import APIRouter, Depends, HTTPException
+
 from web_app.db.crud import DepositDBConnector, UserDBConnector
-from web_app.schemas.vault import VaultDepositRequest, VaultDepositResponse
+from web_app.schemas.vault import (
+    UpdateVaultBalanceRequest,
+    UpdateVaultBalanceResponse,
+    VaultBalanceResponse,
+    VaultDepositRequest,
+    VaultDepositResponse,
+)
 
 logger = logging.getLogger(__name__)
 router = APIRouter(prefix="/api/vault", tags=["vault"])
 
+
 @router.post("/deposit", response_model=VaultDepositResponse)
 async def deposit_to_vault(
     request: VaultDepositRequest,
-    deposit_connector: DepositDBConnector = Depends(DepositDBConnector)
+    deposit_connector: DepositDBConnector = Depends(DepositDBConnector),
 ) -> VaultDepositResponse:
     """
     Process a vault deposit request.
     """
     logger.info(f"Processing deposit request for wallet {request.wallet_id}")
-    
+
     try:
         user_db = UserDBConnector()
         user = user_db.get_user_by_wallet_id(request.wallet_id)
         if not user:
             raise HTTPException(status_code=404, detail="User not found")
 
         vault = deposit_connector.create_vault(
-            user=user,
-            symbol=request.symbol,
-            amount=request.amount
+            user=user, symbol=request.symbol, amount=request.amount
         )
-        
+
         return VaultDepositResponse(
             deposit_id=vault.id,
             wallet_id=request.wallet_id,
             amount=request.amount,
-            symbol=request.symbol
+            symbol=request.symbol,
         )
     except (ValueError, TypeError) as e:
         logger.error(f"Invalid input data: {str(e)}")
         raise HTTPException(status_code=400, detail=str(e))
+
+
+@router.get("/api/balance", response_model=VaultBalanceResponse)
+async def get_user_vault_balance(
+    wallet_id: str,
+    symbol: str,
+    deposit_connector: DepositDBConnector = Depends(DepositDBConnector),
+) -> VaultBalanceResponse:
+    """
+    Get the balance of a user's vault for a specific token.
+    """
+    balance = deposit_connector.get_vault_balance(wallet_id=wallet_id, symbol=symbol)
+    if balance is None:
+        raise HTTPException(
+            status_code=404, detail="Vault not found or user does not exist"
+        )
+    return VaultBalanceResponse(wallet_id=wallet_id, symbol=symbol, amount=balance)
+
+
+@router.post("/api/add_balance", response_model=UpdateVaultBalanceResponse)
+async def add_vault_balance(
+    request: UpdateVaultBalanceRequest,
+    deposit_connector: DepositDBConnector = Depends(DepositDBConnector),
+) -> UpdateVaultBalanceResponse:
+    """
+    Add balance to a user's vault for a specific token.
+    """
+    try:
+        updated_vault = deposit_connector.add_vault_balance(
+            wallet_id=request.wallet_id, symbol=request.symbol, amount=request.amount
+        )
+        return UpdateVaultBalanceResponse(
+            wallet_id=request.wallet_id,
+            symbol=request.symbol,
+            amount=updated_vault.amount,
+        )
+    except (ValueError, TypeError) as e:
+        raise HTTPException(
+            status_code=400, detail=f"Failed to update vault balance: {str(e)}"
+        )

web_app/db/crud.py

@@ -680,23 +680,25 @@ def get_vault(self, wallet_id: str, symbol: str) -> Vault | None:
             vault = db.query(Vault).filter_by(user_id=user.id, symbol=symbol).first()
         return vault
 
-    def add_vault_balance(self, wallet_id: str, symbol: str, amount: str) -> None:
+    def add_vault_balance(self, wallet_id: str, symbol: str, amount: str) -> Vault:
         """
         Adds balance to user vault for token symbol
 
         :param wallet_id: Wallet id of user
         :param symbol: Token symbol or address
         :param amount: An amount in string
 
-        :return: None
+        :return: Updated Vault instance
         """
         vault = self.get_vault(wallet_id, symbol)
         if not vault:
-            return
+            raise ValueError("Vault not found")
         with self.Session() as db:
             new_amount = Decimal(vault.amount) + Decimal(amount)
             db.query(Vault).filter_by(id=vault.id).update(amount=str(new_amount))
             db.commit()
+            vault = self.get_vault(wallet_id, symbol)
+        return vault
 
     def get_vault_balance(self, wallet_id: str, symbol: str) -> str | None:
         """

web_app/schemas/vault.py

@@ -39,3 +39,48 @@ class VaultDepositResponse(BaseModel):
     wallet_id: str
     amount: str
     symbol: str
+
+
+class VaultBalanceResponse(BaseModel):
+    """
+    Schema for the response when retrieving a vault balance.
+
+    Attributes:
+        wallet_id (str): The StarkNet wallet address.
+        symbol (str): Token symbol or address
+        amount (str): The current balance of the vault for the specified token.
+    """
+
+    wallet_id: str
+    symbol: str
+    amount: str
+
+
+class UpdateVaultBalanceRequest(BaseModel):
+    """
+    Schema for the request to update a user's vault balance.
+
+    Attributes:
+        wallet_id (str): The StarkNet wallet address.
+        symbol (str): Token symbol or address
+        amount (str): The amount to be added to the user's vault balance.
+    """
+
+    wallet_id: str
+    symbol: str
+    amount: str
+
+
+class UpdateVaultBalanceResponse(BaseModel):
+    """
+    Schema for the response when updating a vault balance.
+
+    Attributes:
+        wallet_id (str): The StarkNet wallet address.
+        symbol (str): Token symbol or address
+        amount (str): The new balance of the vault for the specified token.
+    """
+
+    wallet_id: str
+    symbol: str
+    amount: str