"""Add column status in Trade

Revision ID: 33c8546d11e2
Revises: 86498a37f807
Create Date: 2022-08-13 18:55:08.104304

"""
from alembic import op
import sqlalchemy as sa
import enum
from sqlalchemy.dialects import postgresql


# revision identifiers, used by Alembic.
revision = '33c8546d11e2'
down_revision = '86498a37f807'
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    class TradeStatus(enum.Enum):
        hold = 'hold'
        open = 'open'
        close = 'close'
        cancel = 'cancel'
        target = 'target'

    trade_status = postgresql.ENUM(TradeStatus, name="tradestatus")
    trade_status.create(op.get_bind(), checkfirst=True)

    op.add_column('trades', sa.Column('status', trade_status, nullable=True, default=TradeStatus.open.value))
    op.execute("UPDATE trades SET status = 'close'")
    op.execute("UPDATE trades SET status = 'open' WHERE sell_price is NULL")
    op.alter_column('trades', 'status', nullable=False)
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###

    op.drop_column('trades', 'status')
    op.execute("DROP TYPE tradestatus;")

    # ### end Alembic commands ###
