learningcircuit--local-deep-research/tests/api_tests/migrate_research_mode.py

#!/usr/bin/env python3
"""
Add research_mode column to existing databases.
This script adds the missing research_mode column to TokenUsage and SearchCall tables.
"""

import os
import sqlite3
import sys
from pathlib import Path

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from loguru import logger

from local_deep_research.config.paths import get_user_data_directory


def add_research_mode_column(db_path: str):
    """Add research_mode column to tables if it doesn't exist."""
    try:
        conn = sqlite3.connect(db_path)
    except Exception:
        logger.exception(f"Error connecting to {db_path}")
        return
    try:
        cursor = conn.cursor()

        # Check if column exists in token_usage
        cursor.execute("PRAGMA table_info(token_usage)")
        columns = [col[1] for col in cursor.fetchall()]

        if "research_mode" not in columns:
            logger.info(
                f"Adding research_mode column to token_usage in {db_path}"
            )
            cursor.execute(
                "ALTER TABLE token_usage ADD COLUMN research_mode VARCHAR(50)"
            )
            conn.commit()
        else:
            logger.info(
                f"research_mode column already exists in token_usage in {db_path}"
            )

        # Check if column exists in search_calls
        cursor.execute("PRAGMA table_info(search_calls)")
        columns = [col[1] for col in cursor.fetchall()]

        if "research_mode" not in columns:
            logger.info(
                f"Adding research_mode column to search_calls in {db_path}"
            )
            cursor.execute(
                "ALTER TABLE search_calls ADD COLUMN research_mode VARCHAR(50)"
            )
            conn.commit()
        else:
            logger.info(
                f"research_mode column already exists in search_calls in {db_path}"
            )

        logger.info(f"Successfully updated {db_path}")
    except Exception:
        logger.exception(f"Error updating {db_path}")
    finally:
        conn.close()


def migrate_all_user_databases():
    """Migrate all user databases to add research_mode column."""
    # Get base data directory
    from local_deep_research.config.paths import get_data_directory

    base_dir = get_data_directory()

    # Look for user directories
    user_dirs = []
    if base_dir.exists():
        for item in base_dir.iterdir():
            if item.is_dir() and item.name.startswith("user_"):
                user_dirs.append(item)

    if not user_dirs:
        logger.warning("No user directories found")
        return

    logger.info(f"Found {len(user_dirs)} user directories to migrate")

    for user_dir in user_dirs:
        db_path = user_dir / f"{user_dir.name}_encrypted.db"
        if db_path.exists():
            logger.info(f"Migrating database for {user_dir.name}")
            add_research_mode_column(str(db_path))
        else:
            logger.warning(f"No database found in {user_dir}")


if __name__ == "__main__":
    # Enable debug logging
    logger.remove()
    # diagnose=False: don't dump frame-local credentials in tracebacks
    # (#4185 / #4384). This script touches user databases and SQLCipher
    # passwords live in frame locals during migration.
    logger.add(sys.stdout, level="INFO", diagnose=False)

    # Allow unencrypted for testing
    os.environ["LDR_BOOTSTRAP_ALLOW_UNENCRYPTED"] = "true"

    # Run migration
    migrate_all_user_databases()

    # Also migrate specific test user if needed
    if len(sys.argv) > 1:
        username = sys.argv[1]
        user_dir = get_user_data_directory(username)
        db_path = user_dir / f"{user_dir.name}_encrypted.db"
        if db_path.exists():
            logger.info(f"Migrating specific user: {username}")
            add_research_mode_column(str(db_path))