"""add_user_id_to_spans

Revision ID: ecc45f321532
Revises: 7fb45748e807
Create Date: 2025-10-21 19:18:37.077936

"""

import sqlalchemy as sa

from alembic import op

# revision identifiers, used by Alembic.
revision = "ecc45f321532"
down_revision = "7fb45748e807"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column("spans", sa.Column("user_id", sa.String(), nullable=True))
    op.create_index(
        "idx_spans_user_task_time",
        "spans",
        ["user_id", "task_id", "start_time"],
        unique=False,
    )
    op.add_column("trace_metadata", sa.Column("user_id", sa.String(), nullable=True))
    op.create_index(
        "idx_traces_session_time",
        "trace_metadata",
        ["session_id", "start_time"],
        unique=False,
    )
    op.create_index(
        "idx_traces_task_session_time",
        "trace_metadata",
        ["task_id", "session_id", "start_time"],
        unique=False,
        postgresql_where=sa.text("session_id IS NOT NULL"),
    )
    op.create_index(
        "idx_traces_user_session",
        "trace_metadata",
        ["user_id", "session_id", "start_time"],
        unique=False,
    )
    op.create_index(
        "idx_traces_user_task_time",
        "trace_metadata",
        ["user_id", "task_id", "start_time"],
        unique=False,
    )
    # ### end Alembic commands ###

    # Backfill user_id in spans from raw_data attributes
    op.execute(
        """
        UPDATE spans
        SET user_id = raw_data->'attributes'->>'user.id'
        WHERE raw_data ? 'attributes'
          AND raw_data->'attributes' ? 'user.id'
          AND raw_data->'attributes'->>'user.id' IS NOT NULL
          AND raw_data->'attributes'->>'user.id' != ''
        """,
    )

    # Backfill user_id in trace_metadata from first span per trace
    op.execute(
        """
        WITH first_spans_per_trace AS (
            SELECT DISTINCT ON (trace_id)
                trace_id,
                user_id,
                created_at
            FROM spans
            WHERE trace_id IS NOT NULL
                AND user_id IS NOT NULL
                AND user_id != ''
            ORDER BY trace_id, created_at ASC
        )
        UPDATE trace_metadata
        SET user_id = first_spans_per_trace.user_id
        FROM first_spans_per_trace
        WHERE trace_metadata.trace_id = first_spans_per_trace.trace_id
        """,
    )


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index("idx_traces_user_task_time", table_name="trace_metadata")
    op.drop_index("idx_traces_user_session", table_name="trace_metadata")
    op.drop_index(
        "idx_traces_task_session_time",
        table_name="trace_metadata",
        postgresql_where=sa.text("session_id IS NOT NULL"),
    )
    op.drop_index("idx_traces_session_time", table_name="trace_metadata")
    op.drop_column("trace_metadata", "user_id")
    op.drop_index("idx_spans_user_task_time", table_name="spans")
    op.drop_column("spans", "user_id")
    # ### end Alembic commands ###
