local-deep-research 0.6.0__py3-none-any.whl → 0.6.4__py3-none-any.whl

local_deep_research/web/database/schema_upgrade.py

@@ -1,519 +0,0 @@
-"""
-Schema upgrade script for Local Deep Research database.
-Handles schema upgrades for existing ldr.db databases.
-"""
-
-import os
-import sys
-
-from loguru import logger
-from sqlalchemy import create_engine, inspect, text
-
-# Add the parent directory to sys.path to allow relative imports
-sys.path.append(
-    os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
-)
-
-try:
-    from src.local_deep_research.web.models.database import DB_PATH
-except ImportError:
-    # Fallback path if import fails
-    current_dir = os.path.dirname(os.path.abspath(__file__))
-    project_root = os.path.abspath(
-        os.path.join(current_dir, "..", "..", "..", "..")
-    )
-    DB_PATH = os.path.join(project_root, "src", "data", "ldr.db")
-
-from .models import Base, ResearchStrategy, RateLimitAttempt, RateLimitEstimate
-
-
-def remove_research_log_table(engine):
-    """
-    Remove the redundant research_log table if it exists
-
-    Args:
-        engine: SQLAlchemy engine
-
-    Returns:
-        bool: True if operation was successful, False otherwise
-    """
-    try:
-        inspector = inspect(engine)
-
-        # Check if table exists
-        if inspector.has_table("research_log"):
-            # For SQLite, we need to use raw SQL for DROP TABLE
-            with engine.connect() as conn:
-                conn.execute("DROP TABLE research_log")
-                conn.commit()
-            logger.info("Successfully removed redundant 'research_log' table")
-            return True
-        else:
-            logger.info("Table 'research_log' does not exist, no action needed")
-            return True
-    except Exception:
-        logger.exception("Error removing research_log table")
-        return False
-
-
-def create_research_strategy_table(engine):
-    """
-    Create the research_strategies table if it doesn't exist
-
-    Args:
-        engine: SQLAlchemy engine
-
-    Returns:
-        bool: True if operation was successful, False otherwise
-    """
-    try:
-        inspector = inspect(engine)
-
-        # Check if table exists
-        if not inspector.has_table("research_strategies"):
-            # Create the table using ORM
-            Base.metadata.create_all(
-                engine, tables=[ResearchStrategy.__table__]
-            )
-            logger.info("Successfully created 'research_strategies' table")
-            return True
-        else:
-            logger.info(
-                "Table 'research_strategies' already exists, no action needed"
-            )
-            return True
-    except Exception:
-        logger.exception("Error creating research_strategies table")
-        return False
-
-
-def create_benchmark_tables(engine):
-    """
-    Create benchmark tables if they don't exist
-
-    Args:
-        engine: SQLAlchemy engine
-
-    Returns:
-        bool: True if operation was successful, False otherwise
-    """
-    try:
-        from .benchmark_schema import create_benchmark_tables_simple
-
-        inspector = inspect(engine)
-
-        # Check if benchmark tables already exist
-        if not inspector.has_table("benchmark_runs"):
-            # Create all benchmark tables using simple schema
-            create_benchmark_tables_simple(engine)
-            logger.info("Successfully created benchmark tables")
-            return True
-        else:
-            logger.info("Benchmark tables already exist, no action needed")
-            return True
-    except Exception:
-        logger.exception("Error creating benchmark tables")
-        return False
-
-
-def create_rate_limiting_tables(engine):
-    """
-    Create rate limiting tables if they don't exist
-
-    Args:
-        engine: SQLAlchemy engine
-
-    Returns:
-        bool: True if operation was successful, False otherwise
-    """
-    try:
-        inspector = inspect(engine)
-
-        tables_to_create = []
-
-        # Check if rate_limit_attempts table exists
-        if not inspector.has_table("rate_limit_attempts"):
-            tables_to_create.append(RateLimitAttempt.__table__)
-            logger.info("Need to create 'rate_limit_attempts' table")
-
-        # Check if rate_limit_estimates table exists
-        if not inspector.has_table("rate_limit_estimates"):
-            tables_to_create.append(RateLimitEstimate.__table__)
-            logger.info("Need to create 'rate_limit_estimates' table")
-
-        if tables_to_create:
-            # Create the tables using ORM
-            Base.metadata.create_all(engine, tables=tables_to_create)
-            logger.info(
-                f"Successfully created {len(tables_to_create)} rate limiting tables"
-            )
-        else:
-            logger.info("Rate limiting tables already exist, no action needed")
-
-        return True
-    except Exception:
-        logger.exception("Error creating rate limiting tables")
-        return False
-
-
-def add_research_id_to_benchmark_results(engine):
-    """
-    Add research_id column to benchmark_results table if it doesn't exist.
-    """
-    try:
-        import sqlite3
-
-        # Get database path from engine
-        db_path = engine.url.database
-
-        logger.info("Checking if benchmark_results needs research_id column...")
-
-        conn = sqlite3.connect(db_path)
-
-        try:
-            cursor = conn.cursor()
-
-            # Check if table exists
-            cursor.execute(
-                "SELECT name FROM sqlite_master WHERE type='table' AND name='benchmark_results'"
-            )
-            if not cursor.fetchone():
-                logger.info("benchmark_results table does not exist, skipping")
-                return True
-
-            # Check if research_id column already exists
-            cursor.execute("PRAGMA table_info(benchmark_results)")
-            columns = cursor.fetchall()
-            has_research_id = any(col[1] == "research_id" for col in columns)
-
-            if has_research_id:
-                logger.info("benchmark_results already has research_id column")
-                return True
-
-            # Add research_id column
-            logger.info("Adding research_id column to benchmark_results table")
-            cursor.execute(
-                "ALTER TABLE benchmark_results ADD COLUMN research_id TEXT"
-            )
-
-            conn.commit()
-            logger.info(
-                "Successfully added research_id column to benchmark_results"
-            )
-            return True
-
-        finally:
-            conn.close()
-
-    except Exception:
-        logger.exception("Error adding research_id column to benchmark_results")
-        return False
-
-
-def add_uuid_id_column_to_research_history(engine):
-    """
-    Adds a new `uuid_id` string column to the `research_history` table if it
-    does not exist already.
-    """
-    try:
-        import sqlite3
-
-        # Get database path from engine
-        db_path = engine.url.database
-
-        logger.info("Checking if research_history needs uuid_id column...")
-
-        conn = sqlite3.connect(db_path)
-
-        try:
-            cursor = conn.cursor()
-
-            # Check if table exists
-            cursor.execute(
-                "SELECT name FROM sqlite_master WHERE type='table' AND name='research_history'"
-            )
-            if not cursor.fetchone():
-                logger.info("research_history table does not exist, skipping")
-                return True
-
-            # Check if uuid_id column already exists
-            cursor.execute("PRAGMA table_info(research_history)")
-            columns = cursor.fetchall()
-            has_uuid_id = any(col[1] == "uuid_id" for col in columns)
-
-            if has_uuid_id:
-                logger.info("research_history already has uuid_id column")
-                return True
-
-            # Add uuid_id column
-            logger.info("Adding uuid_id column to research_history table")
-            cursor.execute(
-                "ALTER TABLE research_history ADD COLUMN uuid_id CHAR(36)"
-            )
-
-            conn.commit()
-            logger.info("Successfully added uuid_id column to research_history")
-            return True
-
-        finally:
-            conn.close()
-
-    except Exception:
-        logger.exception("Error adding uuid_id column to research_history")
-        return False
-
-
-def convert_research_id_to_string_if_needed(engine):
-    """
-    Convert research_id columns from Integer to String in all tables.
-    Preserves existing data by converting integer IDs to string format.
-    Only runs if integer research_id columns are detected.
-    """
-    try:
-        import sqlite3
-
-        # Get database path from engine
-        db_path = engine.url.database
-
-        logger.info(
-            "Checking if research_id columns need conversion to string..."
-        )
-
-        conn = sqlite3.connect(db_path)
-        conn.execute("PRAGMA foreign_keys = OFF")
-
-        try:
-            cursor = conn.cursor()
-
-            # List of tables that might have research_id columns
-            tables_to_check = [
-                "token_usage",
-                "model_usage",
-                "search_calls",
-                "benchmark_results",
-            ]
-
-            tables_needing_conversion = []
-
-            # Check which tables need conversion
-            for table_name in tables_to_check:
-                # Check if table exists
-                cursor.execute(
-                    "SELECT name FROM sqlite_master WHERE type='table' AND name=?",
-                    (table_name,),
-                )
-                if not cursor.fetchone():
-                    continue
-
-                # Check if research_id column exists and is integer type
-                cursor.execute(f"PRAGMA table_info({table_name})")
-                columns = cursor.fetchall()
-
-                for col in columns:
-                    col_name, col_type = col[1], col[2]
-                    if col_name == "research_id" and (
-                        "INTEGER" in col_type.upper()
-                        or "INT" in col_type.upper()
-                    ):
-                        tables_needing_conversion.append(table_name)
-                        break
-
-            if not tables_needing_conversion:
-                logger.info(
-                    "All research_id columns are already string type, no conversion needed"
-                )
-                return True
-
-            logger.info(
-                f"Converting research_id to string in tables: {tables_needing_conversion}"
-            )
-
-            # Convert each table
-            for table_name in tables_needing_conversion:
-                logger.info(f"Converting {table_name} table...")
-
-                # Get the current table schema
-                cursor.execute(
-                    f"SELECT sql FROM sqlite_master WHERE type='table' AND name='{table_name}'"
-                )
-                create_sql = cursor.fetchone()[0]
-
-                # Create new table name
-                new_table_name = f"{table_name}_new"
-
-                # Modify the CREATE TABLE statement to change research_id to TEXT
-                new_create_sql = create_sql.replace(
-                    f"CREATE TABLE {table_name}",
-                    f"CREATE TABLE {new_table_name}",
-                )
-                new_create_sql = new_create_sql.replace(
-                    "research_id INTEGER", "research_id TEXT"
-                )
-                new_create_sql = new_create_sql.replace(
-                    "research_id INT", "research_id TEXT"
-                )
-
-                # Create the new table
-                cursor.execute(new_create_sql)
-
-                # Copy data from old table to new table, converting research_id to string
-                cursor.execute(f"SELECT * FROM {table_name}")
-                old_rows = cursor.fetchall()
-
-                if old_rows:
-                    # Get column names
-                    cursor.execute(f"PRAGMA table_info({table_name})")
-                    columns = cursor.fetchall()
-                    column_names = [col[1] for col in columns]
-                    research_id_index = (
-                        column_names.index("research_id")
-                        if "research_id" in column_names
-                        else -1
-                    )
-
-                    # Prepare insert statement
-                    placeholders = ",".join(["?" for _ in column_names])
-                    insert_sql = f"INSERT INTO {new_table_name} ({','.join(column_names)}) VALUES ({placeholders})"
-
-                    # Convert rows and insert
-                    converted_rows = []
-                    for row in old_rows:
-                        row_list = list(row)
-                        # Convert research_id to string if it's not None
-                        if (
-                            research_id_index >= 0
-                            and row_list[research_id_index] is not None
-                        ):
-                            row_list[research_id_index] = str(
-                                row_list[research_id_index]
-                            )
-                        converted_rows.append(tuple(row_list))
-
-                    cursor.executemany(insert_sql, converted_rows)
-                    logger.info(
-                        f"Converted {len(converted_rows)} rows in {table_name}"
-                    )
-
-                # Drop old table and rename new table
-                cursor.execute(f"DROP TABLE {table_name}")
-                cursor.execute(
-                    f"ALTER TABLE {new_table_name} RENAME TO {table_name}"
-                )
-
-                logger.info(
-                    f"Successfully converted {table_name} research_id to string"
-                )
-
-            # Commit all changes
-            conn.commit()
-            logger.info(
-                "All research_id columns converted to string successfully!"
-            )
-            return True
-
-        finally:
-            conn.execute("PRAGMA foreign_keys = ON")
-            conn.close()
-
-    except Exception:
-        logger.exception("Error converting research_id columns to string")
-        return False
-
-
-def create_provider_models_table(engine):
-    """Create provider_models table for caching available models"""
-    with engine.connect() as conn:
-        result = conn.execute(
-            text(
-                "SELECT name FROM sqlite_master WHERE type='table' AND name='provider_models'"
-            )
-        )
-        if result.fetchone():
-            logger.info(
-                "Table 'provider_models' already exists, no action needed"
-            )
-            return
-
-        logger.info("Creating 'provider_models' table...")
-
-        # Create the table
-        conn.execute(
-            text(
-                """
-            CREATE TABLE provider_models (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                provider VARCHAR(50) NOT NULL,
-                model_key VARCHAR(255) NOT NULL,
-                model_label VARCHAR(255) NOT NULL,
-                model_metadata JSON,
-                last_updated DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL,
-                UNIQUE(provider, model_key)
-            )
-            """
-            )
-        )
-
-        # Create index on provider
-        conn.execute(
-            text(
-                "CREATE INDEX ix_provider_models_provider ON provider_models (provider)"
-            )
-        )
-
-        conn.commit()
-        logger.info("Table 'provider_models' created successfully")
-
-
-def run_schema_upgrades():
-    """
-    Run all schema upgrade operations on the database
-
-    Returns:
-        bool: True if all upgrades successful, False otherwise
-    """
-    # Check if database exists
-    if not os.path.exists(DB_PATH):
-        logger.warning(
-            f"Database not found at {DB_PATH}, skipping schema upgrades"
-        )
-        return False
-
-    logger.info(f"Running schema upgrades on {DB_PATH}")
-
-    try:
-        # Create SQLAlchemy engine
-        engine = create_engine(f"sqlite:///{DB_PATH}")
-
-        # 1. Remove the redundant research_log table
-        remove_research_log_table(engine)
-
-        # 2. Create research_strategies table
-        create_research_strategy_table(engine)
-
-        # 3. Create benchmark tables
-        create_benchmark_tables(engine)
-
-        # 4. Create rate limiting tables
-        create_rate_limiting_tables(engine)
-
-        # 5. Add research_id column to benchmark_results if missing
-        add_research_id_to_benchmark_results(engine)
-
-        # 6. Convert research_id columns from integer to string
-        convert_research_id_to_string_if_needed(engine)
-
-        # 7. Add uuid_id column to research_history if missing
-        add_uuid_id_column_to_research_history(engine)
-
-        # 8. Create provider_models table for caching
-        create_provider_models_table(engine)
-
-        logger.info("Schema upgrades completed successfully")
-        return True
-    except Exception:
-        logger.exception("Error during schema upgrades")
-        return False
-
-
-if __name__ == "__main__":
-    run_schema_upgrades()