storage module

.gitignore

@@ -28,3 +28,4 @@ crypto/sol/logs/token_info.json
 crypto/sol/logs/transation_details.json
 .env
 app_data.db
+crypto/sol/.vs/*

crypto/sol/app.py

@@ -1,233 +1,180 @@
-import concurrent.futures
+import asyncio
-import threading
-import queue
-import uvicorn
-import os
-from dotenv import load_dotenv
 import datetime
 import json
-import websockets
 import logging
-from modules.webui import init_app
+import os
-from modules.utils import telegram_utils, logging, get_pk
+from typing import Dict, Any
-from modules.log_processor import watch_for_new_logs
-from modules.SolanaAPI import SAPI
-from config import DO_WATCH_WALLET
-from asgiref.wsgi import WsgiToAsgi
-from multiprocessing import Process
-import time
 
+import uvicorn
+from asgiref.wsgi import WsgiToAsgi
+from dotenv import load_dotenv
+
+from config import DO_WATCH_WALLET
+from modules.SolanaAPI import SAPI
+from modules.log_processor import watch_for_new_logs
+from modules.utils import telegram_utils
+from modules.webui import init_app, teardown_app
+
+# Load environment variables
 load_dotenv()
 load_dotenv('.env.secret')
 
-def save_log(log):
+# Configure logging
-    try:
-        os.makedirs('./logs', exist_ok=True)
-        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
-        filename = f"./logs/log_{timestamp}.json"
-        with open(filename, 'w') as f:
-            json.dump(log, f, indent=2)
-    except Exception as e:
-        logging.error(f"Error saving RPC log: {e}")
-
-PROCESSING_LOG = False
-log_queue = queue.Queue()
-
-def process_log(log_result):
-    global PROCESSING_LOG
-    tr_details = {
-        "order_id": None,
-        "token_in": None,
-        "token_out": None,
-        "amount_in": 0,
-        "amount_out": 0,
-        "amount_in_USD": 0,
-        "amount_out_USD": 0,
-        "percentage_swapped": 0
-    }
-                
-    if log_result['value']['err']:
-        return
-
-    logs = log_result['value']['logs']
-    try:
-        PROCESSING_LOG = True
-        swap_operations = ['Program log: Instruction: Swap', 'Program log: Instruction: Swap2', 
-                         'Program log: Instruction: SwapExactAmountIn', 'Program log: Instruction: SwapV2']
-
-        if any(op in logs for op in swap_operations):
-            save_log(log_result)
-            tx_signature_str = log_result['value']['signature']
-
-            before_source_balance = 0
-            source_token_change = 0
-          
-            i = 0
-            while i < len(logs):
-                log_entry = logs[i]
-
-                if tr_details["order_id"] is None and "order_id" in log_entry:
-                    tr_details["order_id"] = log_entry.split(":")[-1].strip()
-                    tr_details["token_in"] = logs[i + 1].split(":")[-1].strip()
-                    tr_details["token_out"] = logs[i + 2].split(":")[-1].strip()
-
-                if "source_token_change" in log_entry:
-                    parts = log_entry.split(", ")
-                    for part in parts:
-                        if "source_token_change" in part:
-                            tr_details["amount_in"] = float(part.split(":")[-1].strip()) / 10 ** 6
-                        elif "destination_token_change" in part:
-                            tr_details["amount_out"] = float(part.split(":")[-1].strip()) / 10 ** 6
-                
-                if "before_source_balance" in log_entry:
-                    parts = log_entry.split(", ")
-                    for part in parts:
-                        if "before_source_balance" in part:
-                            before_source_balance = float(part.split(":")[-1].strip()) / 10 ** 6
-                if "source_token_change" in log_entry:
-                    parts = log_entry.split(", ")
-                    for part in parts:
-                        if "source_token_change" in part:
-                            source_token_change = float(part.split(":")[-1].strip()) / 10 ** 6
-                i += 1
-
-            try:
-                if tr_details["token_in"] is None or tr_details["token_out"] is None or \
-                   tr_details["amount_in"] == 0 or tr_details["amount_out"] == 0:
-                    logging.warning("Incomplete swap details found in logs. Getting details from transaction")
-                    tr_details = SAPI.get_transaction_details_info(tx_signature_str, logs)
-
-                
-                if before_source_balance > 0 and source_token_change > 0:
-                    tr_details["percentage_swapped"] = (source_token_change / before_source_balance) * 100   
-                if tr_details["percentage_swapped"] > 100:
-                    tr_details["percentage_swapped"] = tr_details["percentage_swapped"] / 1000
-
-                try:
-                    token_in = SAPI.dex.TOKENS_INFO[tr_details["token_in"]]
-                    token_out = SAPI.dex.TOKENS_INFO[tr_details["token_out"]]
-                    
-                    tr_details["symbol_in"] = token_in.get('symbol')
-                    tr_details["symbol_out"] = token_out.get('symbol')
-                    tr_details['amount_in_USD'] = tr_details['amount_in'] * token_in.get('price', 0)
-                    tr_details['amount_out_USD'] = tr_details['amount_out'] * token_out.get('price', 0)
-                    
-                except Exception as e:
-                    logging.error(f"Error fetching token prices: {e}")
-                
-                message_text = (
-                    f"<b>Swap detected: </b>\n"
-                    f"{tr_details['amount_in_USD']:.2f} worth of {tr_details['symbol_in']} "
-                    f"({tr_details['percentage_swapped']:.2f}% ) swapped for {tr_details['symbol_out']} \n"
-                )
-                telegram_utils.send_telegram_message(message_text)
-                SAPI.follow_move(tr_details)
-                SAPI.save_token_info()
-                
-            except Exception as e:
-                logging.error(f"Error acquiring log details and following: {e}")
-                telegram_utils.send_telegram_message(f"Not followed! Error following move.")
-                
-    except Exception as e:
-        logging.error(f"Error processing log: {e}")
-
-    PROCESSING_LOG = False
-    return tr_details
-
-# def process_messages(websocket):
-#     try:
-#         while True:
-#             response = websocket.recv()
-#             response_data = json.loads(response)
-#             logger.debug(f"Received response: {response_data}")
-
-#             if 'result' in response_data:
-#                 new_sub_id = response_data['result']
-#                 if int(new_sub_id) > 1:
-#                     subscription_id = new_sub_id
-#                     logger.info(f"Subscription successful. New id: {subscription_id}")
-#                 elif new_sub_id:
-#                     logger.info(f"Existing subscription confirmed: {subscription_id}")
-#                 else:
-#                     return None
-#                 return subscription_id
-#             elif 'params' in response_data:
-#                 log = response_data['params']['result']
-#                 logger.debug(f"Received transaction log: {log}")
-#                 log_queue.put(log)
-#             else:
-#                 logger.warning(f"Unexpected response: {response_data}")
-
-#     except Exception as e:
-#         logger.error(f"An error occurred: {e}")
-
-def log_processor_worker():
-    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
-        while True:
-            try:
-                log = log_queue.get()
-                executor.submit(process_log, log)
-            except Exception as e:
-                logger.error(f"Error in log processor worker: {e}")
-            finally:
-                log_queue.task_done()
-
 logger = logging.getLogger(__name__)
-app = init_app()
-asgi_app = WsgiToAsgi(app)
 
-def run_with_retry(task_func, *args, **kwargs):
+class LogProcessor:
-    while True:
+    @staticmethod
+    def save_log(log: Dict[str, Any]) -> None:
+        """Save log to JSON file with timestamp."""
         try:
-            task_func(*args, **kwargs)
+            os.makedirs('./logs', exist_ok=True)
+            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
+            filename = f"./logs/log_{timestamp}.json"
+            with open(filename, 'w') as f:
+                json.dump(log, f, indent=2)
         except Exception as e:
-            error_msg = f"Error in task {task_func.__name__}: {e}"
+            logger.error(f"Error saving RPC log: {e}")
-            logger.error(error_msg)
-            telegram_utils.send_telegram_message(error_msg)
-            time.sleep(5)
 
-def init_bot():
+    @staticmethod
-    # Initialize bot components
+    def extract_transaction_details(logs: list) -> Dict[str, Any]:
-    # pk = get_pk()
+        """Extract transaction details from logs."""
-    telegram_utils.initialize()
+        tr_details = {
-    telegram_utils.send_telegram_message("Solana Agent Started. Connecting to mainnet...")
+            "order_id": None,
-    
+            "token_in": None,
-    # Start monitoring tasks
+            "token_out": None,
-    threading.Thread(target=run_with_retry, args=(watch_for_new_logs,), daemon=True).start()
+            "amount_in": 0,
-    
+            "amount_out": 0,
-    if DO_WATCH_WALLET:
+            "amount_in_USD": 0,
-        threading.Thread(target=run_with_retry, args=(SAPI.wallet_watch_loop,), daemon=True).start()
+            "amount_out_USD": 0,
+            "percentage_swapped": 0
+        }
 
-def run_server():
+        before_source_balance = 0
-    uvicorn.run(
+        source_token_change = 0
+
+        for i, log_entry in enumerate(logs):
+            if tr_details["order_id"] is None and "order_id" in log_entry:
+                tr_details["order_id"] = log_entry.split(":")[-1].strip()
+                tr_details["token_in"] = logs[i + 1].split(":")[-1].strip()
+                tr_details["token_out"] = logs[i + 2].split(":")[-1].strip()
+
+            if "source_token_change" in log_entry:
+                parts = log_entry.split(", ")
+                for part in parts:
+                    if "source_token_change" in part:
+                        tr_details["amount_in"] = float(part.split(":")[-1].strip()) / 10 ** 6
+                    elif "destination_token_change" in part:
+                        tr_details["amount_out"] = float(part.split(":")[-1].strip()) / 10 ** 6
+
+            if "before_source_balance" in log_entry:
+                before_source_balance = float(log_entry.split(":")[-1].strip()) / 10 ** 6
+            if "source_token_change" in log_entry:
+                source_token_change = float(log_entry.split(":")[-1].strip()) / 10 ** 6
+
+        if before_source_balance > 0 and source_token_change > 0:
+            tr_details["percentage_swapped"] = (source_token_change / before_source_balance) * 100
+            if tr_details["percentage_swapped"] > 100:
+                tr_details["percentage_swapped"] /= 1000
+
+        return tr_details
+
+    @staticmethod
+    async def process_log(log_result: Dict[str, Any]) -> Dict[str, Any]:
+        """Process a single log entry."""
+        if log_result['value']['err']:
+            return
+
+        logs = log_result['value']['logs']
+        swap_operations = [
+            'Program log: Instruction: Swap',
+            'Program log: Instruction: Swap2',
+            'Program log: Instruction: SwapExactAmountIn',
+            'Program log: Instruction: SwapV2'
+        ]
+
+        try:
+            if not any(op in logs for op in swap_operations):
+                return
+
+            LogProcessor.save_log(log_result)
+            tx_signature = log_result['value']['signature']
+            tr_details = LogProcessor.extract_transaction_details(logs)
+
+            if not all([tr_details["token_in"], tr_details["token_out"], 
+                       tr_details["amount_in"], tr_details["amount_out"]]):
+                tr_details = await SAPI.get_transaction_details_info(tx_signature, logs)
+
+            # Update token information
+            token_in = SAPI.dex.TOKENS_INFO[tr_details["token_in"]]
+            token_out = SAPI.dex.TOKENS_INFO[tr_details["token_out"]]
+            
+            tr_details.update({
+                "symbol_in": token_in.get('symbol'),
+                "symbol_out": token_out.get('symbol'),
+                "amount_in_USD": tr_details['amount_in'] * token_in.get('price', 0),
+                "amount_out_USD": tr_details['amount_out'] * token_out.get('price', 0)
+            })
+
+            # Send notification
+            message = (
+                f"<b>Swap detected: </b>\n"
+                f"{tr_details['amount_in_USD']:.2f} worth of {tr_details['symbol_in']} "
+                f"({tr_details['percentage_swapped']:.2f}% ) swapped for {tr_details['symbol_out']}"
+            )
+            await telegram_utils.send_telegram_message(message)
+
+            # Follow up actions
+            await SAPI.follow_move(tr_details)
+            await SAPI.save_token_info()
+
+        except Exception as e:
+            logger.error(f"Error processing log: {e}")
+            await telegram_utils.send_telegram_message("Not followed! Error following move.")
+
+        return tr_details
+
+class Bot:
+    @staticmethod
+    async def initialize():
+        """Initialize bot and start monitoring."""
+        await telegram_utils.initialize()
+        await telegram_utils.send_telegram_message("Solana Agent Started. Connecting to mainnet...")
+        
+        asyncio.create_task(watch_for_new_logs())
+        if DO_WATCH_WALLET:
+            asyncio.create_task(SAPI.wallet_watch_loop())
+
+async def start_server():
+    """Run the ASGI server."""
+    config = uvicorn.Config(
         "app:asgi_app",
         host="0.0.0.0",
         port=3001,
         log_level="info",
         reload=True
     )
+    server = uvicorn.Server(config)
+    await server.serve()
 
-def main():
+async def main():
-    # Start log processor worker
+    """Main application entry point."""
-    log_processor_thread = threading.Thread(target=log_processor_worker, daemon=True)
+    # Initialize app and create ASGI wrapper
-    log_processor_thread.start()
+    app = await init_app()
+    global asgi_app
+    asgi_app = WsgiToAsgi(app)
     
     # Initialize bot
-    init_bot()
+    await Bot.initialize()
-    
-    # Start server in a separate process
-    server_process = Process(target=run_server)
-    server_process.start()
     
+    # Start server
     try:
-        # Keep main process running
+        await start_server()
-        while True:
-            time.sleep(1)
     except KeyboardInterrupt:
         logger.info("Shutting down...")
-    finally:
+        await teardown_app()
-        server_process.terminate()
-        server_process.join()
 
 if __name__ == '__main__':
-    main()
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        logger.info("Application terminated by user")

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()
+            try:
-            await self.process_transaction(message)
+                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:
+            for attempt in range(retries):
-                async with session.get(base_url, params=params) as response:
+                try:
-                    if response.status == 200:
+                    async with session.get(base_url, params=params) as response:
-                        data = await response.json()
+                        if response.status == 200:
-                        if address in data and self.DISPLAY_CURRENCY.lower() in data[address]:
+                            data = await response.json()
-                            return address, data[address][self.DISPLAY_CURRENCY.lower()]
+                            if address in data and self.DISPLAY_CURRENCY.lower() in data[address]:
-                    else:
+                                return address, data[address][self.DISPLAY_CURRENCY.lower()]
-                        logging.warning(f"Failed to get price for {address} from CoinGecko. Status: {response.status}")
+                        elif response.status == 429:
-            except Exception as e:
+                            logging.warning(f"Rate limit exceeded for {address}. Retrying...")
-                logging.error(f"Error fetching price for {address} from CoinGecko: {str(e)}")
+                            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):

crypto/sol/modules/log_processor.py

@@ -1,7 +1,7 @@
 import os
 import asyncio
 from pathlib import Path
-from .storage import store_transaction, prisma_client
+from .storage import Storage
 from .SolanaAPI import SolanaAPI
 
 LOG_DIRECTORY = "./logs"
@@ -28,7 +28,7 @@ async def process_log_file(file_path):
     transaction_data = await solana_api.process_wh(data)
 
     # Check if the transaction already exists
-    existing_transaction = await prisma_client.transaction.find_first(
+    existing_transaction = await Storage.get_prisma().transaction.find_first(
         where={'solana_signature': solana_signature}
     )
 
@@ -46,7 +46,8 @@ async def process_log_file(file_path):
             'solana_signature': solana_signature,
             'details': details
         }
-        await store_transaction(transaction_data)
+        storage = Storage()
+        await storage.store_transaction(transaction_data)
 
     # Rename the file to append '_saved'
     new_file_path = file_path.with_name(file_path.stem + "_saved" + file_path.suffix)

crypto/sol/modules/storage.py

@@ -1,8 +1,10 @@
-import sys
-import os
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
 from typing import NamedTuple
+from enum import Enum
+from datetime import datetime
+import json
+from prisma import Prisma
+from typing import List, Optional, Dict
+import asyncio
 
 class Transaction(NamedTuple):
     wallet: str
@@ -14,218 +16,243 @@ class Transaction(NamedTuple):
     amount_out: float
     value_out_usd: float
     tx_signature: str
-    
-from enum import Enum
-from datetime import datetime
-from enum import Enum
-import json
-from prisma import Prisma
 
+class TransactionType(Enum):
+    BUY = "BUY"
+    SELL = "SELL"
 class TransactionStatus(Enum):
     PENDING = "PENDING"
     SENT = "SENT"
     CONFIRMED = "CONFIRMED"
 
-# Initialize the Prisma client
-prisma_client = Prisma()
 
-async def init_db():
+class Storage:
-    await prisma_client.connect()
+    _instance: Optional['Storage'] = None
-
+    _lock = asyncio.Lock()
-async def store_transaction(transaction: Transaction):
+    _initialized = False
-    """
+    prisma = Prisma()  # Class-level Prisma instance
-    Store a transaction record in the database.
-    """
-    await prisma_client.transaction.create(
-        data={
-            'wallet_id': transaction.wallet,
-            'timestamp': datetime.now().isoformat(),
-            'type': transaction.transaction_type,
-            'sell_currency': transaction.symbol_in,
-            'sell_amount': transaction.amount_in,
-            'sell_value': transaction.value_in_usd,
-            'buy_currency': transaction.symbol_out,
-            'buy_amount': transaction.amount_out,
-            'buy_value': transaction.value_out_usd,
-            'solana_signature': transaction.tx_signature,
-            'details': json.dumps({}),
-            'status': TransactionStatus.PENDING.value
-        }
-    )
-
-async def update_holdings(wallet_id, currency, amount_change):
-    holding = await prisma_client.holding.find_first(
-        where={
-            'wallet_id': wallet_id,
-            'currency': currency
-        }
-    )
-    if holding:
-        new_amount = holding.amount + amount_change
-        await prisma_client.holding.update(
-            where={'id': holding.id},
-            data={
-                'amount': new_amount,
-                'last_updated': datetime.now().isoformat()
-            }
-        )
-    else:
-        await prisma_client.holding.create(
-            data={
-                'wallet_id': wallet_id,
-                'currency': currency,
-                'amount': amount_change,
-                'last_updated': datetime.now().isoformat()
-            }
-        )
-
-async def get_wallet_holdings(wallet_id):
-    return await prisma_client.holding.find_many(
-        where={'wallet_id': wallet_id},
-        select={'currency': True, 'amount': True}
-    )
-
-async def get_transaction_history(wallet_id, start_date=None, end_date=None, include_closed=False):
-    filters = {'wallet_id': wallet_id}
-    if not include_closed:
-        filters['closed'] = False
-    if start_date:
-        filters['timestamp'] = {'gte': start_date}
-    if end_date:
-        filters['timestamp'] = {'lte': end_date}
     
-    return await prisma_client.transaction.find_many(
+    STABLECOINS = ['USDC', 'USDT', 'SOL']
-        where=filters,
-        order={'timestamp': 'desc'}
-    )
 
-async def close_transaction(transaction_id):
+    def __new__(cls):
-    await prisma_client.transaction.update(
+        if cls._instance is None:
-        where={'id': transaction_id},
+            cls._instance = super(Storage, cls).__new__(cls)
-        data={'closed': True}
+        return cls._instance
-    )
 
-async def get_open_transactions(wallet_id, currency):
+    async def __ainit__(self):
-    return await prisma_client.transaction.find_many(
+        if self._initialized:
-        where={
+            return
-            'wallet_id': wallet_id,
+        
-            'buy_currency': currency,
+        await self.prisma.connect()
-            'closed': False
+        self._initialized = True
-        },
-        order={'timestamp': 'asc'}
-    )
 
-async def calculate_current_holdings(wallet_id):
+    @classmethod
-    transactions = await prisma_client.transaction.group_by(
+    async def get_instance(cls) -> 'Storage':
-        by=['buy_currency'],
+        if not cls._instance:
-        where={'wallet_id': wallet_id, 'closed': False},
+            async with cls._lock:
-        _sum={'buy_amount': True, 'sell_amount': True}
+                if not cls._instance:
-    )
+                    cls._instance = cls()
-    return [
+                    await cls._instance.__ainit__()
-        {
+        return cls._instance
-            'currency': t.buy_currency,
+    @classmethod
-            'amount': t._sum.buy_amount - (t._sum.sell_amount or 0)
+    def get_prisma(cls):
+        return cls.prisma
+    async def disconnect(self):
+        if self._initialized:
+            await self.prisma.disconnect()
+            self._initialized = False
+            
+    def __init__(self):
+        self.prisma = Prisma()
+        self.users = {
+            "db": {"id": 1, "username": "db", "email": "user1@example.com", "password": "db"},
+            "popov": {"id": 2, "username": "popov", "email": "user2@example.com", "password": "popov"}
         }
-        for t in transactions if t._sum.buy_amount > (t._sum.sell_amount or 0)
-    ]
 
-STABLECOINS = ['USDC', 'USDT', 'SOL']
+    def is_connected(self):
+        return self.prisma.is_connected()
 
-async def is_transaction_closed(wallet_id, transaction_id):
+    async def connect(self):
-    transaction = await prisma_client.transaction.find_unique(
+        await self.prisma.connect()
-        where={'id': transaction_id}
+
-    )
+    async def disconnect(self):
-    if transaction:
+        await self.prisma.disconnect()
-        sold_amount = await prisma_client.transaction.aggregate(
+
-            _sum={'sell_amount': True},
+    async def store_transaction(self, transaction: Transaction):
+        return await self.prisma.transaction.create(
+            data={
+                'wallet_id': transaction.wallet,
+                'timestamp': datetime.now().isoformat(),
+                'type': transaction.transaction_type,
+                'sell_currency': transaction.symbol_in,
+                'sell_amount': transaction.amount_in,
+                'sell_value': transaction.value_in_usd,
+                'buy_currency': transaction.symbol_out,
+                'buy_amount': transaction.amount_out,
+                'buy_value': transaction.value_out_usd,
+                'solana_signature': transaction.tx_signature,
+                'details': json.dumps({}),
+                'status': TransactionStatus.PENDING.value
+            }
+        )
+
+    async def update_holdings(self, wallet_id: str, currency: str, amount_change: float):
+        holding = await self.prisma.holding.find_first(
             where={
                 'wallet_id': wallet_id,
-                'sell_currency': transaction.buy_currency,
+                'currency': currency
-                'timestamp': {'gt': transaction.timestamp}
             }
         )
-        return sold_amount._sum.sell_amount >= transaction.buy_amount
+        if holding:
-    return False
+            new_amount = holding.amount + amount_change
+            return await self.prisma.holding.update(
+                where={'id': holding.id},
+                data={
+                    'amount': new_amount,
+                    'last_updated': datetime.now().isoformat()
+                }
+            )
+        else:
+            return await self.prisma.holding.create(
+                data={
+                    'wallet_id': wallet_id,
+                    'currency': currency,
+                    'amount': amount_change,
+                    'last_updated': datetime.now().isoformat()
+                }
+            )
 
-async def close_completed_transactions(wallet_id):
+    async def get_wallet_holdings(self, wallet_id: str):
-    transactions = await prisma_client.transaction.find_many(
+        return await self.prisma.holding.find_many(
-        where={
+            where={'wallet_id': wallet_id},
-            'wallet_id': wallet_id,
+            select={'currency': True, 'amount': True}
-            'closed': False,
+        )
-            'buy_currency': {'notIn': STABLECOINS}
-        }
-    )
-    for transaction in transactions:
-        if await is_transaction_closed(wallet_id, transaction.id):
-            await close_transaction(transaction.id)
 
-async def get_profit_loss(wallet_id, currency, start_date=None, end_date=None):
+    async def get_transaction_history(self, wallet_id: str, start_date: Optional[str] = None, 
-    filters = {
+                                    end_date: Optional[str] = None, include_closed: bool = False):
-        'wallet_id': wallet_id,
+        filters = {'wallet_id': wallet_id}
-        'OR': [
+        if not include_closed:
-            {'sell_currency': currency},
+            filters['closed'] = False
-            {'buy_currency': currency}
+        if start_date:
+            filters['timestamp'] = {'gte': start_date}
+        if end_date:
+            filters['timestamp'] = {'lte': end_date}
+        
+        return await self.prisma.transaction.find_many(
+            where=filters,
+            order={'timestamp': 'desc'}
+        )
+
+    async def close_transaction(self, transaction_id: int):
+        return await self.prisma.transaction.update(
+            where={'id': transaction_id},
+            data={'closed': True}
+        )
+
+    async def get_open_transactions(self, wallet_id: str, currency: str):
+        return await self.prisma.transaction.find_many(
+            where={
+                'wallet_id': wallet_id,
+                'buy_currency': currency,
+                'closed': False
+            },
+            order={'timestamp': 'asc'}
+        )
+
+    async def calculate_current_holdings(self, wallet_id: str):
+        transactions = await self.prisma.transaction.group_by(
+            by=['buy_currency'],
+            where={'wallet_id': wallet_id, 'closed': False},
+            _sum={'buy_amount': True, 'sell_amount': True}
+        )
+        return [
+            {
+                'currency': t.buy_currency,
+                'amount': t._sum.buy_amount - (t._sum.sell_amount or 0)
+            }
+            for t in transactions if t._sum.buy_amount > (t._sum.sell_amount or 0)
         ]
-    }
-    if start_date:
-        filters['timestamp'] = {'gte': start_date}
-    if end_date:
-        filters['timestamp'] = {'lte': end_date}
-    
-    result = await prisma_client.transaction.aggregate(
-        _sum={
-            'sell_value': True,
-            'buy_value': True
-        },
-        where=filters
-    )
-    return (result._sum.sell_value or 0) - (result._sum.buy_value or 0)
 
-# # # # # # USERS
+    async def is_transaction_closed(self, wallet_id: str, transaction_id: int) -> bool:
+        transaction = await self.prisma.transaction.find_unique(
+            where={'id': transaction_id}
+        )
+        if transaction:
+            sold_amount = await self.prisma.transaction.aggregate(
+                _sum={'sell_amount': True},
+                where={
+                    'wallet_id': wallet_id,
+                    'sell_currency': transaction.buy_currency,
+                    'timestamp': {'gt': transaction.timestamp}
+                }
+            )
+            return sold_amount._sum.sell_amount >= transaction.buy_amount
+        return False
 
-# For this example, we'll use a simple dictionary to store users
+    async def close_completed_transactions(self, wallet_id: str):
-users = {
+        transactions = await self.prisma.transaction.find_many(
-    "db": {"id": 1, "username": "db", "email": "user1@example.com", "password": "db"},
+            where={
-    "popov": {"id": 2, "username": "popov", "email": "user2@example.com", "password": "popov"}
+                'wallet_id': wallet_id,
-}
+                'closed': False,
+                'buy_currency': {'notIn': self.STABLECOINS}
+            }
+        )
+        for transaction in transactions:
+            if await self.is_transaction_closed(wallet_id, transaction.id):
+                await self.close_transaction(transaction.id)
 
-def get_or_create_user(email, google_id):
+    async def get_profit_loss(self, wallet_id: str, currency: str, 
-    user = next((u for u in users.values() if u['email'] == email), None)
+                             start_date: Optional[str] = None, end_date: Optional[str] = None):
-    if not user:
+        filters = {
-        user_id = max(u['id'] for u in users.values()) + 1
+            'wallet_id': wallet_id,
-        username = email.split('@')[0]  # Use the part before @ as username
+            'OR': [
-        user = {
+                {'sell_currency': currency},
-            'id': user_id,
+                {'buy_currency': currency}
-            'username': username,
+            ]
-            'email': email,
-            'google_id': google_id
         }
-        users[username] = user
+        if start_date:
-    return user
+            filters['timestamp'] = {'gte': start_date}
+        if end_date:
+            filters['timestamp'] = {'lte': end_date}
+        
+        result = await self.prisma.transaction.aggregate(
+            _sum={
+                'sell_value': True,
+                'buy_value': True
+            },
+            where=filters
+        )
+        return (result._sum.sell_value or 0) - (result._sum.buy_value or 0)
 
-def authenticate_user(username, password):
+    # User management methods
-    """
+    def get_or_create_user(self, email: str, google_id: str):
-    Authenticate a user based on username and password.
+        user = next((u for u in self.users.values() if u['email'] == email), None)
-    Returns user data if authentication is successful, None otherwise.
+        if not user:
-    """
+            user_id = max(u['id'] for u in self.users.values()) + 1
-    user = users.get(username)
+            username = email.split('@')[0]
-    if user and user['password'] == password:
+            user = {
-        return {"id": user['id'], "username": user['username'], "email": user['email']}
+                'id': user_id,
-    return None
+                'username': username,
+                'email': email,
+                'google_id': google_id
+            }
+            self.users[username] = user
+        return user
 
-def get_user_by_id(user_id):
+    def authenticate_user(self, username: str, password: str):
-    """
+        user = self.users.get(username)
-    Retrieve a user by their ID.
+        if user and user['password'] == password:
-    """
-    for user in users.values():
-        if user['id'] == int(user_id):
             return {"id": user['id'], "username": user['username'], "email": user['email']}
-    return None
+        return None
 
-def store_api_key(user_id, api_key):
+    def get_user_by_id(self, user_id: int):
-    """
+        for user in self.users.values():
-    Store the generated API key for a user.
+            if user['id'] == int(user_id):
-    """
+                return {"id": user['id'], "username": user['username'], "email": user['email']}
-    # In a real application, you would store this in a database
+        return None
-    # For this example, we'll just print it
+
-    print(f"Storing API key {api_key} for user {user_id}")
+    def store_api_key(self, user_id: int, api_key: str):
+        print(f"Storing API key {api_key} for user {user_id}")
+
+
+storage = Storage()

crypto/sol/modules/webui.py

@@ -14,13 +14,14 @@ from config import LIQUIDITY_TOKENS, YOUR_WALLET
 
 from modules import storage, utils, SolanaAPI
 from modules.utils import  async_safe_call, decode_instruction_data
+from modules.storage import Storage
 import os
 import logging
 from datetime import datetime
 
 on_transaction = None
 
-def init_app(tr_handler=None):
+async def init_app(tr_handler=None):
     global on_transaction
     on_transaction = tr_handler
     app = Flask(__name__, template_folder='../templates', static_folder='../static')
@@ -29,6 +30,15 @@ def init_app(tr_handler=None):
     executor = ThreadPoolExecutor(max_workers=10)  # Adjust the number of workers as needed
     login_manager = LoginManager(app)
     login_manager.login_view = 'login'
+    
+    storage = Storage()
+
+    # Ensure database connection
+    async def ensure_storage_connection():
+        if not storage.is_connected():
+            await storage.connect()
+
+    asyncio.run(ensure_storage_connection())
 
     # oauth = OAuth(app)
     # google = oauth.remote_app(
@@ -481,6 +491,9 @@ def init_app(tr_handler=None):
         
     return app 
 
+def teardown_app():
+    # Close the database connection
+    storage.disconnect()
 
 # Function to find the latest log file
 def get_latest_log_file(wh:bool):
@@ -499,4 +512,4 @@ def get_latest_log_file(wh:bool):
         utils.log.error(f"Error fetching latest log file: {e}")
         return None
 
-export = init_app
+export = init_app, teardown_app