"""add job_id to vertex_build, create job status table.

Revision ID: 369268b9af8b
Revises: 182e5471b900
Create Date: 2026-01-28 13:00:52.967282

Phase: EXPAND
"""

from collections.abc import Sequence

import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision: str = "369268b9af8b"  # pragma: allowlist secret
down_revision: str | None = "182e5471b900"  # pragma: allowlist secret
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    from langflow.utils import migration

    conn = op.get_bind()
    if not migration.table_exists("job", conn):
        op.create_table(
            "job",
            sa.Column("job_id", sa.Uuid(), nullable=False),
            sa.Column("flow_id", sa.Uuid(), nullable=False),
            sa.Column(
                "status",
                sa.Enum(
                    "queued",
                    "in_progress",
                    "completed",
                    "failed",
                    "cancelled",
                    "timed_out",
                    name="job_status_enum",
                ),
                nullable=False,
            ),
            sa.Column("created_timestamp", sa.DateTime(timezone=True), nullable=False),
            sa.Column("finished_timestamp", sa.DateTime(timezone=True), nullable=True),
            sa.PrimaryKeyConstraint("job_id"),
        )
        with op.batch_alter_table("job", schema=None) as batch_op:
            batch_op.create_index(batch_op.f("ix_job_flow_id"), ["flow_id"], unique=False)
            batch_op.create_index(batch_op.f("ix_job_job_id"), ["job_id"], unique=False)
            batch_op.create_index(batch_op.f("ix_job_status"), ["status"], unique=False)

    if not migration.column_exists("vertex_build", "job_id", conn):
        with op.batch_alter_table("vertex_build", schema=None) as batch_op:
            batch_op.add_column(sa.Column("job_id", sa.Uuid(), nullable=True))
            batch_op.create_index(batch_op.f("ix_vertex_build_job_id"), ["job_id"], unique=False)

    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    from langflow.utils import migration

    conn = op.get_bind()

    # Drop vertex_build column if it exists
    if migration.column_exists("vertex_build", "job_id", conn):
        with op.batch_alter_table("vertex_build", schema=None) as batch_op:
            try:
                batch_op.drop_index(batch_op.f("ix_vertex_build_job_id"))
            except Exception:  # noqa: S110
                pass  # Index might not exist
            batch_op.drop_column("job_id")

    # Drop job table if it exists (this will automatically drop all indexes)
    if migration.table_exists("job", conn):
        op.drop_table("job")
    # ### end Alembic commands ###
