File size: 6,255 Bytes
fc115d5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 |
import os
import json
import logging
from abc import ABC, abstractmethod
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional
from dotenv import load_dotenv
import pymongo
# Load environment variables
load_dotenv()
logger = logging.getLogger(__name__)
class StorageInterface(ABC):
"""Abstract base class for storage backends."""
@abstractmethod
def save_state(self, state: Dict):
"""Save the current application state."""
pass
@abstractmethod
def load_state(self) -> Dict:
"""Load the application state."""
pass
@abstractmethod
def log_trade(self, trade: Dict):
"""Append a trade to the trade log."""
pass
@abstractmethod
def get_trades(self, limit: int = 100) -> List[Dict]:
"""Get recent trades."""
pass
class JsonFileStorage(StorageInterface):
"""Legacy storage using local JSON files."""
def __init__(self, base_dir: Path = Path("logs")):
self.base_dir = base_dir
self.state_file = base_dir / "multi_asset_state.json"
self.trade_file = base_dir / "trading_log.json"
self.base_dir.mkdir(parents=True, exist_ok=True)
def save_state(self, state: Dict):
try:
with open(self.state_file, 'w') as f:
json.dump(state, f, indent=2)
except Exception as e:
logger.error(f"Failed to save state to file: {e}")
def load_state(self) -> Dict:
if not self.state_file.exists():
return {}
try:
with open(self.state_file, 'r') as f:
return json.load(f)
except Exception as e:
logger.error(f"Failed to load state from file: {e}")
return {}
def log_trade(self, trade: Dict):
try:
with open(self.trade_file, 'a') as f:
f.write(json.dumps(trade) + '\n')
except Exception as e:
logger.error(f"Failed to log trade to file: {e}")
def get_trades(self, limit: int = 100) -> List[Dict]:
if not self.trade_file.exists():
return []
trades = []
try:
with open(self.trade_file, 'r') as f:
for line in f:
if line.strip():
trades.append(json.loads(line))
return trades[-limit:]
except Exception as e:
logger.error(f"Failed to load trades from file: {e}")
return []
class MongoStorage(StorageInterface):
"""Cloud-native storage using MongoDB Atlas."""
def __init__(self, connection_string: str = None):
self.uri = connection_string or os.getenv("MONGO_URI")
if not self.uri:
raise ValueError("MONGO_URI environment variable is not set")
# Create client with SSL certificate support
try:
import certifi
self.client = pymongo.MongoClient(self.uri, tlsCAFile=certifi.where())
except ImportError:
self.client = pymongo.MongoClient(self.uri)
# Use environment-specific database name to separate prod/dev data
environment = os.getenv("ENVIRONMENT", "production").lower()
db_name = "trading_system" if environment == "production" else f"trading_system_{environment}"
try:
self.db = self.client.get_database(db_name)
self.state_collection = self.db.get_collection("state")
self.trades_collection = self.db.get_collection("trades")
# Verify connection
self.client.admin.command('ping')
logger.info(f"✅ Connected to MongoDB Atlas (database: {db_name}, environment: {environment})")
except Exception as e:
logger.error(f"Failed to connect to MongoDB: {e}")
raise
def save_state(self, state: Dict):
try:
# Upsert the single state document (ID='current_state')
self.state_collection.update_one(
{"_id": "current_state"},
{"$set": state},
upsert=True
)
except Exception as e:
logger.error(f"Failed to save state to MongoDB: {e}")
def load_state(self) -> Dict:
try:
state = self.state_collection.find_one({"_id": "current_state"})
if state:
# Remove _id which is not part of the app state
del state['_id']
return state
return {}
except Exception as e:
logger.error(f"Failed to load state from MongoDB: {e}")
return {}
def log_trade(self, trade: Dict):
try:
self.trades_collection.insert_one(trade)
except Exception as e:
logger.error(f"Failed to log trade to MongoDB: {e}")
def get_trades(self, limit: int = 100) -> List[Dict]:
try:
cursor = self.trades_collection.find().sort("timestamp", -1).limit(limit)
trades = list(cursor)
# Remove _id and reverse to match file order (oldest first) ?
# Actually dashboard expects newest first usually but file reading was messy.
# Let's clean up _id
for t in trades:
if '_id' in t:
del t['_id']
return trades[::-1] # Return oldest first to match file behavior if needed, or check app usage
except Exception as e:
logger.error(f"Failed to fetch trades from MongoDB: {e}")
return []
def get_storage() -> StorageInterface:
"""Factory to get the configured storage backend."""
storage_type = os.getenv("STORAGE_TYPE", "json").lower()
if storage_type == "mongo":
try:
logger.info("💽 Attempting MongoDB Storage...")
return MongoStorage()
except Exception as e:
logger.warning(f"⚠️ MongoDB connection failed: {e}")
logger.info("📁 Falling back to Local JSON File Storage")
return JsonFileStorage()
else:
logger.info("📁 Using Local JSON File Storage")
return JsonFileStorage()
|