"""Create tables for v1 schema

Revision ID: 148c032ffbd8
Revises: b835738272a8
Create Date: 2024-04-08 08:08:39.717348

"""

import sqlalchemy as sa
from alembic import op

from dioptra.restapi.db.custom_types import GUID, TZDateTime

# revision identifiers, used by Alembic.
revision = "148c032ffbd8"
down_revision = "b835738272a8"
branch_labels = None
depends_on = None

# Declare the database ontology
ENTRY_POINT_PARAMETER_TYPES = [
    {"parameter_type": "string"},
    {"parameter_type": "float"},
    {"parameter_type": "path"},
    {"parameter_type": "uri"},
]
JOB_STATUS_TYPES = [
    {"status": "queued"},
    {"status": "started"},
    {"status": "deferred"},
    {"status": "finished"},
    {"status": "failed"},
]
USER_LOCK_TYPES = [
    {"user_lock_type": "delete"},
]
GROUP_LOCK_TYPES = [
    {"group_lock_type": "delete"},
]
RESOURCE_LOCK_TYPES = [
    {"resource_lock_type": "delete"},
]
RESOURCE_TYPES = [
    {"resource_type": "artifact"},
    {"resource_type": "entry_point"},
    {"resource_type": "experiment"},
    {"resource_type": "job"},
    {"resource_type": "ml_model"},
    {"resource_type": "plugin"},
    {"resource_type": "plugin_file"},
    {"resource_type": "plugin_task_parameter_type"},
    {"resource_type": "queue"},
    {"resource_type": "resource_snapshot"},
]
RESOURCE_DEPENDENCY_TYPES = [
    {"parent_resource_type": "experiment", "child_resource_type": "entry_point"},
    {"parent_resource_type": "entry_point", "child_resource_type": "plugin"},
    {"parent_resource_type": "entry_point", "child_resource_type": "queue"},
    {"parent_resource_type": "plugin", "child_resource_type": "plugin_file"},
    {"parent_resource_type": "job", "child_resource_type": "artifact"},
    {"parent_resource_type": "job", "child_resource_type": "job"},
]
LEGACY_JOB_STATUS_TYPES = [
    {"status": "queued"},
    {"status": "started"},
    {"status": "deferred"},
    {"status": "finished"},
    {"status": "failed"},
]


# Helper function for inserting the ontology definitions
def _insert_ontology_data(
    entry_point_parameter_types_table,
    job_status_types_table,
    user_lock_types_table,
    group_lock_types_table,
    resource_lock_types_table,
    resource_types_table,
    resource_dependency_types_table,
) -> None:
    table_data_mappings = [
        (entry_point_parameter_types_table, ENTRY_POINT_PARAMETER_TYPES),
        (job_status_types_table, JOB_STATUS_TYPES),
        (user_lock_types_table, USER_LOCK_TYPES),
        (group_lock_types_table, GROUP_LOCK_TYPES),
        (resource_lock_types_table, RESOURCE_LOCK_TYPES),
        (resource_types_table, RESOURCE_TYPES),
        (resource_dependency_types_table, RESOURCE_DEPENDENCY_TYPES),
    ]

    for table, data in table_data_mappings:
        op.bulk_insert(table, data)


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    entry_point_parameter_types_table = op.create_table(
        "entry_point_parameter_types",
        sa.Column("parameter_type", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint(
            "parameter_type", name=op.f("pk_entry_point_parameter_types")
        ),
    )
    group_lock_types_table = op.create_table(
        "group_lock_types",
        sa.Column("group_lock_type", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint("group_lock_type", name=op.f("pk_group_lock_types")),
    )
    op.create_table(
        "job_mlflow_runs",
        sa.Column(
            "job_resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("mlflow_run_id", GUID(), nullable=False),
        sa.PrimaryKeyConstraint("job_resource_id", name=op.f("pk_job_mlflow_runs")),
    )
    with op.batch_alter_table("job_mlflow_runs", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_job_mlflow_runs_mlflow_run_id"),
            ["mlflow_run_id"],
            unique=True,
        )

    job_status_types_table = op.create_table(
        "job_status_types",
        sa.Column("status", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint("status", name=op.f("pk_job_status_types")),
    )
    resource_lock_types_table = op.create_table(
        "resource_lock_types",
        sa.Column("resource_lock_type", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint(
            "resource_lock_type", name=op.f("pk_resource_lock_types")
        ),
    )
    resource_types_table = op.create_table(
        "resource_types",
        sa.Column("resource_type", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint("resource_type", name=op.f("pk_resource_types")),
    )
    user_lock_types_table = op.create_table(
        "user_lock_types",
        sa.Column("user_lock_type", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint("user_lock_type", name=op.f("pk_user_lock_types")),
    )
    op.create_table(
        "users",
        sa.Column(
            "user_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("alternative_id", GUID(), nullable=False),
        sa.Column("username", sa.Text(), nullable=False),
        sa.Column("password", sa.Text(), nullable=True),
        sa.Column("email_address", sa.Text(), nullable=False),
        sa.Column("created_on", TZDateTime(), nullable=False),
        sa.Column("last_modified_on", TZDateTime(), nullable=False),
        sa.Column("last_login_on", TZDateTime(), nullable=True),
        sa.Column("password_expire_on", TZDateTime(), nullable=False),
        sa.PrimaryKeyConstraint("user_id", name=op.f("pk_users")),
    )
    with op.batch_alter_table("users", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_users_alternative_id"), ["alternative_id"], unique=True
        )
        batch_op.create_index(
            batch_op.f("ix_users_email_address"), ["email_address"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_users_username"), ["username"], unique=False
        )

    op.create_table(
        "groups",
        sa.Column(
            "group_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("name", sa.Text(), nullable=False),
        sa.Column(
            "user_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("created_on", TZDateTime(), nullable=False),
        sa.Column("last_modified_on", TZDateTime(), nullable=False),
        sa.PrimaryKeyConstraint("group_id", name=op.f("pk_groups")),
    )
    with op.batch_alter_table("groups", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_groups_user_id"), ["user_id"], unique=False
        )

    resource_dependency_types_table = op.create_table(
        "resource_dependency_types",
        sa.Column("parent_resource_type", sa.Text(), nullable=False),
        sa.Column("child_resource_type", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint(
            "parent_resource_type",
            "child_resource_type",
            name=op.f("pk_resource_dependency_types"),
        ),
    )
    op.create_table(
        "user_locks",
        sa.Column(
            "user_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("user_lock_type", sa.Text(), nullable=False),
        sa.Column("created_on", TZDateTime(), nullable=False),
        sa.PrimaryKeyConstraint(
            "user_id", "user_lock_type", name=op.f("pk_user_locks")
        ),
    )
    op.create_table(
        "draft_resources",
        sa.Column(
            "draft_resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "group_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("resource_type", sa.Text(), nullable=False),
        sa.Column(
            "user_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("payload", sa.JSON(), nullable=False),
        sa.Column("created_on", TZDateTime(), nullable=False),
        sa.Column("last_modified_on", TZDateTime(), nullable=False),
        sa.PrimaryKeyConstraint("draft_resource_id", name=op.f("pk_draft_resources")),
    )
    with op.batch_alter_table("draft_resources", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_draft_resources_group_id"),
            ["group_id", "resource_type", "user_id"],
            unique=False,
        )
        batch_op.create_index(
            batch_op.f("ix_draft_resources_user_id"), ["user_id"], unique=False
        )

    op.create_table(
        "group_locks",
        sa.Column(
            "group_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("group_lock_type", sa.Text(), nullable=False),
        sa.Column("created_on", TZDateTime(), nullable=False),
        sa.PrimaryKeyConstraint(
            "group_id", "group_lock_type", name=op.f("pk_group_locks")
        ),
    )
    op.create_table(
        "group_managers",
        sa.Column(
            "user_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "group_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("owner", sa.Boolean(), nullable=False),
        sa.Column("admin", sa.Boolean(), nullable=False),
        sa.PrimaryKeyConstraint("user_id", "group_id", name=op.f("pk_group_managers")),
    )
    op.create_table(
        "group_members",
        sa.Column(
            "user_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "group_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("read", sa.Boolean(), nullable=False),
        sa.Column("write", sa.Boolean(), nullable=False),
        sa.Column("share_read", sa.Boolean(), nullable=False),
        sa.Column("share_write", sa.Boolean(), nullable=False),
        sa.PrimaryKeyConstraint("user_id", "group_id", name=op.f("pk_group_members")),
    )
    op.create_table(
        "resource_dependencies",
        sa.Column(
            "parent_resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "child_resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("parent_resource_type", sa.Text(), nullable=False),
        sa.Column("child_resource_type", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint(
            "parent_resource_id",
            "child_resource_id",
            name=op.f("pk_resource_dependencies"),
        ),
    )
    with op.batch_alter_table("resource_dependencies", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_resource_dependencies_child_resource_id"),
            ["child_resource_id", "parent_resource_id"],
            unique=False,
        )
        batch_op.create_index(
            batch_op.f("ix_resource_dependencies_parent_resource_type"),
            ["parent_resource_type", "child_resource_type"],
            unique=False,
        )

    op.create_table(
        "resources",
        sa.Column(
            "resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "group_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("resource_type", sa.Text(), nullable=False),
        sa.Column("created_on", TZDateTime(), nullable=False),
        sa.PrimaryKeyConstraint("resource_id", name=op.f("pk_resources")),
    )
    with op.batch_alter_table("resources", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_resources_group_id"),
            ["group_id", "resource_type"],
            unique=False,
        )
        batch_op.create_index(
            batch_op.f("ix_resources_resource_id"),
            ["resource_id", "resource_type"],
            unique=True,
        )
        batch_op.create_index(
            batch_op.f("ix_resources_resource_type"), ["resource_type"], unique=False
        )

    op.create_table(
        "tags",
        sa.Column(
            "tag_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "group_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "user_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("name", sa.Text(), nullable=False),
        sa.Column("created_on", TZDateTime(), nullable=False),
        sa.Column("last_modified_on", TZDateTime(), nullable=False),
        sa.PrimaryKeyConstraint("tag_id", name=op.f("pk_tags")),
    )
    with op.batch_alter_table("tags", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_tags_group_id"), ["group_id", "name"], unique=True
        )

    op.create_table(
        "resource_locks",
        sa.Column(
            "resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("resource_lock_type", sa.Text(), nullable=False),
        sa.Column("created_on", TZDateTime(), nullable=False),
        sa.PrimaryKeyConstraint(
            "resource_id", "resource_lock_type", name=op.f("pk_resource_locks")
        ),
    )
    op.create_table(
        "resource_snapshots",
        sa.Column(
            "resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("resource_type", sa.Text(), nullable=False),
        sa.Column(
            "user_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("description", sa.Text(), nullable=True),
        sa.Column("created_on", TZDateTime(), nullable=False),
        sa.PrimaryKeyConstraint(
            "resource_snapshot_id", name=op.f("pk_resource_snapshots")
        ),
        sa.UniqueConstraint(
            "resource_snapshot_id",
            "resource_id",
            name=op.f("uq_resource_snapshots_resource_snapshot_id"),
        ),
    )
    with op.batch_alter_table("resource_snapshots", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_resource_snapshots_resource_id"),
            ["resource_id", "created_on"],
            unique=False,
        )
        batch_op.create_index(
            batch_op.f("ix_resource_snapshots_resource_snapshot_id"),
            ["resource_snapshot_id", "resource_id", "resource_type"],
            unique=True,
        )
        batch_op.create_index(
            batch_op.f("ix_resource_snapshots_resource_type"),
            ["resource_type"],
            unique=False,
        )
        batch_op.create_index(
            batch_op.f("ix_resource_snapshots_description"),
            ["description"],
            unique=False,
        )

    op.create_table(
        "resource_tags",
        sa.Column(
            "resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "tag_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint("resource_id", "tag_id", name=op.f("pk_resource_tags")),
    )
    op.create_table(
        "shared_resources",
        sa.Column(
            "shared_resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "group_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "user_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("read", sa.Boolean(), nullable=False),
        sa.Column("write", sa.Boolean(), nullable=False),
        sa.Column("created_on", TZDateTime(), nullable=False),
        sa.PrimaryKeyConstraint("shared_resource_id", name=op.f("pk_shared_resources")),
    )
    with op.batch_alter_table("shared_resources", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_shared_resources_group_id"),
            ["group_id", "resource_id"],
            unique=False,
        )
        batch_op.create_index(
            batch_op.f("ix_shared_resources_resource_id"), ["resource_id"], unique=False
        )

    op.create_table(
        "artifacts",
        sa.Column(
            "resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("uri", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint("resource_snapshot_id", name=op.f("pk_artifacts")),
    )
    with op.batch_alter_table("artifacts", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_artifacts_resource_id"), ["resource_id"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_artifacts_resource_snapshot_id"),
            ["resource_snapshot_id", "resource_id"],
            unique=True,
        )

    op.create_table(
        "entry_points",
        sa.Column(
            "resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("name", sa.Text(), nullable=False),
        sa.Column("task_graph", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint("resource_snapshot_id", name=op.f("pk_entry_points")),
    )
    with op.batch_alter_table("entry_points", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_entry_points_name"), ["name"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_entry_points_resource_id"), ["resource_id"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_entry_points_resource_snapshot_id"),
            ["resource_snapshot_id", "resource_id"],
            unique=True,
        )

    op.create_table(
        "experiments",
        sa.Column(
            "resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("name", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint("resource_snapshot_id", name=op.f("pk_experiments")),
    )
    with op.batch_alter_table("experiments", schema=None) as batch_op:
        batch_op.create_index(batch_op.f("ix_experiments_name"), ["name"], unique=False)
        batch_op.create_index(
            batch_op.f("ix_experiments_resource_id"), ["resource_id"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_experiments_resource_snapshot_id"),
            ["resource_snapshot_id", "resource_id"],
            unique=True,
        )

    op.create_table(
        "jobs",
        sa.Column(
            "resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("timeout", sa.Text(), nullable=False),
        sa.Column("status", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint("resource_snapshot_id", name=op.f("pk_jobs")),
    )
    with op.batch_alter_table("jobs", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_jobs_resource_id"), ["resource_id"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_jobs_resource_snapshot_id"),
            ["resource_snapshot_id", "resource_id"],
            unique=True,
        )
        batch_op.create_index(batch_op.f("ix_jobs_status"), ["status"], unique=False)

    op.create_table(
        "ml_models",
        sa.Column(
            "resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "artifact_resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("name", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint("resource_snapshot_id", name=op.f("pk_ml_models")),
        sa.UniqueConstraint(
            "resource_snapshot_id",
            "resource_id",
            name=op.f("uq_ml_models_resource_snapshot_id"),
        ),
    )
    with op.batch_alter_table("ml_models", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_ml_models_resource_id"), ["resource_id"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_ml_models_resource_snapshot_id"),
            ["resource_snapshot_id", "resource_id", "artifact_resource_snapshot_id"],
            unique=True,
        )

    op.create_table(
        "plugin_files",
        sa.Column(
            "resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("filename", sa.Text(), nullable=False),
        sa.Column("contents", sa.Text(), nullable=True),
        sa.PrimaryKeyConstraint("resource_snapshot_id", name=op.f("pk_plugin_files")),
    )
    with op.batch_alter_table("plugin_files", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_plugin_files_filename"), ["filename"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_plugin_files_resource_id"), ["resource_id"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_plugin_files_resource_snapshot_id"),
            ["resource_snapshot_id", "resource_id"],
            unique=True,
        )

    op.create_table(
        "plugin_task_parameter_types",
        sa.Column(
            "resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("name", sa.Text(), nullable=False),
        sa.Column("structure", sa.JSON(), nullable=True),
        sa.PrimaryKeyConstraint(
            "resource_snapshot_id", name=op.f("pk_plugin_task_parameter_types")
        ),
    )
    with op.batch_alter_table("plugin_task_parameter_types", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_plugin_task_parameter_types_name"), ["name"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_plugin_task_parameter_types_resource_id"),
            ["resource_id"],
            unique=False,
        )
        batch_op.create_index(
            batch_op.f("ix_plugin_task_parameter_types_resource_snapshot_id"),
            ["resource_snapshot_id", "resource_id"],
            unique=True,
        )

    op.create_table(
        "plugins",
        sa.Column(
            "resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("name", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint("resource_snapshot_id", name=op.f("pk_plugins")),
    )
    with op.batch_alter_table("plugins", schema=None) as batch_op:
        batch_op.create_index(batch_op.f("ix_plugins_name"), ["name"], unique=False)
        batch_op.create_index(
            batch_op.f("ix_plugins_resource_id"), ["resource_id"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_plugins_resource_snapshot_id"),
            ["resource_snapshot_id", "resource_id"],
            unique=True,
        )

    op.create_table(
        "queues",
        sa.Column(
            "resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("name", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint("resource_snapshot_id", name=op.f("pk_queues")),
    )
    with op.batch_alter_table("queues", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_queues_resource_id"), ["resource_id"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_queues_resource_snapshot_id"),
            ["resource_snapshot_id", "resource_id"],
            unique=True,
        )

    op.create_table(
        "shared_resource_tags",
        sa.Column(
            "shared_resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "tag_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint(
            "shared_resource_id", "tag_id", name=op.f("pk_shared_resource_tags")
        ),
    )
    op.create_table(
        "entry_point_jobs",
        sa.Column(
            "entry_point_resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "job_resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint(
            "entry_point_resource_snapshot_id",
            "job_resource_id",
            name=op.f("pk_entry_point_jobs"),
        ),
    )
    with op.batch_alter_table("entry_point_jobs", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_entry_point_jobs_job_resource_id"),
            ["job_resource_id"],
            unique=True,
        )

    op.create_table(
        "entry_point_parameters",
        sa.Column(
            "entry_point_resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "parameter_number",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("parameter_type", sa.Text(), nullable=False),
        sa.Column("name", sa.Text(), nullable=False),
        sa.Column("default_value", sa.Text(), nullable=True),
        sa.ForeignKeyConstraint(
            ["entry_point_resource_snapshot_id"],
            ["entry_points.resource_snapshot_id"],
            name=op.f(
                "fk_entry_point_parameters_entry_point_resource_snapshot_id_entry_points"
            ),
        ),
        sa.ForeignKeyConstraint(
            ["parameter_type"],
            ["entry_point_parameter_types.parameter_type"],
            name=op.f(
                "fk_entry_point_parameters_parameter_type_entry_point_parameter_types"
            ),
        ),
        sa.PrimaryKeyConstraint(
            "entry_point_resource_snapshot_id",
            "parameter_number",
            name=op.f("pk_entry_point_parameters"),
        ),
    )
    with op.batch_alter_table("entry_point_parameters", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_entry_point_parameters_entry_point_resource_snapshot_id"),
            ["entry_point_resource_snapshot_id", "name"],
            unique=True,
        )
        batch_op.create_index(
            batch_op.f("ix_entry_point_parameters_parameter_type"),
            ["parameter_type"],
            unique=False,
        )

    op.create_table(
        "entry_point_plugin_files",
        sa.Column(
            "entry_point_resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "plugin_resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "plugin_file_resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint(
            "entry_point_resource_snapshot_id",
            "plugin_resource_snapshot_id",
            "plugin_file_resource_snapshot_id",
            name=op.f("pk_entry_point_plugin_files"),
        ),
    )
    op.create_table(
        "experiment_jobs",
        sa.Column(
            "experiment_resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "job_resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint(
            "experiment_resource_snapshot_id",
            "job_resource_id",
            name=op.f("pk_experiment_jobs"),
        ),
    )
    with op.batch_alter_table("experiment_jobs", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_experiment_jobs_job_resource_id"),
            ["job_resource_id"],
            unique=True,
        )

    op.create_table(
        "plugin_tasks",
        sa.Column(
            "plugin_file_resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("plugin_task_name", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint(
            "plugin_file_resource_snapshot_id",
            "plugin_task_name",
            name=op.f("pk_plugin_tasks"),
        ),
    )
    op.create_table(
        "queue_jobs",
        sa.Column(
            "queue_resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "job_resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.PrimaryKeyConstraint(
            "queue_resource_snapshot_id", "job_resource_id", name=op.f("pk_queue_jobs")
        ),
    )
    with op.batch_alter_table("queue_jobs", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_queue_jobs_job_resource_id"),
            ["job_resource_id"],
            unique=True,
        )

    op.create_table(
        "entry_point_parameter_values",
        sa.Column(
            "job_resource_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "entry_point_resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "parameter_number",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("value", sa.Text(), nullable=True),
        sa.PrimaryKeyConstraint(
            "job_resource_id",
            "entry_point_resource_snapshot_id",
            "parameter_number",
            name=op.f("pk_entry_point_parameter_values"),
        ),
    )
    with op.batch_alter_table("entry_point_parameter_values", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_entry_point_parameter_values_job_resource_id"),
            ["job_resource_id", "parameter_number"],
            unique=True,
        )

    op.create_table(
        "plugin_task_input_parameters",
        sa.Column(
            "plugin_file_resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("plugin_task_name", sa.Text(), nullable=False),
        sa.Column(
            "parameter_number",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "plugin_task_parameter_type_resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("name", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint(
            "plugin_file_resource_snapshot_id",
            "plugin_task_name",
            "parameter_number",
            name=op.f("pk_plugin_task_input_parameters"),
        ),
    )
    with op.batch_alter_table("plugin_task_input_parameters", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f(
                "ix_plugin_task_input_parameters_plugin_file_resource_snapshot_id"
            ),
            ["plugin_file_resource_snapshot_id", "plugin_task_name"],
            unique=True,
        )
        batch_op.create_index(
            batch_op.f(
                "ix_plugin_task_input_parameters_plugin_task_parameter_type_resource_snapshot_id"
            ),
            ["plugin_task_parameter_type_resource_snapshot_id"],
            unique=False,
        )

    op.create_table(
        "plugin_task_output_parameters",
        sa.Column(
            "plugin_file_resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("plugin_task_name", sa.Text(), nullable=False),
        sa.Column(
            "parameter_number",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column(
            "plugin_task_parameter_type_resource_snapshot_id",
            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
            nullable=False,
        ),
        sa.Column("name", sa.Text(), nullable=False),
        sa.PrimaryKeyConstraint(
            "plugin_file_resource_snapshot_id",
            "plugin_task_name",
            "parameter_number",
            name=op.f("pk_plugin_task_output_parameters"),
        ),
    )
    with op.batch_alter_table("plugin_task_output_parameters", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f(
                "ix_plugin_task_output_parameters_plugin_file_resource_snapshot_id"
            ),
            ["plugin_file_resource_snapshot_id", "plugin_task_name"],
            unique=True,
        )
        batch_op.create_index(
            batch_op.f(
                "ix_plugin_task_output_parameters_plugin_task_parameter_type_resource_snapshot_id"
            ),
            ["plugin_task_parameter_type_resource_snapshot_id"],
            unique=False,
        )

    # Apply foreign keys after creating tables and indexes
    with op.batch_alter_table("groups", schema=None) as batch_op:
        batch_op.create_foreign_key(
            "fk_groups_user_id_users", "users", ["user_id"], ["user_id"]
        )

    with op.batch_alter_table("resource_dependency_types", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f(
                "fk_resource_dependency_types_parent_resource_type_resource_types"
            ),
            "resource_types",
            ["parent_resource_type"],
            ["resource_type"],
        )
        batch_op.create_foreign_key(
            batch_op.f(
                "fk_resource_dependency_types_child_resource_type_resource_types"
            ),
            "resource_types",
            ["child_resource_type"],
            ["resource_type"],
        )

    with op.batch_alter_table("user_locks", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_user_locks_user_id_users"), "users", ["user_id"], ["user_id"]
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_user_locks_user_lock_type_user_lock_types"),
            "user_lock_types",
            ["user_lock_type"],
            ["user_lock_type"],
        )

    with op.batch_alter_table("draft_resources", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_draft_resources_group_id_groups"),
            "groups",
            ["group_id"],
            ["group_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_draft_resources_resource_type_resource_types"),
            "resource_types",
            ["resource_type"],
            ["resource_type"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_draft_resources_user_id_users"),
            "users",
            ["user_id"],
            ["user_id"],
        )

    with op.batch_alter_table("group_locks", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_group_locks_group_id_groups"),
            "groups",
            ["group_id"],
            ["group_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_group_locks_group_lock_type_group_lock_types"),
            "group_lock_types",
            ["group_lock_type"],
            ["group_lock_type"],
        )

    with op.batch_alter_table("group_managers", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_group_managers_group_id_groups"),
            "groups",
            ["group_id"],
            ["group_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_group_managers_user_id_users"),
            "users",
            ["user_id"],
            ["user_id"],
        )

    with op.batch_alter_table("group_members", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_group_members_group_id_groups"),
            "groups",
            ["group_id"],
            ["group_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_group_members_user_id_users"),
            "users",
            ["user_id"],
            ["user_id"],
        )

    with op.batch_alter_table("resource_dependencies", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f(
                "fk_resource_dependencies_parent_resource_type_resource_dependency_types"
            ),
            "resource_dependency_types",
            ["parent_resource_type", "child_resource_type"],
            ["parent_resource_type", "child_resource_type"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_resource_dependencies_parent_resource_id_resources"),
            "resources",
            ["parent_resource_id", "parent_resource_type"],
            ["resource_id", "resource_type"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_resource_dependencies_child_resource_id_resources"),
            "resources",
            ["child_resource_id", "child_resource_type"],
            ["resource_id", "resource_type"],
        )

    with op.batch_alter_table("resources", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_resources_group_id_groups"),
            "groups",
            ["group_id"],
            ["group_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_resources_resource_type_resource_types"),
            "resource_types",
            ["resource_type"],
            ["resource_type"],
        )

    with op.batch_alter_table("tags", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_tags_group_id_groups"), "groups", ["group_id"], ["group_id"]
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_tags_user_id_users"), "users", ["user_id"], ["user_id"]
        )

    with op.batch_alter_table("resource_locks", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_resource_locks_resource_id_resources"),
            "resources",
            ["resource_id"],
            ["resource_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_resource_locks_resource_lock_type_resource_lock_types"),
            "resource_lock_types",
            ["resource_lock_type"],
            ["resource_lock_type"],
        )

    with op.batch_alter_table("resource_snapshots", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_resource_snapshots_resource_id_resources"),
            "resources",
            ["resource_id", "resource_type"],
            ["resource_id", "resource_type"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_resource_snapshots_resource_type_resource_types"),
            "resource_types",
            ["resource_type"],
            ["resource_type"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_resource_snapshots_user_id_users"),
            "users",
            ["user_id"],
            ["user_id"],
        )

    with op.batch_alter_table("resource_tags", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_resource_tags_resource_id_resources"),
            "resources",
            ["resource_id"],
            ["resource_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_resource_tags_tag_id_tags"), "tags", ["tag_id"], ["tag_id"]
        )

    with op.batch_alter_table("shared_resources", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_shared_resources_group_id_groups"),
            "groups",
            ["group_id"],
            ["group_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_shared_resources_resource_id_resources"),
            "resources",
            ["resource_id"],
            ["resource_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_shared_resources_user_id_users"),
            "users",
            ["user_id"],
            ["user_id"],
        )

    with op.batch_alter_table("artifacts", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_artifacts_resource_snapshot_id_resource_snapshots"),
            "resource_snapshots",
            ["resource_snapshot_id", "resource_id"],
            ["resource_snapshot_id", "resource_id"],
        )

    with op.batch_alter_table("entry_points", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_entry_points_resource_snapshot_id_resource_snapshots"),
            "resource_snapshots",
            ["resource_snapshot_id", "resource_id"],
            ["resource_snapshot_id", "resource_id"],
        )

    with op.batch_alter_table("experiments", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_experiments_resource_snapshot_id_resource_snapshots"),
            "resource_snapshots",
            ["resource_snapshot_id", "resource_id"],
            ["resource_snapshot_id", "resource_id"],
        )

    with op.batch_alter_table("jobs", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_jobs_resource_snapshot_id_resource_snapshots"),
            "resource_snapshots",
            ["resource_snapshot_id", "resource_id"],
            ["resource_snapshot_id", "resource_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_jobs_status_job_status_types"),
            "job_status_types",
            ["status"],
            ["status"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_jobs_resource_id_entry_point_jobs"),
            "entry_point_jobs",
            ["resource_id"],
            ["job_resource_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_jobs_resource_id_experiment_jobs"),
            "experiment_jobs",
            ["resource_id"],
            ["job_resource_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_jobs_resource_id_queue_jobs"),
            "queue_jobs",
            ["resource_id"],
            ["job_resource_id"],
        )

    with op.batch_alter_table("ml_models", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_ml_models_resource_snapshot_id_resource_snapshots"),
            "resource_snapshots",
            ["resource_snapshot_id", "resource_id"],
            ["resource_snapshot_id", "resource_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_ml_models_artifact_resource_snapshot_id_artifacts"),
            "artifacts",
            ["artifact_resource_snapshot_id"],
            ["resource_snapshot_id"],
        )

    with op.batch_alter_table("plugin_files", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_plugin_files_resource_snapshot_id_resource_snapshots"),
            "resource_snapshots",
            ["resource_snapshot_id", "resource_id"],
            ["resource_snapshot_id", "resource_id"],
        )

    with op.batch_alter_table("plugin_task_parameter_types", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f(
                "fk_plugin_task_parameter_types_resource_snapshot_id_resource_snapshots"
            ),
            "resource_snapshots",
            ["resource_snapshot_id", "resource_id"],
            ["resource_snapshot_id", "resource_id"],
        )

    with op.batch_alter_table("plugins", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_plugins_resource_snapshot_id_resource_snapshots"),
            "resource_snapshots",
            ["resource_snapshot_id", "resource_id"],
            ["resource_snapshot_id", "resource_id"],
        )

    with op.batch_alter_table("queues", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_queues_resource_snapshot_id_resource_snapshots"),
            "resource_snapshots",
            ["resource_snapshot_id", "resource_id"],
            ["resource_snapshot_id", "resource_id"],
        )

    with op.batch_alter_table("shared_resource_tags", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_shared_resource_tags_shared_resource_id_shared_resources"),
            "shared_resources",
            ["shared_resource_id"],
            ["shared_resource_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_shared_resource_tags_tag_id_tags"),
            "tags",
            ["tag_id"],
            ["tag_id"],
        )

    with op.batch_alter_table("entry_point_jobs", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f(
                "fk_entry_point_jobs_entry_point_resource_snapshot_id_entry_points"
            ),
            "entry_points",
            ["entry_point_resource_snapshot_id"],
            ["resource_snapshot_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_entry_point_jobs_job_resource_id_resources"),
            "resources",
            ["job_resource_id"],
            ["resource_id"],
        )

    with op.batch_alter_table("entry_point_plugin_files", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f(
                "fk_entry_point_plugin_files_entry_point_resource_snapshot_id_entry_points"
            ),
            "entry_points",
            ["entry_point_resource_snapshot_id"],
            ["resource_snapshot_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f(
                "fk_entry_point_plugin_files_plugin_file_resource_snapshot_id_plugin_files"
            ),
            "plugin_files",
            ["plugin_file_resource_snapshot_id"],
            ["resource_snapshot_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f(
                "fk_entry_point_plugin_files_plugin_resource_snapshot_id_plugins"
            ),
            "plugins",
            ["plugin_resource_snapshot_id"],
            ["resource_snapshot_id"],
        )

    with op.batch_alter_table("experiment_jobs", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f(
                "fk_experiment_jobs_experiment_resource_snapshot_id_experiments"
            ),
            "experiments",
            ["experiment_resource_snapshot_id"],
            ["resource_snapshot_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_experiment_jobs_job_resource_id_resources"),
            "resources",
            ["job_resource_id"],
            ["resource_id"],
        )

    with op.batch_alter_table("job_mlflow_runs", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_job_mlflow_runs_job_resource_id_resources"),
            "resources",
            ["job_resource_id"],
            ["resource_id"],
        )

    with op.batch_alter_table("plugin_tasks", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_plugin_tasks_plugin_file_resource_snapshot_id_plugin_files"),
            "plugin_files",
            ["plugin_file_resource_snapshot_id"],
            ["resource_snapshot_id"],
        )

    with op.batch_alter_table("queue_jobs", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f("fk_queue_jobs_queue_resource_snapshot_id_queues"),
            "queues",
            ["queue_resource_snapshot_id"],
            ["resource_snapshot_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_queue_jobs_job_resource_id_resources"),
            "resources",
            ["job_resource_id"],
            ["resource_id"],
        )

    with op.batch_alter_table("entry_point_parameter_values", schema=None) as batch_op:
        batch_op.create_foreign_key(
            op.f(
                "fk_entry_point_parameter_values_entry_point_resource_snapshot_id_entry_point_parameters"
            ),
            "entry_point_parameters",
            ["entry_point_resource_snapshot_id", "parameter_number"],
            ["entry_point_resource_snapshot_id", "parameter_number"],
        )
        batch_op.create_foreign_key(
            op.f(
                "fk_entry_point_parameter_values_entry_point_resource_snapshot_id_entry_point_jobs"
            ),
            "entry_point_jobs",
            ["entry_point_resource_snapshot_id", "job_resource_id"],
            ["entry_point_resource_snapshot_id", "job_resource_id"],
        )
        batch_op.create_foreign_key(
            batch_op.f("fk_entry_point_parameter_values_job_resource_id_resources"),
            "resources",
            ["job_resource_id"],
            ["resource_id"],
        )

    with op.batch_alter_table("plugin_task_input_parameters", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f(
                "fk_plugin_task_input_parameters_plugin_file_resource_snapshot_id_plugin_tasks"
            ),
            "plugin_tasks",
            ["plugin_file_resource_snapshot_id", "plugin_task_name"],
            ["plugin_file_resource_snapshot_id", "plugin_task_name"],
        )
        batch_op.create_foreign_key(
            batch_op.f(
                "fk_plugin_task_input_parameters_plugin_task_parameter_type_resource_snapshot_id_plugin_task_parameter_types"
            ),
            "plugin_task_parameter_types",
            ["plugin_task_parameter_type_resource_snapshot_id"],
            ["resource_snapshot_id"],
        )

    with op.batch_alter_table("plugin_task_output_parameters", schema=None) as batch_op:
        batch_op.create_foreign_key(
            batch_op.f(
                "fk_plugin_task_output_parameters_plugin_file_resource_snapshot_id_plugin_tasks"
            ),
            "plugin_tasks",
            ["plugin_file_resource_snapshot_id", "plugin_task_name"],
            ["plugin_file_resource_snapshot_id", "plugin_task_name"],
        )
        batch_op.create_foreign_key(
            batch_op.f(
                "fk_plugin_task_output_parameters_plugin_task_parameter_type_resource_snapshot_id_plugin_task_parameter_types"
            ),
            "plugin_task_parameter_types",
            ["plugin_task_parameter_type_resource_snapshot_id"],
            ["resource_snapshot_id"],
        )

    # Insert the ontology data
    _insert_ontology_data(
        entry_point_parameter_types_table=entry_point_parameter_types_table,
        job_status_types_table=job_status_types_table,
        user_lock_types_table=user_lock_types_table,
        group_lock_types_table=group_lock_types_table,
        resource_lock_types_table=resource_lock_types_table,
        resource_types_table=resource_types_table,
        resource_dependency_types_table=resource_dependency_types_table,
    )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("plugin_task_output_parameters", schema=None) as batch_op:
        batch_op.drop_index(
            batch_op.f(
                "ix_plugin_task_output_parameters_plugin_task_parameter_type_resource_snapshot_id"
            )
        )
        batch_op.drop_index(
            batch_op.f(
                "ix_plugin_task_output_parameters_plugin_file_resource_snapshot_id"
            )
        )

    op.drop_table("plugin_task_output_parameters")
    with op.batch_alter_table("plugin_task_input_parameters", schema=None) as batch_op:
        batch_op.drop_index(
            batch_op.f(
                "ix_plugin_task_input_parameters_plugin_task_parameter_type_resource_snapshot_id"
            )
        )
        batch_op.drop_index(
            batch_op.f(
                "ix_plugin_task_input_parameters_plugin_file_resource_snapshot_id"
            )
        )

    op.drop_table("plugin_task_input_parameters")
    with op.batch_alter_table("entry_point_parameter_values", schema=None) as batch_op:
        batch_op.drop_index(
            batch_op.f("ix_entry_point_parameter_values_job_resource_id")
        )

    op.drop_table("entry_point_parameter_values")
    op.drop_table("plugin_tasks")
    op.drop_table("entry_point_plugin_files")
    with op.batch_alter_table("entry_point_parameters", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_entry_point_parameters_parameter_type"))
        batch_op.drop_index(
            batch_op.f("ix_entry_point_parameters_entry_point_resource_snapshot_id")
        )

    op.drop_table("entry_point_parameters")
    op.drop_table("shared_resource_tags")
    with op.batch_alter_table("queues", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_queues_resource_snapshot_id"))
        batch_op.drop_index(batch_op.f("ix_queues_resource_id"))

    with op.batch_alter_table("jobs", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_jobs_status"))
        batch_op.drop_index(batch_op.f("ix_jobs_resource_snapshot_id"))
        batch_op.drop_index(batch_op.f("ix_jobs_resource_id"))

    op.drop_table("jobs")
    with op.batch_alter_table("queue_jobs", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_queue_jobs_job_resource_id"))

    op.drop_table("queue_jobs")
    with op.batch_alter_table("experiment_jobs", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_experiment_jobs_job_resource_id"))

    op.drop_table("experiment_jobs")
    with op.batch_alter_table("entry_point_jobs", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_entry_point_jobs_job_resource_id"))

    op.drop_table("entry_point_jobs")
    with op.batch_alter_table("job_mlflow_runs", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_job_mlflow_runs_mlflow_run_id"))

    op.drop_table("job_mlflow_runs")
    op.drop_table("queues")
    with op.batch_alter_table("plugins", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_plugins_resource_snapshot_id"))
        batch_op.drop_index(batch_op.f("ix_plugins_resource_id"))
        batch_op.drop_index(batch_op.f("ix_plugins_name"))

    op.drop_table("plugins")
    with op.batch_alter_table("plugin_task_parameter_types", schema=None) as batch_op:
        batch_op.drop_index(
            batch_op.f("ix_plugin_task_parameter_types_resource_snapshot_id")
        )
        batch_op.drop_index(batch_op.f("ix_plugin_task_parameter_types_resource_id"))
        batch_op.drop_index(batch_op.f("ix_plugin_task_parameter_types_name"))

    op.drop_table("plugin_task_parameter_types")
    with op.batch_alter_table("plugin_files", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_plugin_files_resource_snapshot_id"))
        batch_op.drop_index(batch_op.f("ix_plugin_files_resource_id"))
        batch_op.drop_index(batch_op.f("ix_plugin_files_filename"))

    op.drop_table("plugin_files")
    with op.batch_alter_table("ml_models", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_ml_models_resource_snapshot_id"))
        batch_op.drop_index(batch_op.f("ix_ml_models_resource_id"))

    op.drop_table("ml_models")
    with op.batch_alter_table("experiments", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_experiments_resource_snapshot_id"))
        batch_op.drop_index(batch_op.f("ix_experiments_resource_id"))
        batch_op.drop_index(batch_op.f("ix_experiments_name"))

    op.drop_table("experiments")
    with op.batch_alter_table("entry_points", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_entry_points_resource_snapshot_id"))
        batch_op.drop_index(batch_op.f("ix_entry_points_resource_id"))
        batch_op.drop_index(batch_op.f("ix_entry_points_name"))

    op.drop_table("entry_points")
    with op.batch_alter_table("artifacts", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_artifacts_resource_snapshot_id"))
        batch_op.drop_index(batch_op.f("ix_artifacts_resource_id"))

    op.drop_table("artifacts")
    with op.batch_alter_table("shared_resources", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_shared_resources_resource_id"))
        batch_op.drop_index(batch_op.f("ix_shared_resources_group_id"))

    op.drop_table("shared_resources")
    op.drop_table("resource_tags")
    with op.batch_alter_table("resource_snapshots", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_resource_snapshots_description"))
        batch_op.drop_index(batch_op.f("ix_resource_snapshots_resource_type"))
        batch_op.drop_index(batch_op.f("ix_resource_snapshots_resource_snapshot_id"))
        batch_op.drop_index(batch_op.f("ix_resource_snapshots_resource_id"))

    op.drop_table("resource_snapshots")
    op.drop_table("resource_locks")
    with op.batch_alter_table("tags", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_tags_group_id"))

    op.drop_table("tags")
    with op.batch_alter_table("resource_dependencies", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_resource_dependencies_parent_resource_type"))
        batch_op.drop_index(batch_op.f("ix_resource_dependencies_child_resource_id"))

    op.drop_table("resource_dependencies")
    with op.batch_alter_table("resources", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_resources_resource_type"))
        batch_op.drop_index(batch_op.f("ix_resources_resource_id"))
        batch_op.drop_index(batch_op.f("ix_resources_group_id"))

    op.drop_table("resources")
    op.drop_table("group_members")
    op.drop_table("group_managers")
    op.drop_table("group_locks")
    with op.batch_alter_table("draft_resources", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_draft_resources_user_id"))
        batch_op.drop_index(batch_op.f("ix_draft_resources_group_id"))

    op.drop_table("draft_resources")
    op.drop_table("user_locks")
    op.drop_table("resource_dependency_types")
    with op.batch_alter_table("groups", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_groups_user_id"))

    op.drop_table("groups")
    with op.batch_alter_table("users", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_users_username"))
        batch_op.drop_index(batch_op.f("ix_users_email_address"))
        batch_op.drop_index(batch_op.f("ix_users_alternative_id"))

    op.drop_table("users")
    op.drop_table("user_lock_types")
    op.drop_table("resource_types")
    op.drop_table("resource_lock_types")
    op.drop_table("job_status_types")
    op.drop_table("group_lock_types")
    op.drop_table("entry_point_parameter_types")
    # ### end Alembic commands ###
