gogo2/crypto/sol/modules/storage.py
2024-10-27 22:59:38 +02:00

261 lines
10 KiB
Python

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
DATABASE_FILE = "./app_data.db"
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()
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()
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()
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()
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
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()
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()
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()
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
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)
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
# # # # # # USERS
# For this example, we'll use a simple dictionary to store users
users = {
"db": {"id": 1, "username": "db", "email": "user1@example.com", "password": "db"},
"popov": {"id": 2, "username": "popov", "email": "user2@example.com", "password": "popov"}
}
def get_or_create_user(email, google_id):
user = next((u for u in users.values() if u['email'] == email), None)
if not user:
user_id = max(u['id'] for u in users.values()) + 1
username = email.split('@')[0] # Use the part before @ as username
user = {
'id': user_id,
'username': username,
'email': email,
'google_id': google_id
}
users[username] = user
return user
def authenticate_user(username, password):
"""
Authenticate a user based on username and password.
Returns user data if authentication is successful, None otherwise.
"""
user = users.get(username)
if user and user['password'] == password:
return {"id": user['id'], "username": user['username'], "email": user['email']}
return None
def get_user_by_id(user_id):
"""
Retrieve a user by their ID.
"""
for user in users.values():
if user['id'] == int(user_id):
return {"id": user['id'], "username": user['username'], "email": user['email']}
return None
def store_api_key(user_id, api_key):
"""
Store the generated API key for a user.
"""
# In a real application, you would store this in a database
# For this example, we'll just print it
print(f"Storing API key {api_key} for user {user_id}")