WebashalarForML/scratch_chat

1#!/usr/bin/env python3
2"""Database management CLI for the chat agent application."""
3
4import sys
5import argparse
6from flask import Flask
7from config import config
8from chat_agent.utils.database import DatabaseManager, get_database_info, check_database_connection
9
10
11def create_app(config_name='development'):
12    """Create Flask app with configuration."""
13    app = Flask(__name__)
14    app.config.from_object(config[config_name])
15    return app
16
17
18def main():
19    """Main CLI interface for database management."""
20    parser = argparse.ArgumentParser(description="Database management tool")
21    parser.add_argument(
22        "command",
23        choices=["init", "reset", "info", "stats", "sample", "cleanup", "check"],
24        help="Database command to run"
25    )
26    parser.add_argument(
27        "--config",
28        default="development",
29        choices=["development", "production", "testing"],
30        help="Configuration environment"
31    )
32    
33    args = parser.parse_args()
34    
35    # Create Flask app
36    app = create_app(args.config)
37    
38    with app.app_context():
39        db_manager = DatabaseManager(app)
40        
41        if args.command == "init":
42            print("Initializing database...")
43            db_manager.create_tables()
44            print("Database initialized successfully")
45            
46        elif args.command == "reset":
47            print("Resetting database...")
48            confirm = input("This will delete all data. Are you sure? (y/N): ")
49            if confirm.lower() == 'y':
50                db_manager.reset_database()
51                print("Database reset completed")
52            else:
53                print("Reset cancelled")
54                
55        elif args.command == "info":
56            print("Database Information:")
57            print("-" * 40)
58            info = get_database_info()
59            if 'error' in info:
60                print(f"Error: {info['error']}")
61            else:
62                print(f"Database URL: {info['database_url']}")
63                print(f"Tables: {info['table_count']}")
64                for table in info['tables']:
65                    count = info['table_counts'].get(table, 'Unknown')
66                    print(f"  - {table}: {count} rows")
67                    
68        elif args.command == "stats":
69            print("Database Statistics:")
70            print("-" * 40)
71            stats = db_manager.get_stats()
72            for key, value in stats.items():
73                if isinstance(value, dict):
74                    print(f"{key}:")
75                    for k, v in value.items():
76                        print(f"  - {k}: {v}")
77                else:
78                    print(f"{key}: {value}")
79                    
80        elif args.command == "sample":
81            print("Creating sample data...")
82            result = db_manager.create_sample_data()
83            print("Sample data created successfully")
84            
85        elif args.command == "cleanup":
86            print("Cleaning up old sessions...")
87            count = db_manager.cleanup_old_sessions()
88            print(f"Cleanup completed: {count} sessions removed")
89            
90        elif args.command == "check":
91            print("Checking database connection...")
92            if check_database_connection():
93                print("✓ Database connection successful")
94            else:
95                print("✗ Database connection failed")
96                sys.exit(1)
97
98
99if __name__ == "__main__":
100    main()