Unify activity status (#513)

Author: eleftherioszisis

alembic/versions/20260120_175931_447c8883c88f_unify_activity_status.py

@@ -0,0 +1,249 @@
+"""unify_activity_status
+
+Revision ID: 447c8883c88f
+Revises: 07064e01c345
+Create Date: 2026-01-20 17:59:31.059678
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+from sqlalchemy import Text, text, inspect
+import app.db.types
+from app.db.model import Activity
+
+# revision identifiers, used by Alembic.
+revision: str = "447c8883c88f"
+down_revision: Union[str, None] = "07064e01c345"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+TABLES_WITHOUT_STATUS = [
+    "analysis_notebook_execution",
+    "calibration",
+    "circuit_extraction_config_generation",
+    "ion_channel_modeling_config_generation",
+    "simulation_generation",
+    "skeletonization_config_generation",
+    "validation",
+]
+
+# move status from these tables to the activity parent table
+TABLES_TO_MOVE_STATUS = [
+    {
+        "table": "circuit_extraction_execution",
+        "column": "status",
+        "enum": {
+            "name": "circuit_extraction_execution_status",
+            "values": ["created", "pending", "running", "done", "error"],
+        },
+    },
+    {
+        "table": "skeletonization_execution",
+        "column": "status",
+        "enum": {
+            "name": "skeletonizationexecutionstatus",
+            "values": ["created", "pending", "running", "done", "error"],
+        },
+    },
+    {
+        "table": "ion_channel_modeling_execution",
+        "column": "status",
+        "enum": {
+            "name": "ion_channel_modeling_execution_status",
+            "values": ["created", "pending", "running", "done", "error"],
+        },
+    },
+    {
+        "table": "simulation_execution",
+        "column": "status",
+        "enum": {
+            "name": "simulation_execution_status",
+            "values": ["created", "pending", "running", "done", "error", "cancelled"],
+        },
+    },
+]
+
+
+# remap table status to activity status without moving it
+TABLES_TO_DROP_STATUS = [
+    {
+        "table": "single_neuron_simulation",
+        "column": "status",
+        "enum": {
+            "name": "singleneuronsimulationstatus",
+            "values": ["started", "failure", "success"],
+            "default": "success",
+        },
+    },
+    {
+        "table": "single_neuron_synaptome_simulation",
+        "column": "status",
+        "enum": {
+            "name": "singleneuronsimulationstatus",
+            "values": ["started", "failure", "success"],
+            "default": "success",
+        },
+    },
+]
+
+
+def _activity_enum():
+    return sa.Enum(
+        "created",
+        "pending",
+        "running",
+        "done",
+        "error",
+        "cancelled",
+        name="activitystatus",
+    )
+
+
+def _create_activity_status_column(op):
+    """Create activity status collumn and fill it with done."""
+    activity_enum = _activity_enum()
+    activity_enum.create(op.get_bind())
+    op.add_column(
+        "activity",
+        sa.Column("status", activity_enum, server_default=sa.text("'done'"), nullable=False),
+    )
+    return activity_enum
+
+
+def _using_expr(col: str, mapping: dict[str, str]) -> str:
+    if not mapping:
+        return f"{col}::text::activitystatus"
+
+    cases = "\n".join(f"WHEN '{old}' THEN '{new}'" for old, new in mapping.items())
+
+    return f"""
+    (
+        CASE {col}::text
+            {cases}
+            ELSE {col}::text
+        END
+    )::activitystatus
+    """
+
+
+def _move_table_statuses(op):
+    def _move_status(op, from_table: str, from_column: str, to_table: str, mapping):
+        cases = _using_expr(f"{from_table}.{from_column}", mapping)
+        op.execute(f"""
+        UPDATE {to_table}
+        SET status = {cases}
+        FROM {from_table}
+        WHERE {to_table}.id = {from_table}.id
+        """)
+        op.drop_column(from_table, "status")
+
+    for t in TABLES_TO_MOVE_STATUS:
+        _move_status(op, t["table"], t["column"], "activity", {})
+        postgresql.ENUM(*t["enum"]["values"], name=t["enum"]["name"]).drop(op.get_bind())
+
+
+def _drop_table_statuses(op):
+    for t in TABLES_TO_DROP_STATUS:
+        op.drop_column(t["table"], t["column"])
+    for t in TABLES_TO_DROP_STATUS:
+        enum = postgresql.ENUM(*t["enum"]["values"], name=t["enum"]["name"])
+        enum.drop(op.get_bind(), checkfirst=True)
+
+
+def upgrade() -> None:
+    activity_enum = _create_activity_status_column(op)
+    _move_table_statuses(op)
+    _drop_table_statuses(op)
+
+    conn = op.get_bind()
+
+    # for the activity tables that did not have a status a "done" default
+    # should be set in the acticity table
+    for table in TABLES_WITHOUT_STATUS:
+        count = conn.execute(
+            text(f"""
+            SELECT COUNT(*)
+            FROM {table} t
+            JOIN activity a ON a.id = t.id
+            WHERE a.status != 'done'
+        """)
+        ).scalar()
+
+        if count:
+            raise RuntimeError(f"{table}: {count} rows have activity.status != 'done'")
+
+
+def _create_table_enums(conn):
+    """Create all old enums."""
+    enums = {}
+    for t in TABLES_TO_MOVE_STATUS:
+        table_enum = postgresql.ENUM(*t["enum"]["values"], name=t["enum"]["name"])
+        table_enum.create(conn, checkfirst=True)
+        enums[t["enum"]["name"]] = table_enum
+    for t in TABLES_TO_DROP_STATUS:
+        table_enum = postgresql.ENUM(*t["enum"]["values"], name=t["enum"]["name"])
+        table_enum.create(conn, checkfirst=True)
+        enums[t["enum"]["name"]] = table_enum
+    return enums
+
+
+def _create_table_columns(enums):
+    for t in TABLES_TO_MOVE_STATUS:
+        op.add_column(
+            t["table"],
+            sa.Column(t["column"], enums[t["enum"]["name"]], nullable=True),
+        )
+    for t in TABLES_TO_DROP_STATUS:
+        op.add_column(
+            t["table"],
+            sa.Column(
+                t["column"],
+                enums[t["enum"]["name"]],
+                nullable=False,
+                server_default=t["enum"]["default"],
+            ),
+        )
+
+
+def _move_status_from_activity_to_tables(conn):
+    for t in TABLES_TO_MOVE_STATUS:
+        table_name = t["table"]
+        enum_name = t["enum"]["name"]
+        conn.execute(
+            text(f"""
+            UPDATE {table_name} t
+            SET status = a.status::text::{enum_name}
+            FROM activity a
+            WHERE a.id = t.id
+        """)
+        )
+
+
+def downgrade() -> None:
+    conn = op.get_bind()
+
+    enums = _create_table_enums(conn)
+
+    # create nullable status columns for tables
+    _create_table_columns(enums)
+
+    # copy values from activity to tables
+    _move_status_from_activity_to_tables(conn)
+
+    # remove nullability
+    for t in TABLES_TO_MOVE_STATUS:
+        op.alter_column(
+            t["table"],
+            t["column"],
+            existing_type=enums[t["enum"]["name"]],
+            nullable=False,
+        )
+
+    op.drop_column("activity", "status")
+    _activity_enum().drop(conn)

app/db/model.py

@@ -36,6 +36,7 @@
     BIGINT,
     JSON_DICT,
     STRING_LIST,
+    ActivityStatus,
     ActivityType,
     AgentType,
     AgePeriod,
@@ -47,7 +48,6 @@
     CellMorphologyGenerationType,
     CellMorphologyProtocolDesign,
     CircuitBuildCategory,
-    CircuitExtractionExecutionStatus,
     CircuitScale,
     ContentType,
     DerivationType,
@@ -61,17 +61,13 @@
     ExecutorType,
     ExternalSource,
     GlobalType,
-    IonChannelModelingExecutionStatus,
     MeasurementStatistic,
     MeasurementUnit,
     PointLocation,
     PointLocationType,
     PublicationType,
     RepairPipelineType,
     Sex,
-    SimulationExecutionStatus,
-    SingleNeuronSimulationStatus,
-    SkeletonizationExecutionStatus,
     SlicingDirectionType,
     StainingType,
     StorageType,
@@ -424,7 +420,7 @@ class Activity(Identifiable):
     authorized_project_id: Mapped[uuid.UUID]
     authorized_public: Mapped[bool] = mapped_column(default=False)
     type: Mapped[ActivityType]
-
+    status: Mapped[ActivityStatus]
     start_time: Mapped[datetime | None]
     end_time: Mapped[datetime | None]
 
@@ -1108,7 +1104,6 @@ class SingleNeuronSimulation(LocationMixin, NameDescriptionVectorMixin, Entity):
     seed: Mapped[int]
     injection_location: Mapped[STRING_LIST] = mapped_column(default=[])
     recording_location: Mapped[STRING_LIST] = mapped_column(default=[])
-    status: Mapped[SingleNeuronSimulationStatus]
     # TODO: called used ?
     me_model_id: Mapped[uuid.UUID] = mapped_column(
         ForeignKey(f"{EntityType.memodel}.id"), index=True
@@ -1123,7 +1118,6 @@ class SingleNeuronSynaptomeSimulation(LocationMixin, NameDescriptionVectorMixin,
     seed: Mapped[int]
     injection_location: Mapped[STRING_LIST] = mapped_column(default=[])
     recording_location: Mapped[STRING_LIST] = mapped_column(default=[])
-    status: Mapped[SingleNeuronSimulationStatus]
     synaptome_id: Mapped[uuid.UUID] = mapped_column(
         ForeignKey(f"{EntityType.single_neuron_synaptome}.id"), index=True
     )
@@ -1347,16 +1341,10 @@ class IonChannelModelingExecution(Activity, ExecutionActivityMixin):
     Attributes:
         id (uuid.UUID): Primary key for the an ion channel modeling execution,
             referencing the ion channel recording IDs.
-        status (IonChannelModelingExecutionStatus): The status of the
-            ion channel modeling execution.
     """
 
     __tablename__ = ActivityType.ion_channel_modeling_execution.value
     id: Mapped[uuid.UUID] = mapped_column(ForeignKey("activity.id"), primary_key=True)
-    status: Mapped[IonChannelModelingExecutionStatus] = mapped_column(
-        Enum(IonChannelModelingExecutionStatus, name="ion_channel_modeling_execution_status"),
-        default=IonChannelModelingExecutionStatus.created,
-    )
 
     __mapper_args__ = {"polymorphic_identity": __tablename__}  # noqa: RUF012
 
@@ -1421,7 +1409,7 @@ class Asset(Identifiable):
     """Asset table."""
 
     __tablename__ = "asset"
-    status: Mapped[AssetStatus] = mapped_column()
+    status: Mapped[AssetStatus] = mapped_column()  # TODO: Remove if postgresql_where below removed
     path: Mapped[str]  # relative path
     full_path: Mapped[str]  # full path on S3
     is_directory: Mapped[bool]
@@ -1586,15 +1574,10 @@ class SimulationExecution(Activity, ExecutionActivityMixin):
 
     Attributes:
         id: Primary key for the simulation execution, referencing the entity ID.
-        status: The status of the simulation execution.
     """
 
     __tablename__ = ActivityType.simulation_execution.value
     id: Mapped[uuid.UUID] = mapped_column(ForeignKey("activity.id"), primary_key=True)
-    status: Mapped[SimulationExecutionStatus] = mapped_column(
-        Enum(SimulationExecutionStatus, name="simulation_execution_status"),
-        default=SimulationExecutionStatus.created,
-    )
 
     __mapper_args__ = {"polymorphic_identity": __tablename__}  # noqa: RUF012
 
@@ -1923,18 +1906,13 @@ class CircuitExtractionExecution(Activity, ExecutionActivityMixin):
 
     Attributes:
         id (uuid.UUID): Primary key.
-        status (CircuitExtractionExecutionStatus): The status of the circuit extraction execution.
 
     Note: The CircuitExtractionExecution activity associates a CircuitExtractionConfig entity with
           its corresponding extracted output Circuit entity.
     """
 
     __tablename__ = ActivityType.circuit_extraction_execution.value
     id: Mapped[uuid.UUID] = mapped_column(ForeignKey("activity.id"), primary_key=True)
-    status: Mapped[CircuitExtractionExecutionStatus] = mapped_column(
-        Enum(CircuitExtractionExecutionStatus, name="circuit_extraction_execution_status"),
-        default=CircuitExtractionExecutionStatus.created,
-    )
 
     __mapper_args__ = {"polymorphic_identity": __tablename__}  # noqa: RUF012
 
@@ -2264,9 +2242,6 @@ class SkeletonizationExecution(Activity, ExecutionActivityMixin):
     __tablename__ = ActivityType.skeletonization_execution.value
 
     id: Mapped[uuid.UUID] = mapped_column(ForeignKey("activity.id"), primary_key=True)
-    status: Mapped[SkeletonizationExecutionStatus] = mapped_column(
-        default=SkeletonizationExecutionStatus.created,
-    )
 
     __mapper_args__ = {"polymorphic_identity": __tablename__}  # noqa: RUF012