3bb0ddf32dfb_add_unique_constraints_per_user_in_flow_.py

"""Add unique constraints per user in flow table

Revision ID: 3bb0ddf32dfb
Revises: a72f5cf9c2f9
Create Date: 2024-05-29 23:08:43.935040

"""

from typing import Sequence, Union

from alembic import op
from sqlalchemy.engine.reflection import Inspector

# revision identifiers, used by Alembic.
revision: str = "3bb0ddf32dfb"
down_revision: Union[str, None] = "a72f5cf9c2f9"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    conn = op.get_bind()
    inspector = Inspector.from_engine(conn)  # type: ignore
    # ### commands auto generated by Alembic - please adjust! ###
    indexes_names = [index["name"] for index in inspector.get_indexes("flow")]
    constraints_names = [constraint["name"] for constraint in inspector.get_unique_constraints("flow")]
    with op.batch_alter_table("flow", schema=None) as batch_op:
        if "ix_flow_endpoint_name" in indexes_names:
            batch_op.drop_index("ix_flow_endpoint_name")
            batch_op.create_index(batch_op.f("ix_flow_endpoint_name"), ["endpoint_name"], unique=False)
        if "unique_flow_endpoint_name" not in constraints_names:
            batch_op.create_unique_constraint("unique_flow_endpoint_name", ["user_id", "endpoint_name"])
        if "unique_flow_name" not in constraints_names:
            batch_op.create_unique_constraint("unique_flow_name", ["user_id", "name"])

    # ### end Alembic commands ###


def downgrade() -> None:
    conn = op.get_bind()
    inspector = Inspector.from_engine(conn)  # type: ignore
    # ### commands auto generated by Alembic - please adjust! ###
    indexes_names = [index["name"] for index in inspector.get_indexes("flow")]
    constraints_names = [constraint["name"] for constraint in inspector.get_unique_constraints("flow")]
    with op.batch_alter_table("flow", schema=None) as batch_op:
        if "unique_flow_name" in constraints_names:
            batch_op.drop_constraint("unique_flow_name", type_="unique")
        if "unique_flow_endpoint_name" in constraints_names:
            batch_op.drop_constraint("unique_flow_endpoint_name", type_="unique")
        if "ix_flow_endpoint_name" in indexes_names:
            batch_op.drop_index(batch_op.f("ix_flow_endpoint_name"))
            batch_op.create_index("ix_flow_endpoint_name", ["endpoint_name"], unique=1)

    # ### end Alembic commands ###