AI-Discord-Bot/migrate_to_database.py

#!/usr/bin/env python3
"""
migrate_to_database.py
Migration script to move from JSON files to database system
"""

import os
import sys
import json
from datetime import datetime

# Add src directory to path for imports
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from database import db_manager
from logger import setup_logger

logger = setup_logger("migration")

def migrate_user_profiles():
    """Migrate user_profiles.json to database"""
    profiles_path = os.path.join("src", "user_profiles.json")
    
    if not os.path.exists(profiles_path):
        logger.info("No user_profiles.json found, skipping user profile migration")
        return
    
    try:
        with open(profiles_path, 'r', encoding='utf-8') as f:
            profiles = json.load(f)
        
        migrated_count = 0
        for user_id, profile in profiles.items():
            db_manager.save_user_profile(user_id, profile)
            migrated_count += 1
        
        logger.info(f"Migrated {migrated_count} user profiles")
        
        # Backup original file
        backup_path = f"{profiles_path}.backup.{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        os.rename(profiles_path, backup_path)
        logger.info(f"Backed up original file to {backup_path}")
        
    except Exception as e:
        logger.error(f"Failed to migrate user profiles: {e}")

def migrate_memory_data():
    """Migrate memory.json to database"""
    memory_path = os.path.join("src", "memory.json")
    
    if not os.path.exists(memory_path):
        logger.info("No memory.json found, skipping memory migration")
        return
    
    try:
        with open(memory_path, 'r', encoding='utf-8') as f:
            memory_data = json.load(f)
        
        migrated_conversations = 0
        migrated_user_memories = 0
        
        # Migrate conversation memories
        conversations = memory_data.get("conversations", {})
        for channel_id, memories in conversations.items():
            for memory in memories:
                db_manager.store_conversation_memory(
                    channel_id=channel_id,
                    user_id=memory.get("user_id", "unknown"),
                    content=memory.get("content", ""),
                    context=memory.get("context", ""),
                    importance_score=memory.get("importance_score", 0.5)
                )
                migrated_conversations += 1
        
        # Migrate user memories
        user_memories = memory_data.get("user_memories", {})
        for user_id, memories in user_memories.items():
            for memory in memories:
                db_manager.store_user_memory(
                    user_id=user_id,
                    memory_type=memory.get("type", "general"),
                    content=memory.get("content", ""),
                    importance_score=memory.get("importance_score", 0.5)
                )
                migrated_user_memories += 1
        
        logger.info(f"Migrated {migrated_conversations} conversation memories and {migrated_user_memories} user memories")
        
        # Backup original file
        backup_path = f"{memory_path}.backup.{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        os.rename(memory_path, backup_path)
        logger.info(f"Backed up original file to {backup_path}")
        
    except Exception as e:
        logger.error(f"Failed to migrate memory data: {e}")

def verify_migration():
    """Verify that migration was successful"""
    logger.info("Verifying migration...")
    
    # Test user profile operations
    test_profile = {"name": "test", "display_name": "Test User", "interactions": 5}
    db_manager.save_user_profile("test_user", test_profile)
    retrieved = db_manager.get_user_profile("test_user")
    
    if retrieved and retrieved["interactions"] == 5:
        logger.info("✓ User profile operations working")
    else:
        logger.error("✗ User profile operations failed")
        return False
    
    # Test memory operations (only if enabled)
    if db_manager.is_memory_enabled():
        db_manager.store_conversation_memory("test_channel", "test_user", "test message", "test context", 0.8)
        memories = db_manager.get_conversation_context("test_channel", hours=1)
        
        if memories and len(memories) > 0:
            logger.info("✓ Memory operations working")
        else:
            logger.error("✗ Memory operations failed")
            return False
    else:
        logger.info("- Memory system disabled, skipping memory tests")
    
    # Clean up test data
    try:
        if hasattr(db_manager.backend, 'conn'):  # SQLite backend
            cursor = db_manager.backend.conn.cursor()
            cursor.execute("DELETE FROM user_profiles WHERE user_id = 'test_user'")
            cursor.execute("DELETE FROM conversation_memory WHERE channel_id = 'test_channel'")
            db_manager.backend.conn.commit()
        else:  # JSON backend
            # Test data will be cleaned up naturally
            pass
    except Exception as e:
        logger.warning(f"Failed to clean up test data: {e}")
    
    logger.info("✓ Migration verification completed successfully")
    return True

def main():
    """Main migration function"""
    print("=== Discord Bot Database Migration ===")
    print()
    
    # Initialize database (it auto-initializes when imported)
    logger.info("Initializing database system...")
    # db_manager auto-initializes when imported
    
    print(f"Current configuration:")
    print(f"  Backend: {db_manager.get_backend_type()}")
    print(f"  Memory enabled: {db_manager.is_memory_enabled()}")
    print()
    
    # Run migrations
    logger.info("Starting migration process...")
    migrate_user_profiles()
    migrate_memory_data()
    
    # Verify
    if verify_migration():
        print()
        print("✓ Migration completed successfully!")
        print()
        print("Next steps:")
        print("1. Update your bot code to use the new system")
        print("2. Test the bot to ensure everything works")
        print("3. Your original JSON files have been backed up")
        print()
        print("Configuration file: src/settings.yml")
        print("You can switch between SQLite and JSON backends in the database section.")
    else:
        print()
        print("✗ Migration verification failed!")
        print("Please check the logs and try again.")
        return 1
    
    return 0

if __name__ == "__main__":
    exit(main())