# Copyright Contributors to the Testing Farm project.
# SPDX-License-Identifier: Apache-2.0

"""
Adds user role and tokens

Revision ID: <KEY>
Revises: 2a205a7d6150
Create Date: 2020-09-16 13:40:01.642000

"""
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from alembic import op

# revision identifiers, used by Alembic.
revision = '<KEY>'
down_revision = '<KEY>'
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    if op.get_bind().dialect.name == 'postgresql':
        userroles_enum = postgresql.ENUM('USER', 'ADMIN', name='userroles', create_type=True)

        # We need an explicit create() here, because the users table is being altered in a batch.
        userroles_enum.create(op.get_bind())

    else:
        userroles_enum = sa.Enum('USER', 'ADMIN', name='userroles')

    with op.batch_alter_table('users', schema=None) as batch_op:
        batch_op.add_column(sa.Column('admin_token', sa.String(length=32), server_default='undefined', nullable=False))
        batch_op.add_column(
            sa.Column('provisioning_token', sa.String(length=32), server_default='undefined', nullable=False)
        )
        batch_op.add_column(
            sa.Column('role', userroles_enum, server_default='USER', nullable=False)
        )

    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('users', schema=None) as batch_op:
        batch_op.drop_column('role')
        batch_op.drop_column('provisioning_token')
        batch_op.drop_column('admin_token')

    if op.get_bind().dialect.name == 'postgresql':
        batch_op.execute("DROP TYPE userroles;")

    # ### end Alembic commands ###
