"""optimizing database [37835ce041d2].

Revision ID: 37835ce041d2
Revises: 0.44.3
Create Date: 2023-10-04 09:28:23.761274

"""

import sqlalchemy as sa
import sqlmodel
from alembic import op
from sqlalchemy.dialects import mysql
from sqlalchemy.engine.base import Connection
from sqlalchemy.sql import text

# revision identifiers, used by Alembic.
revision = "37835ce041d2"
down_revision = "0.44.3"
branch_labels = None
depends_on = None


def upgrade() -> None:
    """Upgrade database schema and/or data, creating a new revision."""
    # ### commands auto generated by Alembic - please adjust! ###
    bind: Connection = op.get_bind()
    column_info = sa.inspect(bind).get_columns("pipeline_deployment")
    columns = [c["name"] for c in column_info]

    with op.batch_alter_table("pipeline_deployment", schema=None) as batch_op:
        if "client_version" not in columns:
            batch_op.add_column(
                sa.Column(
                    "client_version",
                    sqlmodel.sql.sqltypes.AutoString(),
                    nullable=True,
                )
            )
        else:
            batch_op.alter_column(
                "client_version",
                existing_type=sqlmodel.sql.sqltypes.AutoString(),
                nullable=True,
            )

        if "server_version" not in columns:
            batch_op.add_column(
                sa.Column(
                    "server_version",
                    sqlmodel.sql.sqltypes.AutoString(),
                    nullable=True,
                )
            )
        else:
            batch_op.alter_column(
                "server_version",
                existing_type=sqlmodel.sql.sqltypes.AutoString(),
                nullable=True,
            )

        batch_op.alter_column(
            "pipeline_configuration",
            existing_type=sa.TEXT(),
            type_=sa.String(length=16777215).with_variant(
                mysql.MEDIUMTEXT, "mysql"
            ),
            existing_nullable=False,
        )

    # Fill in the values and make fields nullable
    connection = op.get_bind()

    update_client_version_query = text(
        """
        UPDATE pipeline_deployment
        SET client_version = (
            SELECT max(pipeline_run.client_version)
            FROM pipeline_run
            WHERE pipeline_run.deployment_id = pipeline_deployment.id
        )
        WHERE EXISTS (
            SELECT 1
            FROM pipeline_run
            WHERE pipeline_run.deployment_id = pipeline_deployment.id
        )
        """
    )
    connection.execute(update_client_version_query)

    update_server_version_query = text(
        """
        UPDATE pipeline_deployment
        SET server_version = (
            SELECT max(pipeline_run.server_version)
            FROM pipeline_run
            WHERE pipeline_run.deployment_id = pipeline_deployment.id
        )
        WHERE EXISTS (
            SELECT 1
            FROM pipeline_run
            WHERE pipeline_run.deployment_id = pipeline_deployment.id
        )
        """
    )
    connection.execute(update_server_version_query)

    with op.batch_alter_table("step_run", schema=None) as batch_op:
        batch_op.add_column(
            sa.Column(
                "deployment_id", sqlmodel.sql.sqltypes.GUID(), nullable=True
            )
        )

        batch_op.create_foreign_key(
            "fk_step_run_deployment_id_pipeline_deployment",
            "pipeline_deployment",
            ["deployment_id"],
            ["id"],
            ondelete="CASCADE",
        )
        batch_op.alter_column(
            "step_configuration",
            existing_type=sa.String(length=16777215).with_variant(
                mysql.MEDIUMTEXT, "mysql"
            ),
            nullable=True,
        )

        batch_op.drop_column("parameters")
        batch_op.drop_column("enable_artifact_metadata")
        batch_op.drop_column("num_outputs")
        batch_op.drop_column("enable_cache")
        batch_op.drop_column("caching_parameters")
        batch_op.drop_column("entrypoint_name")

    connection = op.get_bind()

    update_deployment_id = text(
        """
        UPDATE step_run
        SET deployment_id = (
            SELECT pipeline_run.deployment_id
            FROM pipeline_run
            WHERE pipeline_run.id = step_run.pipeline_run_id
        )
        WHERE EXISTS (
            SELECT 1
            FROM pipeline_run
            WHERE pipeline_run.id = step_run.pipeline_run_id
        )
        """
    )
    connection.execute(update_deployment_id)

    with op.batch_alter_table("pipeline_run", schema=None) as batch_op:
        batch_op.drop_constraint(
            "fk_pipeline_run_deployment_id_pipeline_deployment",
            type_="foreignkey",
        )
        batch_op.create_foreign_key(
            "fk_pipeline_run_deployment_id_pipeline_deployment",
            "pipeline_deployment",
            ["deployment_id"],
            ["id"],
            ondelete="CASCADE",
        )
        batch_op.drop_column("client_version")
        batch_op.drop_column("git_sha")
        batch_op.drop_column("server_version")
        batch_op.drop_column("num_steps")
        batch_op.drop_column("enable_artifact_metadata")
        batch_op.drop_column("enable_cache")

    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade database schema and/or data back to the previous revision."""
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("step_run", schema=None) as batch_op:
        batch_op.add_column(
            sa.Column("entrypoint_name", sa.VARCHAR(), nullable=False)
        )
        batch_op.add_column(
            sa.Column("caching_parameters", sa.TEXT(), nullable=True)
        )
        batch_op.add_column(
            sa.Column("enable_cache", sa.BOOLEAN(), nullable=True)
        )
        batch_op.add_column(
            sa.Column("num_outputs", sa.INTEGER(), nullable=True)
        )
        batch_op.add_column(
            sa.Column("enable_artifact_metadata", sa.BOOLEAN(), nullable=True)
        )
        batch_op.add_column(
            sa.Column(
                "parameters", sa.VARCHAR(length=16777215), nullable=False
            )
        )
        batch_op.drop_constraint(
            "fk_step_run_deployment_id_pipeline_deployment", type_="foreignkey"
        )
        batch_op.drop_column("deployment_id")

    with op.batch_alter_table("pipeline_run", schema=None) as batch_op:
        batch_op.add_column(
            sa.Column("build_id", sa.CHAR(length=32), nullable=True)
        )
        batch_op.add_column(
            sa.Column("enable_cache", sa.BOOLEAN(), nullable=True)
        )
        batch_op.add_column(
            sa.Column("enable_artifact_metadata", sa.BOOLEAN(), nullable=True)
        )
        batch_op.add_column(
            sa.Column("pipeline_configuration", sa.TEXT(), nullable=False)
        )
        batch_op.add_column(
            sa.Column("client_environment", sa.TEXT(), nullable=True)
        )
        batch_op.add_column(
            sa.Column("num_steps", sa.INTEGER(), nullable=True)
        )
        batch_op.add_column(
            sa.Column("server_version", sa.VARCHAR(), nullable=True)
        )
        batch_op.add_column(
            sa.Column("schedule_id", sa.CHAR(length=32), nullable=True)
        )
        batch_op.add_column(
            sa.Column("stack_id", sa.CHAR(length=32), nullable=True)
        )
        batch_op.add_column(sa.Column("git_sha", sa.VARCHAR(), nullable=True))
        batch_op.add_column(
            sa.Column("client_version", sa.VARCHAR(), nullable=False)
        )
        batch_op.drop_constraint(
            "fk_pipeline_run_deployment_id_pipeline_deployment",
            type_="foreignkey",
        )
        batch_op.create_foreign_key(
            "fk_pipeline_run_deployment_id_pipeline_deployment",
            "pipeline_deployment",
            ["deployment_id"],
            ["id"],
            ondelete="SET NULL",
        )
        batch_op.create_foreign_key(
            "fk_pipeline_run_schedule_id_schedule",
            "schedule",
            ["schedule_id"],
            ["id"],
            ondelete="SET NULL",
        )
        batch_op.create_foreign_key(
            "fk_pipeline_run_build_id_pipeline_build",
            "pipeline_build",
            ["build_id"],
            ["id"],
            ondelete="SET NULL",
        )
        batch_op.create_foreign_key(
            "fk_pipeline_run_stack_id_stack",
            "stack",
            ["stack_id"],
            ["id"],
            ondelete="SET NULL",
        )

    with op.batch_alter_table("pipeline_deployment", schema=None) as batch_op:
        batch_op.alter_column(
            "pipeline_configuration",
            existing_type=sa.String(length=16777215).with_variant(
                mysql.MEDIUMTEXT, "mysql"
            ),
            type_=sa.TEXT(),
            existing_nullable=False,
        )
        batch_op.drop_column("server_version")
        batch_op.drop_column("client_version")

    # ### end Alembic commands ###
