# type: ignore
"""add checkpoints and run.lock_transaction column

Revision ID: bc2e195788ba
Revises: 0b45de1602a6
Create Date: 2025-03-13 16:57:24.338393

"""

import sqlalchemy as sa
from alembic import op

from ixmp4.db.migrations import utils

# Revision identifiers, used by Alembic.
revision = "bc2e195788ba"
down_revision = "0b45de1602a6"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "checkpoint",
        sa.Column("run__id", sa.Integer(), nullable=False),
        sa.Column("transaction__id", sa.Integer(), nullable=False),
        sa.Column("message", sa.String(), nullable=False),
        sa.Column(
            "id",
            sa.Integer(),
            sa.Identity(always=False, on_null=True, start=1, increment=1),
            nullable=False,
        ),
        sa.ForeignKeyConstraint(
            ["run__id"], ["run.id"], name=op.f("fk_checkpoint_run__id_run")
        ),
        sa.ForeignKeyConstraint(
            ["transaction__id"],
            ["transaction.id"],
            name=op.f("fk_checkpoint_transaction__id_transaction"),
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_checkpoint")),
    )
    with op.batch_alter_table("checkpoint", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_checkpoint_run__id"), ["run__id"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_checkpoint_transaction__id"),
            ["transaction__id"],
            unique=False,
        )

    with op.batch_alter_table("run", schema=None) as batch_op:
        batch_op.add_column(sa.Column("lock_transaction", sa.Integer(), nullable=True))
        batch_op.create_index(
            batch_op.f("ix_run_lock_transaction"), ["lock_transaction"], unique=False
        )

    with op.batch_alter_table("run_version", schema=None) as batch_op:
        batch_op.add_column(
            sa.Column(
                "lock_transaction", sa.Integer(), autoincrement=False, nullable=True
            )
        )
        batch_op.create_index(
            batch_op.f("ix_run_version_lock_transaction"),
            ["lock_transaction"],
            unique=False,
        )

    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("run_version", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_run_version_lock_transaction"))
        batch_op.drop_column("lock_transaction")

    with utils.pause_foreign_key_checks():
        with op.batch_alter_table("run", schema=None) as batch_op:
            batch_op.drop_index(batch_op.f("ix_run_lock_transaction"))
            batch_op.drop_column("lock_transaction")

    with op.batch_alter_table("checkpoint", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_checkpoint_transaction__id"))
        batch_op.drop_index(batch_op.f("ix_checkpoint_run__id"))

    op.drop_table("checkpoint")
    # ### end Alembic commands ###
