bdlg2023-sms-back/alembic/versions/2024052425_add_slot_template_4ca2e4cc7b95.py

"""add_slot_template

Revision ID: 4ca2e4cc7b95
Revises: 2f90ef72e3b9
Create Date: 2024-05-24 09:25:42.163966

"""
from datetime import datetime
from typing import List, Optional
from alembic import op
import sqlalchemy as sa
from sqlalchemy import orm
from sqlalchemy.ext.declarative import declarative_base

import uuid
from sqlalchemy.dialects.postgresql import UUID


# revision identifiers, used by Alembic.
revision = "4ca2e4cc7b95"
down_revision = "2f90ef72e3b9"
branch_labels = None
depends_on = None


def uid_column() -> orm.Mapped[str]:
    """Returns a postgreSQL UUID column for SQL ORM"""
    return orm.mapped_column(
        UUID(as_uuid=False), primary_key=True, default=lambda _: str(uuid.uuid4())
    )


Base = declarative_base()


class Project(Base):
    __tablename__ = "projects"
    id: orm.Mapped[str] = uid_column()


class Slot(Base):
    __tablename__ = "slots"
    id = uid_column()
    project_id = orm.mapped_column(sa.ForeignKey("projects.id", ondelete="CASCADE"))
    title: orm.Mapped[str] = orm.mapped_column(sa.String(128), nullable=False)

    description: orm.Mapped[str] = orm.mapped_column(sa.String(), default="")
    place: orm.Mapped[str] = orm.mapped_column(sa.String(), default="")
    responsible_contact: orm.Mapped[str] = orm.mapped_column(sa.String(), default="")

    template_id: orm.Mapped[Optional[str]] = orm.mapped_column(sa.ForeignKey("slot_templates.id"))
    template: orm.Mapped["SlotTemplate"] = orm.relationship(back_populates="slots")


class SlotTemplate(Base):
    __tablename__ = "slot_templates"

    id = uid_column()
    project_id = orm.mapped_column(sa.ForeignKey("projects.id", ondelete="CASCADE"))
    created_at = orm.mapped_column(sa.DateTime(timezone=True), server_default=sa.func.now())
    updated_at = orm.mapped_column(
        sa.DateTime(timezone=True), default=datetime.now, onupdate=sa.func.now()
    )

    title: orm.Mapped[str] = orm.mapped_column(sa.String(), default="")

    description: orm.Mapped[str] = orm.mapped_column(sa.String(), default="")
    place: orm.Mapped[str] = orm.mapped_column(sa.String(), default="")
    responsible_contact: orm.Mapped[str] = orm.mapped_column(sa.String(), default="")

    slots: orm.Mapped[List[Slot]] = orm.relationship(back_populates="template")


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "slot_tags",
        sa.Column("id", sa.UUID(as_uuid=False), nullable=False),
        sa.Column("project_id", sa.UUID(as_uuid=False), nullable=False),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
        sa.Column("title", sa.String(), nullable=False),
        sa.ForeignKeyConstraint(["project_id"], ["projects.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "slot_templates",
        sa.Column("id", sa.UUID(as_uuid=False), nullable=False),
        sa.Column("project_id", sa.UUID(as_uuid=False), nullable=False),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            server_default=sa.text("now()"),
            nullable=False,
        ),
        sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
        sa.Column("title", sa.String(), nullable=False),
        sa.Column("description", sa.String(), nullable=False),
        sa.Column("place", sa.String(), nullable=False),
        sa.Column("responsible_contact", sa.String(), nullable=False),
        sa.ForeignKeyConstraint(["project_id"], ["projects.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "association_description_tag",
        sa.Column("description_id", sa.UUID(as_uuid=False), nullable=False),
        sa.Column("tag_id", sa.UUID(as_uuid=False), nullable=False),
        sa.ForeignKeyConstraint(["description_id"], ["slot_templates.id"], ondelete="CASCADE"),
        sa.ForeignKeyConstraint(["tag_id"], ["slot_tags.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("description_id", "tag_id"),
    )
    op.add_column(
        "slots", sa.Column("required_volunteers", sa.Integer(), nullable=False, server_default="0")
    )
    op.alter_column("slots", "required_volunteers", server_default=None)
    op.add_column("slots", sa.Column("template_id", sa.UUID(as_uuid=False), nullable=True))
    op.create_foreign_key(None, "slots", "slot_templates", ["template_id"], ["id"])

    bind = op.get_bind()
    session = orm.Session(bind=bind)
    # Create 1 template for each slot
    for slot in session.query(Slot).all():
        template = SlotTemplate(
            project_id=slot.project_id,
            title="Migration template for " + slot.title,
            description=slot.description,
            place=slot.place,
            responsible_contact=slot.responsible_contact,
        )
        session.add(template)
        slot.template = template
    session.commit()

    op.drop_column("slots", "description")
    op.drop_column("slots", "place")
    op.drop_column("slots", "responsible_contact")
    # ### end Alembic commands ###


def downgrade():
    op.add_column(
        "slots",
        sa.Column(
            "responsible_contact",
            sa.VARCHAR(),
            server_default=sa.text("''::character varying"),
            autoincrement=False,
            nullable=False,
        ),
    )
    op.add_column(
        "slots",
        sa.Column(
            "place",
            sa.VARCHAR(),
            server_default=sa.text("''::character varying"),
            autoincrement=False,
            nullable=False,
        ),
    )
    op.add_column(
        "slots", sa.Column("description", sa.VARCHAR(), autoincrement=False, nullable=False)
    )

    # Revert data from template to slot
    bind = op.get_bind()
    session = orm.Session(bind=bind)

    for slot in session.query(Slot).all():
        slot.place = slot.template.place
        slot.description = slot.template.description
        slot.responsible_contact = slot.template.responsible_contact
    session.commit()

    op.drop_column("slots", "template_id")
    op.drop_column("slots", "required_volunteers")
    op.drop_table("association_description_tag")
    op.drop_table("slot_templates")
    op.drop_table("slot_tags")
    # ### end Alembic commands ###