manage_db.py

#!/usr/bin/env python3
"""Database management CLI for the chat agent application."""

import sys
import argparse
from flask import Flask
from config import config
from chat_agent.utils.database import DatabaseManager, get_database_info, check_database_connection


def create_app(config_name='development'):
    """Create Flask app with configuration."""
    app = Flask(__name__)
    app.config.from_object(config[config_name])
    return app


def main():
    """Main CLI interface for database management."""
    parser = argparse.ArgumentParser(description="Database management tool")
    parser.add_argument(
        "command",
        choices=["init", "reset", "info", "stats", "sample", "cleanup", "check"],
        help="Database command to run"
    )
    parser.add_argument(
        "--config",
        default="development",
        choices=["development", "production", "testing"],
        help="Configuration environment"
    )
    
    args = parser.parse_args()
    
    # Create Flask app
    app = create_app(args.config)
    
    with app.app_context():
        db_manager = DatabaseManager(app)
        
        if args.command == "init":
            print("Initializing database...")
            db_manager.create_tables()
            print("Database initialized successfully")
            
        elif args.command == "reset":
            print("Resetting database...")
            confirm = input("This will delete all data. Are you sure? (y/N): ")
            if confirm.lower() == 'y':
                db_manager.reset_database()
                print("Database reset completed")
            else:
                print("Reset cancelled")
                
        elif args.command == "info":
            print("Database Information:")
            print("-" * 40)
            info = get_database_info()
            if 'error' in info:
                print(f"Error: {info['error']}")
            else:
                print(f"Database URL: {info['database_url']}")
                print(f"Tables: {info['table_count']}")
                for table in info['tables']:
                    count = info['table_counts'].get(table, 'Unknown')
                    print(f"  - {table}: {count} rows")
                    
        elif args.command == "stats":
            print("Database Statistics:")
            print("-" * 40)
            stats = db_manager.get_stats()
            for key, value in stats.items():
                if isinstance(value, dict):
                    print(f"{key}:")
                    for k, v in value.items():
                        print(f"  - {k}: {v}")
                else:
                    print(f"{key}: {value}")
                    
        elif args.command == "sample":
            print("Creating sample data...")
            result = db_manager.create_sample_data()
            print("Sample data created successfully")
            
        elif args.command == "cleanup":
            print("Cleaning up old sessions...")
            count = db_manager.cleanup_old_sessions()
            print(f"Cleanup completed: {count} sessions removed")
            
        elif args.command == "check":
            print("Checking database connection...")
            if check_database_connection():
                print("✓ Database connection successful")
            else:
                print("✗ Database connection failed")
                sys.exit(1)


if __name__ == "__main__":
    main()