"""Add PII Type enum to column

Revision ID: d2c711b84996
Revises: d1daff1715f7
Create Date: 2021-08-15 19:51:09.656527

"""
import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
from sqlalchemy.dialects import postgresql

revision = "d2c711b84996"
down_revision = "d1daff1715f7"
branch_labels = None
depends_on = None


def upgrade():
    bind = op.get_bind()
    if bind.engine.name == "postgresql":
        pii_types = postgresql.ENUM(
            "NONE",
            "UNSUPPORTED",
            "PHONE",
            "EMAIL",
            "CREDIT_CARD",
            "ADDRESS",
            "PERSON",
            "LOCATION",
            "BIRTH_DATE",
            "GENDER",
            "NATIONALITY",
            "IP_ADDRESS",
            "SSN",
            "USER_NAME",
            "PASSWORD",
            name="piitypes",
        )
        pii_types.create(op.get_bind())
        op.add_column(
            "columns",
            sa.Column(
                "pii_type",
                postgresql.ENUM(
                    "NONE",
                    "UNSUPPORTED",
                    "PHONE",
                    "EMAIL",
                    "CREDIT_CARD",
                    "ADDRESS",
                    "PERSON",
                    "LOCATION",
                    "BIRTH_DATE",
                    "GENDER",
                    "NATIONALITY",
                    "IP_ADDRESS",
                    "SSN",
                    "USER_NAME",
                    "PASSWORD",
                    name="piitypes",
                ),
                nullable=True,
            ),
        )
    else:
        # ### commands auto generated by Alembic - please adjust! ###
        op.add_column(
            "columns",
            sa.Column(
                "pii_type",
                sa.Enum(
                    "NONE",
                    "UNSUPPORTED",
                    "PHONE",
                    "EMAIL",
                    "CREDIT_CARD",
                    "ADDRESS",
                    "PERSON",
                    "LOCATION",
                    "BIRTH_DATE",
                    "GENDER",
                    "NATIONALITY",
                    "IP_ADDRESS",
                    "SSN",
                    "USER_NAME",
                    "PASSWORD",
                    name="piitypes",
                ),
                nullable=True,
            ),
        )
        # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column("columns", "pii_type")
    # ### end Alembic commands ###
    bind = op.get_bind()
    if bind.engine.name == "postgresql":
        pii_types = postgresql.ENUM(
            "NONE",
            "UNSUPPORTED",
            "PHONE",
            "EMAIL",
            "CREDIT_CARD",
            "ADDRESS",
            "PERSON",
            "LOCATION",
            "BIRTH_DATE",
            "GENDER",
            "NATIONALITY",
            "IP_ADDRESS",
            "SSN",
            "USER_NAME",
            "PASSWORD",
            name="piitypes",
        )
        pii_types.drop(op.get_bind())
