"""predictor_version

Revision ID: 1e60096fc817
Revises: 43c52d23845a
Create Date: 2022-11-07 16:43:47.301692

"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.sql import text


# revision identifiers, used by Alembic.
revision = '1e60096fc817'
down_revision = '43c52d23845a'
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('predictor', schema=None) as batch_op:
        batch_op.add_column(sa.Column('label', sa.String(), nullable=True))
        batch_op.add_column(sa.Column('version', sa.Integer(), nullable=True))

    # update current predictor versions
    conn = op.get_bind()
    session = sa.orm.Session(bind=conn)

    key0 = (None, None, None)
    models = conn.execute(sa.text('''
        select company_id, project_id, name, id
        from predictor
        order by company_id, project_id, name, created_at
    ''')).fetchall()
    for model in models:
        key = (model['company_id'], model['project_id'], model['name'].lower())

        # it is different name or project or company
        if key != key0:
            version = 1
            key0 = key

        conn.execute(
            text("""
                update predictor
                set version = :version
                where id = :id
            """), {
                'version': version,
                'id': model['id']
            }
        )
        version += 1

    session.commit()
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('predictor', schema=None) as batch_op:
        batch_op.drop_column('version')
        batch_op.drop_column('label')

    # ### end Alembic commands ###
