src/backend/base/langflow/alembic/versions/012fb73ac359_add_folder_table.py

"""Add Folder table

Revision ID: 012fb73ac359
Revises: c153816fd85f
Create Date: 2024-05-07 12:52:16.954691

"""

from typing import Sequence, Union

import sqlalchemy as sa
import sqlmodel
from alembic import op
from sqlalchemy.engine.reflection import Inspector

# revision identifiers, used by Alembic.
revision: str = "012fb73ac359"
down_revision: Union[str, None] = "c153816fd85f"
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
    table_names = inspector.get_table_names()
    # ### commands auto generated by Alembic - please adjust! ###
    if "folder" not in table_names:
        op.create_table(
            "folder",
            sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
            sa.Column("description", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
            sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
            sa.Column("parent_id", sqlmodel.sql.sqltypes.GUID(), nullable=True),
            sa.Column("user_id", sqlmodel.sql.sqltypes.GUID(), nullable=True),
            sa.ForeignKeyConstraint(
                ["parent_id"],
                ["folder.id"],
            ),
            sa.ForeignKeyConstraint(
                ["user_id"],
                ["user.id"],
            ),
            sa.PrimaryKeyConstraint("id"),
        )
    indexes = inspector.get_indexes("folder")
    if "ix_folder_name" not in [index["name"] for index in indexes]:
        with op.batch_alter_table("folder", schema=None) as batch_op:
            batch_op.create_index(batch_op.f("ix_folder_name"), ["name"], unique=False)

    column_names = [column["name"] for column in inspector.get_columns("flow")]
    with op.batch_alter_table("flow", schema=None) as batch_op:
        if "folder_id" not in column_names:
            batch_op.add_column(sa.Column("folder_id", sqlmodel.sql.sqltypes.GUID(), nullable=True))
            batch_op.create_foreign_key("flow_folder_id_fkey", "folder", ["folder_id"], ["id"])
        if "folder" in column_names:
            batch_op.drop_column("folder")

    # ### end Alembic commands ###


def downgrade() -> None:
    conn = op.get_bind()
    inspector = Inspector.from_engine(conn)  # type: ignore
    table_names = inspector.get_table_names()
    # ### commands auto generated by Alembic - please adjust! ###
    column_names = [column["name"] for column in inspector.get_columns("flow")]
    with op.batch_alter_table("flow", schema=None) as batch_op:
        if "folder" not in column_names:
            batch_op.add_column(sa.Column("folder", sa.VARCHAR(), nullable=True))
        if "folder_id" in column_names:
            batch_op.drop_column("folder_id")
            batch_op.drop_constraint("flow_folder_id_fkey", type_="foreignkey")

    indexes = inspector.get_indexes("folder")
    if "ix_folder_name" in [index["name"] for index in indexes]:
        with op.batch_alter_table("folder", schema=None) as batch_op:
            batch_op.drop_index(batch_op.f("ix_folder_name"))

    if "folder" in table_names:
        op.drop_table("folder")
    # ### end Alembic commands ###