Add generic task_config and task_activity

alembic/versions/20260304_174059_38e60731b8fc_add_generic_task_models.py

@@ -0,0 +1,370 @@
+"""Add generic task models
+
+Revision ID: 38e60731b8fc
+Revises: 523e523531a7
+Create Date: 2026-03-04 17:40:59.423728
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from alembic_postgresql_enum import TableReference
+from sqlalchemy.dialects import postgresql
+
+from sqlalchemy import Text
+import app.db.types
+
+# revision identifiers, used by Alembic.
+revision: str = "38e60731b8fc"
+down_revision: Union[str, None] = "523e523531a7"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    sa.Enum(
+        "circuit_simulation__config_generation",
+        "circuit_simulation__execution",
+        "circuit_extraction__config_generation",
+        "circuit_extraction__execution",
+        "ion_channel_modeling__config_generation",
+        "ion_channel_modeling__execution",
+        "skeletonization__config_generation",
+        "skeletonization__execution",
+        "ion_channel_simulation__config_generation",
+        "ion_channel_simulation__execution",
+        "em_synapse_mapping__config_generation",
+        "em_synapse_mapping__execution",
+        name="taskactivitytype",
+    ).create(op.get_bind())
+    sa.Enum(
+        "circuit_simulation__campaign",
+        "circuit_simulation__config",
+        "circuit_extraction__campaign",
+        "circuit_extraction__config",
+        "ion_channel_modeling__campaign",
+        "ion_channel_modeling__config",
+        "skeletonization__campaign",
+        "skeletonization__config",
+        "ion_channel_simulation__campaign",
+        "ion_channel_simulation__config",
+        "em_synapse_mapping__campaign",
+        "em_synapse_mapping__config",
+        name="taskconfigtype",
+    ).create(op.get_bind())
+    op.create_table(
+        "task_activity",
+        sa.Column("id", sa.Uuid(), nullable=False),
+        sa.Column(
+            "task_activity_type",
+            postgresql.ENUM(
+                "circuit_simulation__config_generation",
+                "circuit_simulation__execution",
+                "circuit_extraction__config_generation",
+                "circuit_extraction__execution",
+                "ion_channel_modeling__config_generation",
+                "ion_channel_modeling__execution",
+                "skeletonization__config_generation",
+                "skeletonization__execution",
+                "ion_channel_simulation__config_generation",
+                "ion_channel_simulation__execution",
+                "em_synapse_mapping__config_generation",
+                "em_synapse_mapping__execution",
+                name="taskactivitytype",
+                create_type=False,
+            ),
+            nullable=False,
+        ),
+        sa.Column(
+            "executor",
+            postgresql.ENUM(
+                "single_node_job",
+                "distributed_job",
+                "jupyter_notebook",
+                name="executortype",
+                create_type=False,
+            ),
+            nullable=True,
+        ),
+        sa.Column("execution_id", sa.Uuid(), nullable=True),
+        sa.ForeignKeyConstraint(["id"], ["activity.id"], name=op.f("fk_task_activity_id_activity")),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_task_activity")),
+    )
+    op.create_index(
+        op.f("ix_task_activity_task_activity_type"),
+        "task_activity",
+        ["task_activity_type"],
+        unique=False,
+    )
+    op.create_table(
+        "task_config",
+        sa.Column("id", sa.Uuid(), nullable=False),
+        sa.Column(
+            "task_config_type",
+            postgresql.ENUM(
+                "circuit_simulation__campaign",
+                "circuit_simulation__config",
+                "circuit_extraction__campaign",
+                "circuit_extraction__config",
+                "ion_channel_modeling__campaign",
+                "ion_channel_modeling__config",
+                "skeletonization__campaign",
+                "skeletonization__config",
+                "ion_channel_simulation__campaign",
+                "ion_channel_simulation__config",
+                "em_synapse_mapping__campaign",
+                "em_synapse_mapping__config",
+                name="taskconfigtype",
+                create_type=False,
+            ),
+            nullable=False,
+        ),
+        sa.Column(
+            "meta", postgresql.JSONB(astext_type=sa.Text()), server_default="{}", nullable=False
+        ),
+        sa.Column("task_config_generator_id", sa.Uuid(), nullable=True),
+        sa.Column("name", sa.String(), nullable=False),
+        sa.Column("description", sa.String(), nullable=False),
+        sa.Column("description_vector", postgresql.TSVECTOR(), nullable=True),
+        sa.ForeignKeyConstraint(["id"], ["entity.id"], name=op.f("fk_task_config_id_entity")),
+        sa.ForeignKeyConstraint(
+            ["task_config_generator_id"],
+            ["task_config.id"],
+            name=op.f("fk_task_config_task_config_generator_id_task_config"),
+        ),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_task_config")),
+    )
+    op.create_index(
+        "ix_task_config_description_vector",
+        "task_config",
+        ["description_vector"],
+        unique=False,
+        postgresql_using="gin",
+    )
+    op.create_index(op.f("ix_task_config_name"), "task_config", ["name"], unique=False)
+    op.create_index(
+        op.f("ix_task_config_task_config_generator_id"),
+        "task_config",
+        ["task_config_generator_id"],
+        unique=False,
+    )
+    op.create_index(
+        op.f("ix_task_config_task_config_type"), "task_config", ["task_config_type"], unique=False
+    )
+    op.create_table(
+        "task_config__entity",
+        sa.Column("task_config_id", sa.Uuid(), nullable=False),
+        sa.Column("entity_id", sa.Uuid(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["entity_id"],
+            ["entity.id"],
+            name=op.f("fk_task_config__entity_entity_id_entity"),
+            ondelete="CASCADE",
+        ),
+        sa.ForeignKeyConstraint(
+            ["task_config_id"],
+            ["task_config.id"],
+            name=op.f("fk_task_config__entity_task_config_id_task_config"),
+            ondelete="CASCADE",
+        ),
+        sa.PrimaryKeyConstraint("task_config_id", "entity_id", name=op.f("pk_task_config__entity")),
+    )
+    op.sync_enum_values(
+        enum_schema="public",
+        enum_name="activitytype",
+        new_values=[
+            "simulation_execution",
+            "simulation_generation",
+            "validation",
+            "calibration",
+            "analysis_notebook_execution",
+            "ion_channel_modeling_execution",
+            "ion_channel_modeling_config_generation",
+            "circuit_extraction_config_generation",
+            "circuit_extraction_execution",
+            "skeletonization_execution",
+            "skeletonization_config_generation",
+            "task_activity",
+        ],
+        affected_columns=[
+            TableReference(table_schema="public", table_name="activity", column_name="type")
+        ],
+        enum_values_to_rename=[],
+    )
+    op.sync_enum_values(
+        enum_schema="public",
+        enum_name="entitytype",
+        new_values=[
+            "analysis_software_source_code",
+            "brain_atlas",
+            "brain_atlas_region",
+            "cell_composition",
+            "cell_morphology",
+            "cell_morphology_protocol",
+            "electrical_cell_recording",
+            "electrical_recording",
+            "electrical_recording_stimulus",
+            "emodel",
+            "experimental_bouton_density",
+            "experimental_neuron_density",
+            "experimental_synapses_per_connection",
+            "external_url",
+            "ion_channel_model",
+            "ion_channel_modeling_campaign",
+            "ion_channel_modeling_config",
+            "ion_channel_recording",
+            "memodel",
+            "memodel_calibration_result",
+            "me_type_density",
+            "simulation",
+            "simulation_campaign",
+            "simulation_result",
+            "scientific_artifact",
+            "single_neuron_simulation",
+            "single_neuron_synaptome",
+            "single_neuron_synaptome_simulation",
+            "subject",
+            "validation_result",
+            "circuit",
+            "circuit_extraction_campaign",
+            "circuit_extraction_config",
+            "em_dense_reconstruction_dataset",
+            "em_cell_mesh",
+            "analysis_notebook_template",
+            "analysis_notebook_environment",
+            "analysis_notebook_result",
+            "skeletonization_config",
+            "skeletonization_campaign",
+            "task_config",
+        ],
+        affected_columns=[
+            TableReference(table_schema="public", table_name="entity", column_name="type"),
+            TableReference(
+                table_schema="public", table_name="measurement_label", column_name="entity_type"
+            ),
+        ],
+        enum_values_to_rename=[],
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.sync_enum_values(
+        enum_schema="public",
+        enum_name="entitytype",
+        new_values=[
+            "analysis_software_source_code",
+            "brain_atlas",
+            "brain_atlas_region",
+            "cell_composition",
+            "cell_morphology",
+            "cell_morphology_protocol",
+            "electrical_cell_recording",
+            "electrical_recording",
+            "electrical_recording_stimulus",
+            "emodel",
+            "experimental_bouton_density",
+            "experimental_neuron_density",
+            "experimental_synapses_per_connection",
+            "external_url",
+            "ion_channel_model",
+            "ion_channel_modeling_campaign",
+            "ion_channel_modeling_config",
+            "ion_channel_recording",
+            "memodel",
+            "memodel_calibration_result",
+            "me_type_density",
+            "simulation",
+            "simulation_campaign",
+            "simulation_result",
+            "scientific_artifact",
+            "single_neuron_simulation",
+            "single_neuron_synaptome",
+            "single_neuron_synaptome_simulation",
+            "subject",
+            "validation_result",
+            "circuit",
+            "circuit_extraction_campaign",
+            "circuit_extraction_config",
+            "em_dense_reconstruction_dataset",
+            "em_cell_mesh",
+            "analysis_notebook_template",
+            "analysis_notebook_environment",
+            "analysis_notebook_result",
+            "skeletonization_config",
+            "skeletonization_campaign",
+        ],
+        affected_columns=[
+            TableReference(table_schema="public", table_name="entity", column_name="type"),
+            TableReference(
+                table_schema="public", table_name="measurement_label", column_name="entity_type"
+            ),
+        ],
+        enum_values_to_rename=[],
+    )
+    op.sync_enum_values(
+        enum_schema="public",
+        enum_name="activitytype",
+        new_values=[
+            "simulation_execution",
+            "simulation_generation",
+            "validation",
+            "calibration",
+            "analysis_notebook_execution",
+            "ion_channel_modeling_execution",
+            "ion_channel_modeling_config_generation",
+            "circuit_extraction_config_generation",
+            "circuit_extraction_execution",
+            "skeletonization_execution",
+            "skeletonization_config_generation",
+        ],
+        affected_columns=[
+            TableReference(table_schema="public", table_name="activity", column_name="type")
+        ],
+        enum_values_to_rename=[],
+    )
+    op.drop_table("task_config__entity")
+    op.drop_index(op.f("ix_task_config_task_config_type"), table_name="task_config")
+    op.drop_index(op.f("ix_task_config_task_config_generator_id"), table_name="task_config")
+    op.drop_index(op.f("ix_task_config_name"), table_name="task_config")
+    op.drop_index(
+        "ix_task_config_description_vector", table_name="task_config", postgresql_using="gin"
+    )
+    op.drop_table("task_config")
+    op.drop_index(op.f("ix_task_activity_task_activity_type"), table_name="task_activity")
+    op.drop_table("task_activity")
+    sa.Enum(
+        "circuit_simulation__campaign",
+        "circuit_simulation__config",
+        "circuit_extraction__campaign",
+        "circuit_extraction__config",
+        "ion_channel_modeling__campaign",
+        "ion_channel_modeling__config",
+        "skeletonization__campaign",
+        "skeletonization__config",
+        "ion_channel_simulation__campaign",
+        "ion_channel_simulation__config",
+        "em_synapse_mapping__campaign",
+        "em_synapse_mapping__config",
+        name="taskconfigtype",
+    ).drop(op.get_bind())
+    sa.Enum(
+        "circuit_simulation__config_generation",
+        "circuit_simulation__execution",
+        "circuit_extraction__config_generation",
+        "circuit_extraction__execution",
+        "ion_channel_modeling__config_generation",
+        "ion_channel_modeling__execution",
+        "skeletonization__config_generation",
+        "skeletonization__execution",
+        "ion_channel_simulation__config_generation",
+        "ion_channel_simulation__execution",
+        "em_synapse_mapping__config_generation",
+        "em_synapse_mapping__execution",
+        name="taskactivitytype",
+    ).drop(op.get_bind())
+    # ### end Alembic commands ###