cloud-reporter/backend/migrations/versions/004_add_cloudshell_scanner_fields.py

"""Add CloudShell Scanner fields to tasks table

Revision ID: 004_cloudshell_scanner
Revises: 003_add_session_token_to_base_role
Create Date: 2026-01-23

This migration adds the source and scan_data_path fields to the tasks table
to support the CloudShell Scanner feature.

Requirements:
    - 4.3: Task model needs source field to distinguish between credential and upload tasks
"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '004_cloudshell_scanner'
down_revision = '003_add_session_token_to_base_role'
branch_labels = None
depends_on = None


def upgrade():
    """Add source and scan_data_path columns to tasks table."""
    # Add source column with default value 'credential'
    with op.batch_alter_table('tasks', schema=None) as batch_op:
        batch_op.add_column(
            sa.Column('source', sa.String(length=20), nullable=True, default='credential')
        )
        batch_op.add_column(
            sa.Column('scan_data_path', sa.String(length=500), nullable=True)
        )
    
    # Update existing rows to have 'credential' as source
    op.execute("UPDATE tasks SET source = 'credential' WHERE source IS NULL")


def downgrade():
    """Remove source and scan_data_path columns from tasks table."""
    with op.batch_alter_table('tasks', schema=None) as batch_op:
        batch_op.drop_column('scan_data_path')
        batch_op.drop_column('source')