"""Fix tier CHECK constraint to allow all plan tiers

Revision ID: 006
Revises: cb71a958ad92
"""
from alembic import op

revision = "006"
down_revision = "a1b2c3d4e5f6"
branch_labels = None
depends_on = None


def upgrade() -> None:
    conn = op.get_bind()
    dialect = conn.dialect.name

    if dialect == "sqlite":
        # Migration 002 skipped the CHECK constraint on SQLite; just create it now.
        with op.batch_alter_table("users") as batch_op:
            batch_op.create_check_constraint(
                "ck_users_tier",
                "tier IN ('free', 'starter', 'pro', 'business', 'enterprise')",
            )
    else:
        op.execute("ALTER TABLE users DROP CONSTRAINT IF EXISTS ck_users_tier")
        op.execute(
            "ALTER TABLE users ADD CONSTRAINT ck_users_tier "
            "CHECK (tier IN ('free', 'starter', 'pro', 'business', 'enterprise'))"
        )


def downgrade() -> None:
    conn = op.get_bind()
    dialect = conn.dialect.name

    if dialect == "sqlite":
        with op.batch_alter_table("users") as batch_op:
            batch_op.drop_constraint("ck_users_tier", type_="check")
            batch_op.create_check_constraint(
                "ck_users_tier",
                "tier IN ('free', 'pro')",
            )
    else:
        op.execute("ALTER TABLE users DROP CONSTRAINT IF EXISTS ck_users_tier")
        op.execute(
            "ALTER TABLE users ADD CONSTRAINT ck_users_tier "
            "CHECK (tier IN ('free', 'pro'))"
        )