migrations/versions/bfcf4182dcb5_.py

"""add slug field to organisations and populate it based on the name field

Revision ID: bfcf4182dcb5
Revises: 86c0f6b6a176
Create Date: 2021-03-01 09:52:03.376234

"""
from alembic import op
import sqlalchemy as sa

from slugify import slugify


# revision identifiers, used by Alembic.
revision = "bfcf4182dcb5"
down_revision = "86c0f6b6a176"
branch_labels = None
depends_on = None


def upgrade():
    conn = op.get_bind()
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "organisations", sa.Column("slug", sa.String(length=255), nullable=True)
    )
    orgs = conn.execute("select name from organisations;")
    for org in orgs:
        name = handle_special_chars(org[0])
        query = (
            f"UPDATE organisations SET slug = '{slugify(name)}' WHERE name = '{name}';"
        )
        op.execute(query)
    op.alter_column("organisations", "slug", nullable=False)
    op.create_unique_constraint("organisations_slug_key", "organisations", ["slug"])
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint("organisations_slug_key", "organisations", type_="unique")
    op.drop_column("organisations", "slug")
    # ### end Alembic commands ###


def handle_special_chars(name: str):
    special_char = name.find("'")
    if special_char >= 0:
        name = name[:special_char] + "'" + name[special_char:]
    return name