storage module

crypto/sol/modules/SolanaAPI.py

@@ -41,6 +41,7 @@ from solders import message
 
 from jupiter_python_sdk.jupiter import Jupiter
 import asyncio
+import contextlib
 import json
 import logging
 import random
@@ -99,6 +100,8 @@ class SolanaWS:
             
     async def connect(self):
         while True:
+            if self.websocket is None or self.websocket.closed:
+                await self.connect()
             try:
                 current_url = random.choice(SOLANA_ENDPOINTS)
                 self.websocket = await websockets.connect(current_url, ping_interval=30, ping_timeout=10)
@@ -262,8 +265,13 @@ class SolanaAPI:
 
     async def process_messages(self, solana_ws):
         while True:
-            message = await solana_ws.message_queue.get()
-            await self.process_transaction(message)
+            try:
+                message = await solana_ws.message_queue.get()
+                await self.process_transaction(message)
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Error processing message: {e}")
 
     
     _first_subscription = True
@@ -1045,27 +1053,31 @@ class SolanaDEX:
         base_url = "https://api.coingecko.com/api/v3/simple/token_price/solana"
         prices = {}
 
-        async def fetch_single_price(session, address):
+        async def fetch_single_price(session, address, retries=3, backoff_factor=0.5):
             params = {
                 "contract_addresses": address,
                 "vs_currencies": self.DISPLAY_CURRENCY.lower()
             }
-            try:
-                async with session.get(base_url, params=params) as response:
-                    if response.status == 200:
-                        data = await response.json()
-                        if address in data and self.DISPLAY_CURRENCY.lower() in data[address]:
-                            return address, data[address][self.DISPLAY_CURRENCY.lower()]
-                    else:
-                        logging.warning(f"Failed to get price for {address} from CoinGecko. Status: {response.status}")
-            except Exception as e:
-                logging.error(f"Error fetching price for {address} from CoinGecko: {str(e)}")
+            for attempt in range(retries):
+                try:
+                    async with session.get(base_url, params=params) as response:
+                        if response.status == 200:
+                            data = await response.json()
+                            if address in data and self.DISPLAY_CURRENCY.lower() in data[address]:
+                                return address, data[address][self.DISPLAY_CURRENCY.lower()]
+                        elif response.status == 429:
+                            logging.warning(f"Rate limit exceeded for {address}. Retrying...")
+                            await asyncio.sleep(backoff_factor * (2 ** attempt))
+                        else:
+                            logging.warning(f"Failed to get price for {address} from CoinGecko. Status: {response.status}")
+                except Exception as e:
+                    logging.error(f"Error fetching price for {address} from CoinGecko: {str(e)}")
             return address, None
 
         async with aiohttp.ClientSession() as session:
             tasks = [fetch_single_price(session, address) for address in token_addresses]
             results = await asyncio.gather(*tasks)
-            
+        
             for address, price in results:
                 if price is not None:
                     prices[address] = price
@@ -1169,7 +1181,11 @@ class SolanaDEX:
 
     async def get_wallet_balances(self, wallet_address, doGetTokenName=True):
         balances = {}
+        if not asyncio.get_event_loop().is_running():
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
         logging.info(f"Getting balances for wallet: {wallet_address}")
+        response = None
         try:
             response = await self.solana_client.get_token_accounts_by_owner_json_parsed(
                 Pubkey.from_string(wallet_address),
@@ -1189,16 +1205,17 @@ class SolanaDEX:
                                 mint = info['mint']
                                 decimals = int(info['tokenAmount']['decimals'])
                                 amount = int(info['tokenAmount']['amount'])
-                                amount = float(amount /10**decimals)
+                                amount = int(amount)
                                 if amount > 1:
+                                    amount = float(amount / 10**decimals)
                                     if mint in self.TOKENS_INFO:
                                         token_name = self.TOKENS_INFO[mint].get('symbol')
                                     elif doGetTokenName:
                                         token_name = await self.get_token_metadata_symbol(mint) or 'N/A'
                                         self.TOKENS_INFO[mint] = {'symbol': token_name}
                                         await asyncio.sleep(2)
-                                        
-                                    self.TOKENS_INFO[mint]['holdedAmount'] = round(amount,decimals)
+
+                                    self.TOKENS_INFO[mint]['holdedAmount'] = round(amount, decimals)
                                     self.TOKENS_INFO[mint]['decimals'] = decimals
                                     balances[mint] = {
                                         'name': token_name or 'N/A',
@@ -1227,7 +1244,10 @@ class SolanaDEX:
 
         except Exception as e:
             logging.error(f"Error getting wallet balances: {str(e)}")
-        logging.info(f"Found {len(response.value)} ({len(balances)} non zero) token accounts for wallet: {wallet_address}")
+        if response and response.value:
+            logging.info(f"Found {len(response.value)} ({len(balances)} non zero) token accounts for wallet: {wallet_address}")
+        else:
+            logging.warning(f"No token accounts found for wallet: {wallet_address}")
         return balances
 
     async def convert_balances_to_currency(self, balances, sol_price):