"""v2

Revision ID: 5bf02391cfef
Revises: 9fd26a2b0de4
Create Date: 2024-04-18 10:35:19.067833

"""
import sqlalchemy as sa
import sqlmodel
from alembic import op


# revision identifiers, used by Alembic.
revision = "5bf02391cfef"
down_revision = "9fd26a2b0de4"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "collectionstatev2",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("data", sa.JSON(), nullable=True),
        sa.Column("timestamp", sa.DateTime(timezone=True), nullable=True),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "projectv2",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column(
            "timestamp_created", sa.DateTime(timezone=True), nullable=False
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "taskv2",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("type", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column(
            "command_non_parallel",
            sqlmodel.sql.sqltypes.AutoString(),
            nullable=True,
        ),
        sa.Column(
            "command_parallel",
            sqlmodel.sql.sqltypes.AutoString(),
            nullable=True,
        ),
        sa.Column(
            "source", sqlmodel.sql.sqltypes.AutoString(), nullable=False
        ),
        sa.Column(
            "meta_non_parallel", sa.JSON(), server_default="{}", nullable=False
        ),
        sa.Column(
            "meta_parallel", sa.JSON(), server_default="{}", nullable=False
        ),
        sa.Column("owner", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column(
            "version", sqlmodel.sql.sqltypes.AutoString(), nullable=True
        ),
        sa.Column("args_schema_non_parallel", sa.JSON(), nullable=True),
        sa.Column("args_schema_parallel", sa.JSON(), nullable=True),
        sa.Column(
            "args_schema_version",
            sqlmodel.sql.sqltypes.AutoString(),
            nullable=True,
        ),
        sa.Column(
            "docs_info", sqlmodel.sql.sqltypes.AutoString(), nullable=True
        ),
        sa.Column(
            "docs_link", sqlmodel.sql.sqltypes.AutoString(), nullable=True
        ),
        sa.Column("input_types", sa.JSON(), nullable=True),
        sa.Column("output_types", sa.JSON(), nullable=True),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("source"),
    )
    op.create_table(
        "datasetv2",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("project_id", sa.Integer(), nullable=False),
        sa.Column("history", sa.JSON(), server_default="[]", nullable=False),
        sa.Column(
            "timestamp_created", sa.DateTime(timezone=True), nullable=False
        ),
        sa.Column(
            "zarr_dir", sqlmodel.sql.sqltypes.AutoString(), nullable=False
        ),
        sa.Column("images", sa.JSON(), server_default="[]", nullable=False),
        sa.Column(
            "filters",
            sa.JSON(),
            server_default='{"attributes": {}, "types": {}}',
            nullable=False,
        ),
        sa.ForeignKeyConstraint(
            ["project_id"],
            ["projectv2.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "linkuserprojectv2",
        sa.Column("project_id", sa.Integer(), nullable=False),
        sa.Column("user_id", sa.Integer(), nullable=False),
        sa.ForeignKeyConstraint(
            ["project_id"],
            ["projectv2.id"],
        ),
        sa.ForeignKeyConstraint(
            ["user_id"],
            ["user_oauth.id"],
        ),
        sa.PrimaryKeyConstraint("project_id", "user_id"),
    )
    op.create_table(
        "workflowv2",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("project_id", sa.Integer(), nullable=False),
        sa.Column(
            "timestamp_created", sa.DateTime(timezone=True), nullable=False
        ),
        sa.ForeignKeyConstraint(
            ["project_id"],
            ["projectv2.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "jobv2",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("project_id", sa.Integer(), nullable=True),
        sa.Column("workflow_id", sa.Integer(), nullable=True),
        sa.Column("dataset_id", sa.Integer(), nullable=True),
        sa.Column(
            "user_email", sqlmodel.sql.sqltypes.AutoString(), nullable=False
        ),
        sa.Column(
            "slurm_account", sqlmodel.sql.sqltypes.AutoString(), nullable=True
        ),
        sa.Column("dataset_dump", sa.JSON(), nullable=False),
        sa.Column("workflow_dump", sa.JSON(), nullable=False),
        sa.Column("project_dump", sa.JSON(), nullable=False),
        sa.Column(
            "worker_init", sqlmodel.sql.sqltypes.AutoString(), nullable=True
        ),
        sa.Column(
            "working_dir", sqlmodel.sql.sqltypes.AutoString(), nullable=True
        ),
        sa.Column(
            "working_dir_user",
            sqlmodel.sql.sqltypes.AutoString(),
            nullable=True,
        ),
        sa.Column("first_task_index", sa.Integer(), nullable=False),
        sa.Column("last_task_index", sa.Integer(), nullable=False),
        sa.Column(
            "start_timestamp", sa.DateTime(timezone=True), nullable=False
        ),
        sa.Column("end_timestamp", sa.DateTime(timezone=True), nullable=True),
        sa.Column(
            "status", sqlmodel.sql.sqltypes.AutoString(), nullable=False
        ),
        sa.Column("log", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.ForeignKeyConstraint(
            ["dataset_id"],
            ["datasetv2.id"],
        ),
        sa.ForeignKeyConstraint(
            ["project_id"],
            ["projectv2.id"],
        ),
        sa.ForeignKeyConstraint(
            ["workflow_id"],
            ["workflowv2.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_table(
        "workflowtaskv2",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("workflow_id", sa.Integer(), nullable=False),
        sa.Column("order", sa.Integer(), nullable=True),
        sa.Column("meta_parallel", sa.JSON(), nullable=True),
        sa.Column("meta_non_parallel", sa.JSON(), nullable=True),
        sa.Column("args_parallel", sa.JSON(), nullable=True),
        sa.Column("args_non_parallel", sa.JSON(), nullable=True),
        sa.Column(
            "input_filters",
            sa.JSON(),
            server_default='{"attributes": {}, "types": {}}',
            nullable=False,
        ),
        sa.Column("is_legacy_task", sa.Boolean(), nullable=False),
        sa.Column(
            "task_type", sqlmodel.sql.sqltypes.AutoString(), nullable=False
        ),
        sa.Column("task_id", sa.Integer(), nullable=True),
        sa.Column("task_legacy_id", sa.Integer(), nullable=True),
        sa.ForeignKeyConstraint(
            ["task_id"],
            ["taskv2.id"],
        ),
        sa.ForeignKeyConstraint(
            ["task_legacy_id"],
            ["task.id"],
        ),
        sa.ForeignKeyConstraint(
            ["workflow_id"],
            ["workflowv2.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.add_column(
        "task",
        sa.Column(
            "is_v2_compatible",
            sa.Boolean(),
            server_default=sa.text("false"),
            nullable=False,
        ),
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column("task", "is_v2_compatible")
    op.drop_table("workflowtaskv2")
    op.drop_table("jobv2")
    op.drop_table("workflowv2")
    op.drop_table("linkuserprojectv2")
    op.drop_table("datasetv2")
    op.drop_table("taskv2")
    op.drop_table("projectv2")
    op.drop_table("collectionstatev2")
    # ### end Alembic commands ###
