database/migrations/003_add_order_driver_fk.py

"""
Migration 003: Add Foreign Key Constraint to orders.assigned_driver_id
Adds FK constraint to ensure referential integrity between orders and drivers
"""

import sys
import os

# Add parent directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from database.connection import get_db_connection

MIGRATION_SQL = """
-- Add foreign key constraint to orders.assigned_driver_id
-- Use ON DELETE SET NULL so that if driver is deleted, order is not lost
-- (The assignment record will be handled separately via RESTRICT constraint)
ALTER TABLE orders
ADD CONSTRAINT fk_orders_assigned_driver
FOREIGN KEY (assigned_driver_id)
REFERENCES drivers(driver_id)
ON DELETE SET NULL;
"""

ROLLBACK_SQL = """
-- Drop foreign key constraint
ALTER TABLE orders
DROP CONSTRAINT IF EXISTS fk_orders_assigned_driver;
"""


def up():
    """Apply migration - add FK constraint"""
    print("Running migration 003: Add FK constraint to orders.assigned_driver_id...")

    try:
        conn = get_db_connection()
        cursor = conn.cursor()

        # Execute migration SQL
        cursor.execute(MIGRATION_SQL)

        conn.commit()
        cursor.close()
        conn.close()

        print("SUCCESS: Migration 003 applied successfully")
        print("  - Added FK constraint: orders.assigned_driver_id -> drivers.driver_id")
        print("  - ON DELETE SET NULL (preserves order history if driver deleted)")
        return True

    except Exception as e:
        print(f"ERROR: Migration 003 failed: {e}")
        print("  Note: This may fail if there are existing invalid driver references")
        print("  Clean up orphaned assigned_driver_id values before running this migration")
        return False


def down():
    """Rollback migration - drop FK constraint"""
    print("Rolling back migration 003: Drop FK constraint from orders.assigned_driver_id...")

    try:
        conn = get_db_connection()
        cursor = conn.cursor()

        # Execute rollback SQL
        cursor.execute(ROLLBACK_SQL)

        conn.commit()
        cursor.close()
        conn.close()

        print("SUCCESS: Migration 003 rolled back successfully")
        return True

    except Exception as e:
        print(f"ERROR: Migration 003 rollback failed: {e}")
        return False


if __name__ == "__main__":
    import sys

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