TCPDashboard/scripts/init_database.py

#!/usr/bin/env python3
"""
Database Initialization Script
This script initializes the database schema and ensures all tables exist
"""

import os
import sys
from pathlib import Path
from dotenv import load_dotenv
from sqlalchemy import text

# Add the project root to the Python path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

# Load environment variables from .env file
env_file = project_root / '.env'
if env_file.exists():
    load_dotenv(env_file)
    print(f"Loaded environment variables from: {env_file}")
else:
    print("No .env file found, using system environment variables")

from database.connection import init_database, DatabaseConfig
from database.models import create_all_tables
import logging

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def check_environment():
    """Check if environment variables are set"""
    required_vars = ['DATABASE_URL']
    missing_vars = []
    
    for var in required_vars:
        if not os.getenv(var):
            missing_vars.append(var)
    
    if missing_vars:
        logger.warning(f"Missing environment variables: {missing_vars}")
        logger.error("Please create a .env file with the required database configuration")
        logger.info("You can copy env.template to .env and update the values")
        return False
    
    return True


def init_schema():
    """Initialize database schema"""
    try:
        logger.info("Starting database initialization...")
        
        # Check environment
        if not check_environment():
            logger.error("Environment validation failed")
            return False
        
        # Initialize database connection
        config = DatabaseConfig()
        logger.info(f"Connecting to database: {config._safe_url()}")
        
        db_manager = init_database(config)
        
        # Test connection
        if not db_manager.test_connection():
            logger.error("Failed to connect to database")
            return False
        
        logger.info("Database connection successful")
        
        # Create all tables using SQLAlchemy models
        logger.info("Creating database tables...")
        db_manager.create_tables()
        
        # Verify that raw_trades table was created
        with db_manager.get_session() as session:
            result = session.execute(
                text("SELECT COUNT(*) FROM information_schema.tables WHERE table_name = 'raw_trades'")
            ).scalar()
            
            if result > 0:
                logger.info("✅ raw_trades table created successfully")
            else:
                logger.error("❌ raw_trades table was not created")
                return False
        
        # Check all expected tables
        expected_tables = [
            'market_data', 'raw_trades', 'bots', 'signals', 
            'trades', 'bot_performance', 'supported_timeframes', 'supported_exchanges'
        ]
        
        with db_manager.get_session() as session:
            for table in expected_tables:
                result = session.execute(
                    text("SELECT COUNT(*) FROM information_schema.tables WHERE table_name = :table_name"),
                    {"table_name": table}
                ).scalar()
                
                if result > 0:
                    logger.info(f"✅ Table '{table}' exists")
                else:
                    logger.warning(f"⚠️ Table '{table}' not found")
        
        # Get connection pool status
        pool_status = db_manager.get_pool_status()
        logger.info(f"Connection pool status: {pool_status}")
        
        logger.info("🎉 Database initialization completed successfully!")
        return True
        
    except Exception as e:
        logger.error(f"Database initialization failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def apply_schema_file():
    """Apply the clean schema file directly"""
    try:
        logger.info("Applying schema file...")
        
        # Check environment
        if not check_environment():
            logger.error("Environment validation failed")
            return False
        
        # Initialize database connection
        config = DatabaseConfig()
        db_manager = init_database(config)
        
        # Execute schema file
        schema_file = project_root / "database" / "schema_clean.sql"
        if not schema_file.exists():
            logger.error(f"Schema file not found: {schema_file}")
            return False
        
        logger.info(f"Executing schema file: {schema_file}")
        db_manager.execute_schema_file(str(schema_file))
        
        logger.info("✅ Schema file applied successfully")
        return True
        
    except Exception as e:
        logger.error(f"Failed to apply schema file: {e}")
        import traceback
        traceback.print_exc()
        return False


def main():
    """Main function"""
    logger.info("=== Database Initialization ===")
    
    if len(sys.argv) > 1 and sys.argv[1] == "--schema-file":
        # Apply schema file directly
        success = apply_schema_file()
    else:
        # Use SQLAlchemy models
        success = init_schema()
    
    if success:
        logger.info("Database is ready for use!")
        sys.exit(0)
    else:
        logger.error("Database initialization failed!")
        sys.exit(1)


if __name__ == "__main__":
    main()