prisma

crypto/sol/modules/storage.py

@@ -2,213 +2,161 @@ import sys
 import os
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
-import aiosqlite
 import json
 from datetime import datetime
+import prisma
+from prisma import Prisma
 
-DATABASE_FILE = "./app_data.db"
+# Initialize the Prisma client
+prisma_client = Prisma()
 
 async def init_db():
-    async with aiosqlite.connect(DATABASE_FILE) as db:
-        await db.executescript("""
-            CREATE TABLE IF NOT EXISTS users (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                username TEXT UNIQUE NOT NULL,
-                password_hash TEXT NOT NULL,
-                email TEXT UNIQUE NOT NULL,
-                api_key TEXT UNIQUE,
-                plan TEXT DEFAULT 'free'
-            );
-
-            CREATE TABLE IF NOT EXISTS wallets (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                user_id INTEGER,
-                address TEXT NOT NULL,
-                secret TEXT NOT NULL,
-                name TEXT,
-                FOREIGN KEY (user_id) REFERENCES users(id)
-            );
-
-            CREATE TABLE IF NOT EXISTS transactions (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                wallet_id INTEGER,
-                timestamp TEXT,
-                type TEXT,
-                sell_currency TEXT,
-                sell_amount REAL,
-                sell_value REAL,
-                buy_currency TEXT,
-                buy_amount REAL,
-                buy_value REAL,
-                details TEXT,
-                solana_signature TEXT UNIQUE,
-                FOREIGN KEY (wallet_id) REFERENCES wallets(id)
-            );
-
-            CREATE TABLE IF NOT EXISTS holdings (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                wallet_id INTEGER,
-                currency TEXT,
-                amount REAL,
-                last_updated TEXT,
-                is_base BOOLEAN DEFAULT 0,
-                FOREIGN KEY (wallet_id) REFERENCES wallets(id)
-            );
-
-            CREATE TABLE IF NOT EXISTS price_alerts (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                user_id INTEGER,
-                currency TEXT,
-                target_price REAL,
-                alert_type TEXT,
-                FOREIGN KEY (user_id) REFERENCES users(id)
-            );
-
-            CREATE TABLE IF NOT EXISTS followed_accounts (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                user_id INTEGER,
-                address TEXT,
-                followed_address TEXT,
-                name TEXT,
-                FOREIGN KEY (address) REFERENCES wallets(address),
-                FOREIGN KEY (followed_address) REFERENCES wallets(address),
-                FOREIGN KEY (user_id) REFERENCES users(id)
-            );
-        """)
-        await db.commit()
+    await prisma_client.connect()
 
 async def store_transaction(wallet_id, transaction_type, sell_currency, sell_amount, sell_value, buy_currency, buy_amount, buy_value, solana_signature, details=None):
-    async with aiosqlite.connect(DATABASE_FILE) as db:
-        await db.execute("""
-            INSERT INTO transactions (wallet_id, timestamp, type, sell_currency, sell_amount, sell_value, buy_currency, buy_amount, buy_value, solana_signature, details)
-            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
-        """, (wallet_id, datetime.now().isoformat(), transaction_type, sell_currency, sell_amount, sell_value, buy_currency, buy_amount, buy_value, solana_signature, json.dumps(details or {})))
-        await db.commit()
-        
+    await prisma_client.transaction.create(
+        data={
+            'wallet_id': wallet_id,
+            'timestamp': datetime.now().isoformat(),
+            'type': transaction_type,
+            'sell_currency': sell_currency,
+            'sell_amount': sell_amount,
+            'sell_value': sell_value,
+            'buy_currency': buy_currency,
+            'buy_amount': buy_amount,
+            'buy_value': buy_value,
+            'solana_signature': solana_signature,
+            'details': json.dumps(details or {})
+        }
+    )
 
 async def update_holdings(wallet_id, currency, amount_change):
-    async with aiosqlite.connect(DATABASE_FILE) as db:
-        cursor = await db.execute("SELECT amount FROM holdings WHERE wallet_id = ? AND currency = ?", (wallet_id, currency))
-        result = await cursor.fetchone()
-        if result:
-            new_amount = result[0] + amount_change
-            await db.execute("UPDATE holdings SET amount = ?, last_updated = ? WHERE wallet_id = ? AND currency = ?",
-                             (new_amount, datetime.now().isoformat(), wallet_id, currency))
-        else:
-            await db.execute("INSERT INTO holdings (wallet_id, currency, amount, last_updated) VALUES (?, ?, ?, ?)",
-                             (wallet_id, currency, amount_change, datetime.now().isoformat()))
-        await db.commit()
+    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):
-    async with aiosqlite.connect(DATABASE_FILE) as db:
-        cursor = await db.execute("SELECT currency, amount FROM holdings WHERE wallet_id = ?", (wallet_id,))
-        return await cursor.fetchall()
+    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):
-    async with aiosqlite.connect(DATABASE_FILE) as db:
-        query = "SELECT * FROM transactions WHERE wallet_id = ?"
-        params = [wallet_id]
-        if not include_closed:
-            query += " AND closed = 0"
-        if start_date:
-            query += " AND timestamp >= ?"
-            params.append(start_date)
-        if end_date:
-            query += " AND timestamp <= ?"
-            params.append(end_date)
-        query += " ORDER BY timestamp DESC"
-        cursor = await db.execute(query, params)
-        return await cursor.fetchall()
-
-# New utility functions
+    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(
+        where=filters,
+        order={'timestamp': 'desc'}
+    )
 
 async def close_transaction(transaction_id):
-    async with aiosqlite.connect(DATABASE_FILE) as db:
-        await db.execute("UPDATE transactions SET closed = 1 WHERE id = ?", (transaction_id,))
-        await db.commit()
+    await prisma_client.transaction.update(
+        where={'id': transaction_id},
+        data={'closed': True}
+    )
 
 async def get_open_transactions(wallet_id, currency):
-    async with aiosqlite.connect(DATABASE_FILE) as db:
-        cursor = await db.execute("""
-            SELECT * FROM transactions 
-            WHERE wallet_id = ? AND buy_currency = ? AND closed = 0
-            ORDER BY timestamp ASC
-        """, (wallet_id, currency))
-        return await cursor.fetchall()
+    return await prisma_client.transaction.find_many(
+        where={
+            'wallet_id': wallet_id,
+            'buy_currency': currency,
+            'closed': False
+        },
+        order={'timestamp': 'asc'}
+    )
 
 async def calculate_current_holdings(wallet_id):
-    async with aiosqlite.connect(DATABASE_FILE) as db:
-        cursor = await db.execute("""
-            SELECT 
-                buy_currency AS currency,
-                SUM(buy_amount) - COALESCE(
-                    (SELECT SUM(sell_amount) 
-                     FROM transactions t2 
-                     WHERE t2.wallet_id = t1.wallet_id 
-                     AND t2.sell_currency = t1.buy_currency 
-                     AND t2.closed = 0),
-                    0
-                ) AS amount
-            FROM transactions t1
-            WHERE wallet_id = ? AND closed = 0
-            GROUP BY buy_currency
-            HAVING amount > 0
-        """, (wallet_id,))
-        return await cursor.fetchall()
+    transactions = await prisma_client.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)
+    ]
 
 STABLECOINS = ['USDC', 'USDT', 'SOL']
 
 async def is_transaction_closed(wallet_id, transaction_id):
-    async with aiosqlite.connect(DATABASE_FILE) as db:
-        cursor = await db.execute("""
-            SELECT t1.buy_currency, t1.buy_amount,
-                   (SELECT SUM(sell_amount) 
-                    FROM transactions t2 
-                    WHERE t2.wallet_id = t1.wallet_id 
-                    AND t2.sell_currency = t1.buy_currency 
-                    AND t2.timestamp > t1.timestamp) AS sold_amount
-            FROM transactions t1
-            WHERE t1.id = ? AND t1.wallet_id = ?
-        """, (transaction_id, wallet_id))
-        result = await cursor.fetchone()
-        
-        if result:
-            buy_currency, buy_amount, sold_amount = result
-            return sold_amount is not None and sold_amount >= buy_amount
-        return False
+    transaction = await prisma_client.transaction.find_unique(
+        where={'id': transaction_id}
+    )
+    if transaction:
+        sold_amount = await prisma_client.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
 
 async def close_completed_transactions(wallet_id):
-    async with aiosqlite.connect(DATABASE_FILE) as db:
-        cursor = await db.execute("""
-            SELECT id FROM transactions
-            WHERE wallet_id = ? AND closed = 0 AND buy_currency NOT IN (?)
-        """, (wallet_id, ','.join(STABLECOINS)))
-        transactions = await cursor.fetchall()
-        
-        for (transaction_id,) in transactions:
-            if await is_transaction_closed(wallet_id, transaction_id):
-                await close_transaction(transaction_id)
+    transactions = await prisma_client.transaction.find_many(
+        where={
+            'wallet_id': wallet_id,
+            '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 with aiosqlite.connect(DATABASE_FILE) as db:
-        query = """
-            SELECT 
-                SUM(CASE WHEN sell_currency = ? THEN sell_value ELSE -buy_value END) as profit_loss
-            FROM transactions
-            WHERE wallet_id = ? AND (sell_currency = ? OR buy_currency = ?)
-        """
-        params = [currency, wallet_id, currency, currency]
-        
-        if start_date:
-            query += " AND timestamp >= ?"
-            params.append(start_date)
-        if end_date:
-            query += " AND timestamp <= ?"
-            params.append(end_date)
-        
-        cursor = await db.execute(query, params)
-        result = await cursor.fetchone()
-        return result[0] if result else 0
+    filters = {
+        'wallet_id': wallet_id,
+        'OR': [
+            {'sell_currency': currency},
+            {'buy_currency': currency}
+        ]
+    }
+    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