"""force user to entity association

Revision ID: 5073b1fa473e
Revises: 1f4cbd713979

"""

# revision identifiers, used by Alembic.
revision = "5073b1fa473e"
down_revision = "1f4cbd713979"

from sqlalchemy import sql

from alembic import op


def upgrade():
    entity = sql.table(
        "entity",
        sql.column("id"),
    )
    user = sql.table(
        "userfeatures",
        sql.column("entityid"),
    )

    entity_id = _default_entity_id(entity)
    _sanitize_entity_ids(entity, user, entity_id)
    _add_foreign_key()


def downgrade():
    _delete_foreign_key()


def _add_foreign_key():
    op.create_foreign_key(
        name="userfeatures_entity_id_fkey",
        source="userfeatures",
        referent="entity",
        local_cols=["entityid"],
        remote_cols=["id"],
        ondelete="RESTRICT",
    )


def _delete_foreign_key():
    op.drop_constraint(
        name="userfeatures_entity_id_fkey",
        table_name="userfeatures",
        type_="foreignkey",
    )


def _default_entity_id(entity):
    query = sql.select([entity.c.id], limit=1)
    entity_id = op.get_bind().execute(query).scalar()
    return entity_id


def _sanitize_entity_ids(entity, user, entity_id):
    valid_entities = sql.select([entity.c.id])
    update_query = (
        user.update()
        .values(entityid=entity_id)
        .where(
            sql.or_(
                user.c.entityid == None, sql.not_(user.c.entityid.in_(valid_entities))
            )
        )
    )
    op.get_bind().execute(update_query)
