database/migrations/001_create_drivers_table.py

"""
Migration: Create drivers table
Created: 2025-11-14
"""

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 get_db_connection
from database.schema import DRIVERS_SCHEMA
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def up():
    """Apply migration - Create drivers table"""
    logger.info("Starting migration: Create drivers table")

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

        # Execute drivers schema
        logger.info("Creating drivers table...")
        cursor.execute(DRIVERS_SCHEMA)

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

        logger.info("✅ Migration completed successfully: Drivers table created")
        return True

    except Exception as e:
        logger.error(f"❌ Migration failed: {e}")
        if conn:
            conn.rollback()
        return False


def down():
    """Rollback migration - Drop drivers table"""
    logger.info("Rolling back migration: Drop drivers table")

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

        logger.info("Dropping drivers table...")
        cursor.execute("DROP TABLE IF EXISTS drivers CASCADE;")

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

        logger.info("✅ Rollback completed successfully: Drivers table dropped")
        return True

    except Exception as e:
        logger.error(f"❌ Rollback failed: {e}")
        if conn:
            conn.rollback()
        return False


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "down":
        # Rollback
        success = down()
    else:
        # Apply migration
        success = up()

    sys.exit(0 if success else 1)