"""kernel_requested_slots_to_not_nullable

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

"""

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

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():
    # ### 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)
        )
        result = connection.execute(select_stmt).fetchall()
        kernel_ids_to_update = [row.id for row in result]

        stmt = (
            sa.update(kernels)
            .where(kernels.c.id.in_(kernel_ids_to_update))
            .values(requested_slots={})
        )
        result = 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():
    # ### commands auto generated by Alembic - please adjust! ###
    op.alter_column("kernels", "requested_slots", nullable=True)
    # ### end Alembic commands ###
