"""Make `trace_call_id` in the `prompt` table required

Revision ID: f433a0a15c7e
Revises: 5a61f247dd07
Create Date: 2024-06-10 11:39:41.919712

"""
import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision = 'f433a0a15c7e'
down_revision = '5a61f247dd07'
branch_labels = None
depends_on = None


def upgrade():
    # Disable the foreign key constraints during the migration.
    # https://alembic.sqlalchemy.org/en/latest/batch.html#dealing-with-referencing-foreign-keys
    op.execute('PRAGMA foreign_keys=OFF;')

    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('prompt', schema=None) as batch_op:
        batch_op.alter_column(
            'trace_call_id', existing_type=sa.INTEGER(), nullable=False
        )

    # ### end Alembic commands ###

    # Re-enable the foreign key constraints
    op.execute('PRAGMA foreign_keys=ON;')


def downgrade():
    op.execute('PRAGMA foreign_keys=OFF;')

    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('prompt', schema=None) as batch_op:
        batch_op.alter_column(
            'trace_call_id', existing_type=sa.INTEGER(), nullable=True
        )

    # ### end Alembic commands ###

    op.execute('PRAGMA foreign_keys=ON;')
