"""kernel_requested_slots_to_not_nullable

Revision ID: 85615e005fa3
Revises: a083c6c962e5
Create Date: 2023-09-22 12:37:31.725324

"""

from collections.abc import Sequence
from typing import Any

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql as pgsql
from sqlalchemy.engine import Row

from ai.backend.manager.models.base import GUID, mapper_registry

# revision identifiers, used by Alembic.
revision = "85615e005fa3"
down_revision = "a083c6c962e5"
branch_labels = None
depends_on = None


BATCH_SIZE = 100


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    connection = op.get_bind()
    kernels = sa.Table(
        "kernels",
        mapper_registry.metadata,
        sa.Column(
            "id",
            GUID,
            primary_key=True,
        ),
        sa.Column("requested_slots", pgsql.JSONB()),
        extend_existing=True,
    )

    while True:
        select_stmt = (
            sa.select(kernels.c.id).where(kernels.c.requested_slots.is_(None)).limit(BATCH_SIZE)
        )
        fetch_result: Sequence[Row[Any]] = connection.execute(select_stmt).fetchall()
        kernel_ids_to_update = [row.id for row in fetch_result]

        stmt = (
            sa.update(kernels)
            .where(kernels.c.id.in_(kernel_ids_to_update))
            .values(requested_slots={})
        )
        connection.execute(stmt)
        if len(kernel_ids_to_update) < BATCH_SIZE:
            break
    op.alter_column("kernels", "requested_slots", nullable=False)
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.alter_column("kernels", "requested_slots", nullable=True)
    # ### end Alembic commands ###
