"""add endpoint and routing tables

Revision ID: f108628f032b
Revises: b5be363ab05c
Create Date: 2023-02-17 05:47:48.060365

"""

import sqlalchemy as sa
from alembic import op

from ai.backend.manager.models.base import GUID, EndpointIDColumnType, ResourceSlotColumn

# revision identifiers, used by Alembic.
revision = "f108628f032b"
down_revision = "b5be363ab05c"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "endpoints",
        sa.Column(
            "id",
            EndpointIDColumnType(),
            server_default=sa.text("uuid_generate_v4()"),
            nullable=False,
        ),
        sa.Column("image", GUID(), nullable=False),
        sa.Column("model", GUID(), nullable=False),
        sa.Column("domain", sa.String(length=64), nullable=False),
        sa.Column("project", GUID(), nullable=False),
        sa.Column("resource_group", sa.String(length=64), nullable=False),
        sa.Column("resource_slots", ResourceSlotColumn(astext_type=sa.Text()), nullable=False),
        sa.Column("url", sa.String(length=1024), nullable=False),
        sa.ForeignKeyConstraint(
            ["domain"],
            ["domains.name"],
            name=op.f("fk_endpoints_domain_domains"),
            ondelete="RESTRICT",
        ),
        sa.ForeignKeyConstraint(
            ["image"], ["images.id"], name=op.f("fk_endpoints_image_images"), ondelete="RESTRICT"
        ),
        sa.ForeignKeyConstraint(
            ["model"],
            ["vfolders.id"],
            name=op.f("fk_endpoints_model_vfolders"),
            ondelete="RESTRICT",
        ),
        sa.ForeignKeyConstraint(
            ["project"],
            ["groups.id"],
            name=op.f("fk_endpoints_project_groups"),
            ondelete="RESTRICT",
        ),
        sa.ForeignKeyConstraint(
            ["resource_group"],
            ["scaling_groups.name"],
            name=op.f("fk_endpoints_resource_group_scaling_groups"),
            ondelete="RESTRICT",
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_endpoints")),
        sa.UniqueConstraint("url", name=op.f("uq_endpoints_url")),
    )
    op.create_index(
        op.f("ix_endpoints_resource_group"), "endpoints", ["resource_group"], unique=False
    )
    op.create_table(
        "routings",
        sa.Column("id", GUID(), server_default=sa.text("uuid_generate_v4()"), nullable=False),
        sa.Column("endpoint", GUID(), nullable=False),
        sa.Column("session", GUID(), nullable=False),
        sa.Column("traffic_ratio", sa.Float(), nullable=False),
        sa.ForeignKeyConstraint(
            ["endpoint"],
            ["endpoints.id"],
            name=op.f("fk_routings_endpoint_endpoints"),
            ondelete="CASCADE",
        ),
        sa.ForeignKeyConstraint(
            ["session"],
            ["sessions.id"],
            name=op.f("fk_routings_session_sessions"),
            ondelete="RESTRICT",
        ),
        sa.PrimaryKeyConstraint("id", name=op.f("pk_routings")),
        sa.UniqueConstraint("endpoint", "session", name="uq_routings_endpoint_session"),
    )
    op.create_index(op.f("ix_sessions_status"), "sessions", ["status"], unique=False)
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f("ix_sessions_status"), table_name="sessions")
    op.drop_table("routings")
    op.drop_index(op.f("ix_endpoints_resource_group"), table_name="endpoints")
    op.drop_table("endpoints")
    # ### end Alembic commands ###
