"""add-keypair-resource-policy

Revision ID: dbc1e053b880
Revises: 2b0931e4a059
Create Date: 2019-02-07 15:30:54.861821

"""
from alembic import op
import sqlalchemy as sa
from ai.backend.common.types import DefaultForUnspecified
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision = 'dbc1e053b880'
down_revision = '2b0931e4a059'
branch_labels = None
depends_on = None


default_for_unspecified_choices = list(map(lambda v: v.name, DefaultForUnspecified))
default_for_unspecified = postgresql.ENUM(
    *default_for_unspecified_choices,
    name='default_for_unspecified',
)


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        'keypair_resource_policies',
        sa.Column('name', sa.String(length=256), nullable=False),
        sa.Column('created_at', sa.DateTime(timezone=True),
                  server_default=sa.text('now()'),
                  nullable=True),
        sa.Column('default_for_unspecified',
                  type_=sa.Enum(*default_for_unspecified_choices,
                                name='default_for_unspecified'),
                  nullable=False),
        sa.Column('total_resource_slots',
                  postgresql.JSONB(astext_type=sa.Text()), nullable=False),
        sa.Column('max_concurrent_sessions', sa.Integer(), nullable=False),
        sa.Column('max_containers_per_session', sa.Integer(), nullable=False),
        sa.Column('max_vfolder_count', sa.Integer(), nullable=False),
        sa.Column('max_vfolder_size', sa.BigInteger(), nullable=False),
        sa.Column('allowed_vfolder_hosts',
                  postgresql.ARRAY(sa.String()), nullable=False),
        sa.PrimaryKeyConstraint('name', name=op.f('pk_keypair_resource_policies'))
    )
    # Create a default resource policy
    query = '''
    INSERT INTO keypair_resource_policies
    VALUES (
        'default',
        now(),
        'UNLIMITED',
        '{}'::jsonb,
        30,
        1,
        10,
        0,
        array['local']
    );
    UPDATE keypairs
    SET resource_policy = 'default';
    '''
    connection = op.get_bind()
    connection.execute(query)
    print('\n!!! NOTICE !!!\n')
    print('Created a default resource policy and linked all keypairs to it.')
    print('Please inspect and adjust it!\n')
    op.alter_column('keypairs', 'resource_policy',
                    existing_type=sa.VARCHAR(),
                    type_=sa.String(length=256),
                    nullable=False)
    op.create_foreign_key(
        op.f('fk_keypairs_resource_policy_keypair_resource_policies'),
        'keypairs', 'keypair_resource_policies',
        ['resource_policy'], ['name'])
    op.drop_column('keypairs', 'concurrency_limit')
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(op.f('fk_keypairs_resource_policy_keypair_resource_policies'),
                       'keypairs', type_='foreignkey')
    op.add_column(
        'keypairs',
        sa.Column('concurrency_limit', sa.INTEGER(),
                  autoincrement=False, nullable=True))
    op.alter_column('keypairs', 'resource_policy',
                    existing_type=sa.String(length=256),
                    type_=sa.VARCHAR(),
                    nullable=True)
    op.drop_table('keypair_resource_policies')
    default_for_unspecified.drop(op.get_bind())
    # ### end Alembic commands ###
