"""add provider to tasks

Revision ID: add_provider_to_tasks
Revises: bfac9b8e32f5
Create Date: 2024-02-11

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = 'add_provider_to_tasks'
down_revision = 'bfac9b8e32f5'
branch_labels = None
depends_on = None


def upgrade() -> None:
    """Add provider column to tasks table"""
    # Add provider column (nullable, indexed)
    op.add_column('tasks', sa.Column('provider', sa.String(), nullable=True))
    
    # Create index on provider column for faster queries
    op.create_index(op.f('ix_tasks_provider'), 'tasks', ['provider'], unique=False)
    
    # Optional: Migrate existing data by extracting provider from params
    # This is a data migration that can be run separately if needed
    op.execute("""
        UPDATE tasks 
        SET provider = params->>'provider'
        WHERE params->>'provider' IS NOT NULL
    """)


def downgrade() -> None:
    """Remove provider column from tasks table"""
    # Drop index first
    op.drop_index(op.f('ix_tasks_provider'), table_name='tasks')
    
    # Drop column
    op.drop_column('tasks', 'provider')