diff --git a/web_app/api/user.py b/web_app/api/user.py index e7f7cfeb..77cb4b20 100644 --- a/web_app/api/user.py +++ b/web_app/api/user.py @@ -309,27 +309,20 @@ async def allow_notification( response_description="Status of withdrawal operations", ) async def withdraw_all(wallet_id: str) -> dict: - """ - Withdraws all supported tokens from the user's contract. - - :param wallet_id: The wallet ID of the user. - :return: detail: "Successfully initiated withdrawals for all tokens" - """ - # Get user's contract address - contract_address = await get_user_contract(wallet_id) - if not contract_address: + """ + Withdraws all supported tokens from the user's contract. + + :param wallet_id: The wallet ID of the user. + :return: detail: "Successfully initiated withdrawals for all tokens" + """ + # Get user's contract address + contract_address = user_db.get_contract_address_by_wallet_id(wallet_id) + if not contract_address: raise HTTPException(status_code=404, detail="Contract not found") - try: - # Perform withdrawals - results = await CLIENT.withdraw_all(contract_address) - return { + # Perform withdrawals + results = await CLIENT.withdraw_all(contract_address) + return { "detail": "Successfully initiated withdrawals for all tokens", "results": results - } - except Exception as e: - logger.error(f"Contract withdrawal failed: {str(e)}") - raise HTTPException( - status_code=500, - detail=f"Failed to withdraw tokens: {str(e)}" - ) + } diff --git a/web_app/contract_tools/blockchain_call.py b/web_app/contract_tools/blockchain_call.py index 28a5d650..c7f6eb83 100644 --- a/web_app/contract_tools/blockchain_call.py +++ b/web_app/contract_tools/blockchain_call.py @@ -362,10 +362,18 @@ async def withdraw_all(self, contract_address: str) -> dict[str, str]: contract_addr_int = self._convert_address(contract_address) results = {} - for token in TokenParams.tokens(): + for token in TokenParams.tokens(): + token_symbol = token.name + try: token_addr_int = self._convert_address(token.address) - token_symbol = token.name + + except ValueError as e: + logger.error(f"Invalid address format for {token_symbol}: {str(e)}") + results[token_symbol] = f"Failed: Invalid address format" + continue + + try: logger.info(f"Withdrawing {token_symbol} from contract {contract_address}") await self._func_call( addr=contract_addr_int, @@ -373,11 +381,6 @@ async def withdraw_all(self, contract_address: str) -> dict[str, str]: calldata=[token_addr_int, 0] ) results[token_symbol] = "Success" - - except ValueError as e: - logger.error(f"Invalid address format for {token_symbol}: {str(e)}") - results[token_symbol] = f"Failed: Invalid address format" - except Exception as e: logger.error(f"Error withdrawing {token_symbol}: {e}") results[token_symbol] = f"Failed: {e}"