File size: 5,936 Bytes
330b6e4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
177
178
179
180
181
182
183
184
#!/usr/bin/env python3
"""Database initialization script for the chat agent application."""

import os
import sys
import argparse
from pathlib import Path

# Add the parent directory to the path so we can import modules
sys.path.append(str(Path(__file__).parent.parent))

from config import config
from migrations.migrate import DatabaseMigrator


def init_database(config_name='development', database_url=None, force=False):
    """Initialize the database with schema and initial data."""
    print(f"Initializing database for environment: {config_name}")
    
    # Create migrator
    migrator = DatabaseMigrator(
        database_url=database_url,
        config_name=config_name
    )
    
    try:
        # Run migrations
        print("Running database migrations...")
        migrator.migrate()
        
        print("Database initialization completed successfully!")
        
        # Show migration status
        print("\nMigration Status:")
        migrator.status()
        
    except Exception as e:
        print(f"Error initializing database: {e}")
        sys.exit(1)


def reset_database(config_name='development', database_url=None):
    """Reset the database by dropping and recreating all tables."""
    print(f"WARNING: This will destroy all data in the {config_name} database!")
    
    if config_name == 'production':
        print("ERROR: Cannot reset production database for safety reasons.")
        sys.exit(1)
    
    confirm = input("Are you sure you want to continue? (yes/no): ")
    if confirm.lower() != 'yes':
        print("Database reset cancelled.")
        return
    
    try:
        from app import create_app
        from chat_agent.models.base import db
        
        # Create app with specified config
        app = create_app(config_name)
        
        with app.app_context():
            print("Dropping all tables...")
            db.drop_all()
            
            print("Creating all tables...")
            db.create_all()
            
            print("Database reset completed!")
            
    except Exception as e:
        print(f"Error resetting database: {e}")
        sys.exit(1)


def seed_database(config_name='development'):
    """Seed the database with initial test data."""
    print(f"Seeding database for environment: {config_name}")
    
    try:
        from app import create_app
        from chat_agent.models.base import db
        from chat_agent.models.chat_session import ChatSession
        from chat_agent.models.message import Message
        from chat_agent.models.language_context import LanguageContext
        import uuid
        from datetime import datetime, timedelta
        
        app = create_app(config_name)
        
        with app.app_context():
            # Create sample chat session
            session_id = uuid.uuid4()
            user_id = uuid.uuid4()
            
            session = ChatSession(
                id=session_id,
                user_id=user_id,
                language='python',
                message_count=2,
                is_active=True
            )
            db.session.add(session)
            
            # Create sample messages
            user_message = Message(
                session_id=session_id,
                role='user',
                content='Hello! Can you help me understand Python functions?',
                language='python'
            )
            db.session.add(user_message)
            
            assistant_message = Message(
                session_id=session_id,
                role='assistant',
                content='Hello! I\'d be happy to help you understand Python functions. A function is a reusable block of code that performs a specific task...',
                language='python'
            )
            db.session.add(assistant_message)
            
            # Create language context
            context = LanguageContext(
                session_id=session_id,
                language='python',
                prompt_template='You are a helpful Python programming assistant.',
                syntax_highlighting='python'
            )
            db.session.add(context)
            
            db.session.commit()
            
            print("Database seeded with sample data!")
            print(f"Sample session ID: {session_id}")
            print(f"Sample user ID: {user_id}")
            
    except Exception as e:
        print(f"Error seeding database: {e}")
        sys.exit(1)


def main():
    """Main CLI interface for database initialization."""
    parser = argparse.ArgumentParser(description="Database initialization tool")
    parser.add_argument(
        "command",
        choices=["init", "reset", "seed", "status"],
        help="Database command to run"
    )
    parser.add_argument(
        "--config",
        default="development",
        choices=["development", "production", "testing"],
        help="Configuration environment"
    )
    parser.add_argument(
        "--database-url",
        help="Database URL (overrides config)"
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Force operation without confirmation"
    )
    
    args = parser.parse_args()
    
    # Run command
    if args.command == "init":
        init_database(args.config, args.database_url, args.force)
    elif args.command == "reset":
        reset_database(args.config, args.database_url)
    elif args.command == "seed":
        seed_database(args.config)
    elif args.command == "status":
        migrator = DatabaseMigrator(
            database_url=args.database_url,
            config_name=args.config
        )
        migrator.status()


if __name__ == "__main__":
    main()