database/migrations/007_add_user_id.py

"""
Migration 007: Add user_id columns for multi-tenant support

This migration adds user_id to all tables to enable user-specific data isolation.
Each user will only see their own orders, drivers, assignments, etc.
"""

import sys
from pathlib import Path

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

from database.connection import execute_write


def up():
    """Add user_id columns and indexes"""

    migrations = [
        # Add user_id column to orders table
        """
        ALTER TABLE orders
        ADD COLUMN IF NOT EXISTS user_id VARCHAR(255);
        """,

        # Add user_id column to drivers table
        """
        ALTER TABLE drivers
        ADD COLUMN IF NOT EXISTS user_id VARCHAR(255);
        """,

        # Add user_id column to assignments table
        """
        ALTER TABLE assignments
        ADD COLUMN IF NOT EXISTS user_id VARCHAR(255);
        """,

        # Add user_id column to exceptions table
        """
        ALTER TABLE exceptions
        ADD COLUMN IF NOT EXISTS user_id VARCHAR(255);
        """,

        # Add user_id column to agent_decisions table
        """
        ALTER TABLE agent_decisions
        ADD COLUMN IF NOT EXISTS user_id VARCHAR(255);
        """,

        # Add user_id column to metrics table
        """
        ALTER TABLE metrics
        ADD COLUMN IF NOT EXISTS user_id VARCHAR(255);
        """,

        # Create indexes for fast user-based filtering
        """
        CREATE INDEX IF NOT EXISTS idx_orders_user_id ON orders(user_id);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_drivers_user_id ON drivers(user_id);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_assignments_user_id ON assignments(user_id);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_exceptions_user_id ON exceptions(user_id);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_agent_decisions_user_id ON agent_decisions(user_id);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_metrics_user_id ON metrics(user_id);
        """,

        # Create composite indexes for common queries
        """
        CREATE INDEX IF NOT EXISTS idx_orders_user_status ON orders(user_id, status);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_orders_user_created ON orders(user_id, created_at DESC);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_drivers_user_status ON drivers(user_id, status);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_assignments_user_driver ON assignments(user_id, driver_id);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_assignments_user_order ON assignments(user_id, order_id);
        """,
    ]

    print("Migration 007: Adding user_id columns...")

    for i, sql in enumerate(migrations, 1):
        try:
            print(f"  [{i}/{len(migrations)}] Executing: {sql.strip()[:60]}...")
            execute_write(sql)
            print(f"  Success")
        except Exception as e:
            print(f"  Warning: {e}")
            # Continue even if column already exists

    print("\nMigration 007 complete!")
    print("\nNext steps:")
    print("  1. Existing data will have NULL user_id (that's OK for now)")
    print("  2. New data will automatically get user_id from authentication")
    print("  3. You can optionally run a data migration to assign existing records to a test user")


def down():
    """Remove user_id columns and indexes (rollback)"""

    rollback_migrations = [
        # Drop indexes first
        "DROP INDEX IF EXISTS idx_assignments_user_order;",
        "DROP INDEX IF EXISTS idx_assignments_user_driver;",
        "DROP INDEX IF EXISTS idx_drivers_user_status;",
        "DROP INDEX IF EXISTS idx_orders_user_created;",
        "DROP INDEX IF EXISTS idx_orders_user_status;",
        "DROP INDEX IF EXISTS idx_metrics_user_id;",
        "DROP INDEX IF EXISTS idx_agent_decisions_user_id;",
        "DROP INDEX IF EXISTS idx_exceptions_user_id;",
        "DROP INDEX IF EXISTS idx_assignments_user_id;",
        "DROP INDEX IF EXISTS idx_drivers_user_id;",
        "DROP INDEX IF EXISTS idx_orders_user_id;",

        # Drop columns
        "ALTER TABLE metrics DROP COLUMN IF EXISTS user_id;",
        "ALTER TABLE agent_decisions DROP COLUMN IF EXISTS user_id;",
        "ALTER TABLE exceptions DROP COLUMN IF EXISTS user_id;",
        "ALTER TABLE assignments DROP COLUMN IF EXISTS user_id;",
        "ALTER TABLE drivers DROP COLUMN IF EXISTS user_id;",
        "ALTER TABLE orders DROP COLUMN IF EXISTS user_id;",
    ]

    print("Rolling back Migration 007...")

    for i, sql in enumerate(rollback_migrations, 1):
        try:
            print(f"  [{i}/{len(rollback_migrations)}] {sql[:60]}...")
            execute_write(sql)
            print(f"  Success")
        except Exception as e:
            print(f"  Warning: {e}")

    print("\nRollback complete!")


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "down":
        down()
    else:
        up()