database/migrations/002_create_assignments_table.py

"""
Migration 002: Create Assignments Table
Creates the assignments table for managing order-driver assignments with route data
"""

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 = """
-- Create assignments table with enhanced route tracking
CREATE TABLE IF NOT EXISTS assignments (
    assignment_id VARCHAR(50) PRIMARY KEY,
    order_id VARCHAR(50) NOT NULL,
    driver_id VARCHAR(50) NOT NULL,

    -- Assignment metadata
    assigned_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    sequence_number INTEGER,

    -- Route data from Google Routes API
    route_distance_meters INTEGER,
    route_duration_seconds INTEGER,
    route_duration_in_traffic_seconds INTEGER,
    route_summary TEXT,
    route_polyline TEXT,

    -- Origin/destination tracking
    driver_start_location_lat DECIMAL(10, 8),
    driver_start_location_lng DECIMAL(11, 8),
    delivery_location_lat DECIMAL(10, 8),
    delivery_location_lng DECIMAL(11, 8),
    delivery_address TEXT,

    -- Timing data
    estimated_arrival TIMESTAMP,
    actual_arrival TIMESTAMP,

    -- Actual vs estimated tracking
    actual_distance_meters INTEGER,

    -- Vehicle context
    vehicle_type VARCHAR(50),

    -- Status management
    status VARCHAR(20) CHECK(status IN ('active', 'in_progress', 'completed', 'cancelled', 'failed')) DEFAULT 'active',

    -- Additional metadata
    traffic_delay_seconds INTEGER,
    weather_conditions JSONB,
    route_confidence VARCHAR(100),
    notes TEXT,

    -- Foreign keys with proper cascade behavior
    FOREIGN KEY (order_id) REFERENCES orders(order_id) ON DELETE CASCADE,
    FOREIGN KEY (driver_id) REFERENCES drivers(driver_id) ON DELETE RESTRICT
);

-- Create indexes for performance
CREATE INDEX IF NOT EXISTS idx_assignments_driver ON assignments(driver_id);
CREATE INDEX IF NOT EXISTS idx_assignments_order ON assignments(order_id);
CREATE INDEX IF NOT EXISTS idx_assignments_status ON assignments(status);
CREATE INDEX IF NOT EXISTS idx_assignments_assigned_at ON assignments(assigned_at);

-- Create trigger for auto-updating updated_at
CREATE OR REPLACE FUNCTION update_updated_at_column()
RETURNS TRIGGER AS $$
BEGIN
    NEW.updated_at = CURRENT_TIMESTAMP;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER update_assignments_timestamp
    BEFORE UPDATE ON assignments
    FOR EACH ROW
    EXECUTE FUNCTION update_updated_at_column();

-- Add unique constraint to prevent multiple active assignments per order
CREATE UNIQUE INDEX IF NOT EXISTS idx_assignments_unique_active_order
    ON assignments(order_id)
    WHERE status IN ('active', 'in_progress');
"""

ROLLBACK_SQL = """
DROP INDEX IF EXISTS idx_assignments_unique_active_order;
DROP TRIGGER IF EXISTS update_assignments_timestamp ON assignments;
DROP INDEX IF EXISTS idx_assignments_assigned_at;
DROP INDEX IF EXISTS idx_assignments_status;
DROP INDEX IF EXISTS idx_assignments_order;
DROP INDEX IF EXISTS idx_assignments_driver;
DROP TABLE IF EXISTS assignments CASCADE;
"""


def up():
    """Apply migration - create assignments table"""
    print("Running migration 002: Create assignments table...")

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

        # Execute migration SQL
        cursor.execute(MIGRATION_SQL)

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

        print("SUCCESS: Migration 002 applied successfully")
        print("  - Created assignments table")
        print("  - Created indexes (driver, order, status, assigned_at)")
        print("  - Created update trigger")
        print("  - Created unique constraint for active assignments per order")
        return True

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


def down():
    """Rollback migration - drop assignments table"""
    print("Rolling back migration 002: Drop assignments table...")

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

        # Execute rollback SQL
        cursor.execute(ROLLBACK_SQL)

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

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

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


if __name__ == "__main__":
    import sys

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