"""add_session_status_columns

Revision ID: 38fb313d0c3a
Revises: bf11fb90a818
Create Date: 2025-10-10 10:56:13.501005

"""

import sqlalchemy as sa

from alembic import op

# revision identifiers, used by Alembic.
revision = "38fb313d0c3a"
down_revision = "bf11fb90a818"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column("spans", sa.Column("session_id", sa.String(), nullable=True))
    op.add_column(
        "spans",
        sa.Column(
            "status_code",
            sa.String(),
            server_default=sa.text("'Unset'"),
            nullable=False,
        ),
    )
    op.add_column("trace_metadata", sa.Column("session_id", sa.String(), nullable=True))
    # ### end Alembic commands ###

    # Backfill session_id in trace_metadata from existing span data
    # Get session_id from the first span (by created_at) for each trace
    op.execute(
        """
        WITH first_spans_per_trace AS (
            SELECT DISTINCT ON (trace_id)
                trace_id,
                raw_data->'attributes'->>'session.id' as session_id,
                created_at
            FROM spans
            WHERE trace_id IS NOT NULL
                AND raw_data IS NOT NULL
                AND raw_data ? 'attributes'
                AND raw_data->'attributes' ? 'session.id'
            ORDER BY trace_id, created_at ASC
        )
        UPDATE trace_metadata
        SET session_id = first_spans_per_trace.session_id
        FROM first_spans_per_trace
        WHERE trace_metadata.trace_id = first_spans_per_trace.trace_id
            AND first_spans_per_trace.session_id IS NOT NULL
            AND first_spans_per_trace.session_id != ''
        """,
    )


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column("trace_metadata", "session_id")
    op.drop_column("spans", "status_code")
    op.drop_column("spans", "session_id")
    # ### end Alembic commands ###
