1466223b85
- Removed the `RawDataManager` class and integrated its functionality directly into the `RawTradeRepository`, streamlining the management of raw trade data. - Implemented the `cleanup_old_raw_data` method to delete outdated records, preventing table bloat and improving performance. - Added the `get_raw_data_stats` method to retrieve statistics about raw data storage, enhancing data management capabilities. - Updated documentation to reflect the new methods and their usage, ensuring clarity for future developers. These changes improve the maintainability and efficiency of the database operations related to raw trade data.
385 lines
13 KiB
Python
385 lines
13 KiB
Python
"""
|
|
Database Connection Utility for Crypto Trading Bot Platform
|
|
Provides connection pooling, session management, and database utilities
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import logging
|
|
from contextlib import contextmanager
|
|
from typing import Optional, Generator, Any, Dict, List, Union
|
|
from pathlib import Path
|
|
|
|
# Load environment variables from .env file if it exists
|
|
try:
|
|
from dotenv import load_dotenv
|
|
env_file = Path(__file__).parent.parent / '.env'
|
|
if env_file.exists():
|
|
load_dotenv(env_file)
|
|
except ImportError:
|
|
# dotenv not available, proceed without it
|
|
pass
|
|
|
|
from sqlalchemy import create_engine, Engine, text, event
|
|
from sqlalchemy.orm import sessionmaker, Session, scoped_session
|
|
from sqlalchemy.pool import QueuePool
|
|
from sqlalchemy.exc import SQLAlchemyError, OperationalError, DisconnectionError
|
|
from sqlalchemy.engine import make_url
|
|
import time
|
|
from functools import wraps
|
|
from datetime import datetime, timedelta
|
|
|
|
from .models import Base, create_all_tables, drop_all_tables, RawTrade
|
|
|
|
|
|
# Configure logging
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DatabaseConfig:
|
|
"""Database configuration class"""
|
|
|
|
def __init__(self):
|
|
self.database_url = os.getenv(
|
|
'DATABASE_URL',
|
|
'postgresql://dashboard@localhost:5434/dashboard'
|
|
)
|
|
|
|
# Connection pool settings
|
|
self.pool_size = int(os.getenv('DB_POOL_SIZE', '5'))
|
|
self.max_overflow = int(os.getenv('DB_MAX_OVERFLOW', '10'))
|
|
self.pool_pre_ping = os.getenv('DB_POOL_PRE_PING', 'true').lower() == 'true'
|
|
self.pool_recycle = int(os.getenv('DB_POOL_RECYCLE', '3600')) # 1 hour
|
|
|
|
# Connection timeout settings
|
|
self.connect_timeout = int(os.getenv('DB_CONNECT_TIMEOUT', '30'))
|
|
self.statement_timeout = int(os.getenv('DB_STATEMENT_TIMEOUT', '30000')) # 30 seconds in ms
|
|
|
|
# Retry settings
|
|
self.max_retries = int(os.getenv('DB_MAX_RETRIES', '3'))
|
|
self.retry_delay = float(os.getenv('DB_RETRY_DELAY', '1.0'))
|
|
|
|
# SSL settings
|
|
self.ssl_mode = os.getenv('DB_SSL_MODE', 'prefer')
|
|
|
|
logger.info(f"Database configuration initialized for: {self._safe_url()}")
|
|
|
|
def _safe_url(self) -> str:
|
|
"""Return database URL with password masked for logging"""
|
|
url = make_url(self.database_url)
|
|
return str(url.set(password="***"))
|
|
|
|
def get_engine_kwargs(self) -> Dict[str, Any]:
|
|
"""Get SQLAlchemy engine configuration"""
|
|
return {
|
|
'poolclass': QueuePool,
|
|
'pool_size': self.pool_size,
|
|
'max_overflow': self.max_overflow,
|
|
'pool_pre_ping': self.pool_pre_ping,
|
|
'pool_recycle': self.pool_recycle,
|
|
'connect_args': {
|
|
'connect_timeout': self.connect_timeout,
|
|
'options': f'-c statement_timeout={self.statement_timeout}',
|
|
'sslmode': self.ssl_mode,
|
|
},
|
|
'echo': False, # Disable SQL logging to reduce verbosity
|
|
'future': True, # Use SQLAlchemy 2.0 style
|
|
}
|
|
|
|
|
|
class DatabaseManager:
|
|
"""
|
|
Database manager with connection pooling and session management
|
|
"""
|
|
|
|
def __init__(self, config: Optional[DatabaseConfig] = None):
|
|
self.config = config or DatabaseConfig()
|
|
self._engine: Optional[Engine] = None
|
|
self._session_factory: Optional[sessionmaker] = None
|
|
self._scoped_session: Optional[scoped_session] = None
|
|
|
|
def initialize(self) -> None:
|
|
"""Initialize database engine and session factory"""
|
|
try:
|
|
logger.info("Initializing database connection...")
|
|
|
|
# Create engine with retry logic
|
|
self._engine = self._create_engine_with_retry()
|
|
|
|
# Setup session factory
|
|
self._session_factory = sessionmaker(
|
|
bind=self._engine,
|
|
autocommit=False,
|
|
autoflush=False,
|
|
expire_on_commit=False
|
|
)
|
|
|
|
# Setup scoped session for thread safety
|
|
self._scoped_session = scoped_session(self._session_factory)
|
|
|
|
# Add connection event listeners
|
|
self._setup_connection_events()
|
|
|
|
logger.info("Database connection initialized successfully")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize database: {e}")
|
|
raise
|
|
|
|
def _create_engine_with_retry(self) -> Engine:
|
|
"""Create database engine with retry logic"""
|
|
for attempt in range(self.config.max_retries):
|
|
try:
|
|
engine = create_engine(
|
|
self.config.database_url,
|
|
**self.config.get_engine_kwargs()
|
|
)
|
|
|
|
# Test connection
|
|
with engine.connect() as conn:
|
|
conn.execute(text("SELECT 1"))
|
|
logger.info("Database connection test successful")
|
|
|
|
return engine
|
|
|
|
except (OperationalError, DisconnectionError) as e:
|
|
attempt_num = attempt + 1
|
|
if attempt_num == self.config.max_retries:
|
|
logger.error(f"Failed to connect to database after {self.config.max_retries} attempts: {e}")
|
|
raise
|
|
|
|
logger.warning(f"Database connection attempt {attempt_num} failed: {e}. Retrying in {self.config.retry_delay}s...")
|
|
time.sleep(self.config.retry_delay)
|
|
|
|
def _setup_connection_events(self) -> None:
|
|
"""Setup SQLAlchemy connection event listeners"""
|
|
if not self._engine:
|
|
return
|
|
|
|
@event.listens_for(self._engine, "connect")
|
|
def set_sqlite_pragma(dbapi_connection, connection_record):
|
|
"""Set connection-level settings"""
|
|
if 'postgresql' in str(self._engine.url):
|
|
with dbapi_connection.cursor() as cursor:
|
|
# Set timezone to UTC
|
|
cursor.execute("SET timezone TO 'UTC'")
|
|
# Set application name for monitoring
|
|
cursor.execute("SET application_name TO 'crypto_trading_bot'")
|
|
|
|
@event.listens_for(self._engine, "checkout")
|
|
def checkout_listener(dbapi_connection, connection_record, connection_proxy):
|
|
"""Log connection checkout"""
|
|
logger.debug("Database connection checked out from pool")
|
|
|
|
@event.listens_for(self._engine, "checkin")
|
|
def checkin_listener(dbapi_connection, connection_record):
|
|
"""Log connection checkin"""
|
|
logger.debug("Database connection returned to pool")
|
|
|
|
@property
|
|
def engine(self) -> Engine:
|
|
"""Get database engine"""
|
|
if not self._engine:
|
|
raise RuntimeError("Database not initialized. Call initialize() first.")
|
|
return self._engine
|
|
|
|
@property
|
|
def session_factory(self) -> sessionmaker:
|
|
"""Get session factory"""
|
|
if not self._session_factory:
|
|
raise RuntimeError("Database not initialized. Call initialize() first.")
|
|
return self._session_factory
|
|
|
|
def create_session(self) -> Session:
|
|
"""Create a new database session"""
|
|
if not self._session_factory:
|
|
raise RuntimeError("Database not initialized. Call initialize() first.")
|
|
return self._session_factory()
|
|
|
|
@contextmanager
|
|
def get_session(self) -> Generator[Session, None, None]:
|
|
"""
|
|
Context manager for database sessions with automatic cleanup
|
|
|
|
Usage:
|
|
with db_manager.get_session() as session:
|
|
# Use session here
|
|
pass
|
|
"""
|
|
session = self.create_session()
|
|
try:
|
|
yield session
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close()
|
|
|
|
@contextmanager
|
|
def get_scoped_session(self) -> Generator[Session, None, None]:
|
|
"""
|
|
Context manager for scoped sessions (thread-safe)
|
|
|
|
Usage:
|
|
with db_manager.get_scoped_session() as session:
|
|
# Use session here
|
|
pass
|
|
"""
|
|
if not self._scoped_session:
|
|
raise RuntimeError("Database not initialized. Call initialize() first.")
|
|
|
|
session = self._scoped_session()
|
|
try:
|
|
yield session
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
self._scoped_session.remove()
|
|
|
|
def test_connection(self) -> bool:
|
|
"""Test database connection"""
|
|
try:
|
|
with self.get_session() as session:
|
|
session.execute(text("SELECT 1"))
|
|
logger.info("Database connection test successful")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Database connection test failed: {e}")
|
|
return False
|
|
|
|
def get_pool_status(self) -> Dict[str, Any]:
|
|
"""Get connection pool status"""
|
|
if not self._engine or not hasattr(self._engine.pool, 'size'):
|
|
return {"status": "Pool not available"}
|
|
|
|
pool = self._engine.pool
|
|
return {
|
|
"pool_size": pool.size(),
|
|
"checked_in": pool.checkedin(),
|
|
"checked_out": pool.checkedout(),
|
|
"overflow": pool.overflow(),
|
|
}
|
|
|
|
def create_tables(self) -> None:
|
|
"""Create all database tables"""
|
|
try:
|
|
create_all_tables(self.engine)
|
|
logger.info("Database tables created successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to create database tables: {e}")
|
|
raise
|
|
|
|
def drop_tables(self) -> None:
|
|
"""Drop all database tables"""
|
|
try:
|
|
drop_all_tables(self.engine)
|
|
logger.info("Database tables dropped successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to drop database tables: {e}")
|
|
raise
|
|
|
|
def execute_schema_file(self, schema_file_path: str) -> None:
|
|
"""Execute SQL schema file"""
|
|
try:
|
|
with open(schema_file_path, 'r') as file:
|
|
schema_sql = file.read()
|
|
|
|
with self.get_session() as session:
|
|
# Split and execute each statement
|
|
statements = [stmt.strip() for stmt in schema_sql.split(';') if stmt.strip()]
|
|
for statement in statements:
|
|
if statement:
|
|
session.execute(text(statement))
|
|
|
|
logger.info(f"Schema file executed successfully: {schema_file_path}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to execute schema file {schema_file_path}: {e}")
|
|
raise
|
|
|
|
def close(self) -> None:
|
|
"""Close database connections and cleanup"""
|
|
try:
|
|
if self._scoped_session:
|
|
self._scoped_session.remove()
|
|
|
|
if self._engine:
|
|
self._engine.dispose()
|
|
logger.info("Database connections closed")
|
|
except Exception as e:
|
|
logger.error(f"Error closing database connections: {e}")
|
|
|
|
|
|
def retry_on_database_error(max_retries: int = 3, delay: float = 1.0):
|
|
"""
|
|
Decorator to retry database operations on transient errors
|
|
|
|
Args:
|
|
max_retries: Maximum number of retry attempts
|
|
delay: Delay between retries in seconds
|
|
"""
|
|
def decorator(func):
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
for attempt in range(max_retries):
|
|
try:
|
|
return func(*args, **kwargs)
|
|
except (OperationalError, DisconnectionError) as e:
|
|
if attempt == max_retries - 1:
|
|
logger.error(f"Database operation failed after {max_retries} attempts: {e}")
|
|
raise
|
|
|
|
logger.warning(f"Database operation failed (attempt {attempt + 1}): {e}. Retrying in {delay}s...")
|
|
time.sleep(delay)
|
|
return None
|
|
return wrapper
|
|
return decorator
|
|
|
|
|
|
# Global database manager instance
|
|
db_manager = DatabaseManager()
|
|
|
|
|
|
def get_db_manager() -> DatabaseManager:
|
|
"""Get global database manager instance"""
|
|
return db_manager
|
|
|
|
|
|
def init_database(config: Optional[DatabaseConfig] = None) -> DatabaseManager:
|
|
"""
|
|
Initialize global database manager
|
|
|
|
Args:
|
|
config: Optional database configuration
|
|
|
|
Returns:
|
|
DatabaseManager instance
|
|
"""
|
|
global db_manager
|
|
if config:
|
|
db_manager = DatabaseManager(config)
|
|
db_manager.initialize()
|
|
return db_manager
|
|
|
|
|
|
# Convenience functions for common operations
|
|
def get_session() -> Generator[Session, None, None]:
|
|
"""Get database session (convenience function)"""
|
|
return db_manager.get_session()
|
|
|
|
|
|
def get_scoped_session() -> Generator[Session, None, None]:
|
|
"""Get scoped database session (convenience function)"""
|
|
return db_manager.get_scoped_session()
|
|
|
|
|
|
def test_connection() -> bool:
|
|
"""Test database connection (convenience function)"""
|
|
return db_manager.test_connection()
|
|
|
|
|
|
def get_pool_status() -> Dict[str, Any]:
|
|
"""Get connection pool status (convenience function)"""
|
|
return db_manager.get_pool_status() |