♻️Maintenance: Add UNKOWN type to DB State Type (🗃️) (#8284)

Author: sanderegg

packages/postgres-database/src/simcore_postgres_database/migration/versions/06eafd25d004_add_state_type_unknown.py

@@ -0,0 +1,53 @@
+"""add state type unknown
+
+Revision ID: 06eafd25d004
+Revises: ec4f62595e0c
+Create Date: 2025-09-01 12:25:25.617790+00:00
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "06eafd25d004"
+down_revision = "ec4f62595e0c"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.execute("ALTER TYPE statetype ADD VALUE 'UNKNOWN'")
+
+
+def downgrade() -> None:
+    # NOTE: PostgreSQL doesn't support removing enum values directly
+    # This downgrades only ensure that StateType.UNKNOWN is not used
+    #
+
+    # Find all tables and columns that use statetype enum
+    result = op.get_bind().execute(
+        sa.DDL(
+            """
+        SELECT t.table_name, c.column_name, c.column_default
+        FROM information_schema.columns c
+        JOIN information_schema.tables t ON c.table_name = t.table_name
+        WHERE c.udt_name = 'statetype'
+        AND t.table_schema = 'public'
+    """
+        )
+    )
+
+    tables_columns = result.fetchall()
+
+    # Update UNKNOWN states to FAILED in all affected tables
+    for table_name, column_name, _ in tables_columns:
+        op.execute(
+            sa.DDL(
+                f"""
+            UPDATE {table_name}
+            SET {column_name} = 'FAILED'
+            WHERE {column_name} = 'UNKNOWN'
+        """
+            )
+        )

packages/postgres-database/src/simcore_postgres_database/models/comp_pipeline.py

@@ -1,6 +1,5 @@
-""" Computational Pipeline Table
+"""Computational Pipeline Table"""
 
-"""
 import enum
 import uuid
 
@@ -24,6 +23,7 @@ class StateType(enum.Enum):
     ABORTED = "ABORTED"
     WAITING_FOR_RESOURCES = "WAITING_FOR_RESOURCES"
     WAITING_FOR_CLUSTER = "WAITING_FOR_CLUSTER"
+    UNKNOWN = "UNKNOWN"
 
 
 def _new_uuid():

services/director-v2/src/simcore_service_director_v2/modules/db/repositories/comp_runs.py

@@ -95,11 +95,12 @@ def _handle_foreign_key_violation(
 
 def _resolve_grouped_state(states: list[RunningState]) -> RunningState:
     # If any state is not a final state, return STARTED
+
     final_states = {
         RunningState.FAILED,
         RunningState.ABORTED,
         RunningState.SUCCESS,
-        RunningState.UNKNOWN,
+        RunningState.UNKNOWN,  # NOTE: this is NOT a final state, but happens when tasks are missing
     }
     if any(state not in final_states for state in states):
         return RunningState.STARTED
@@ -399,7 +400,6 @@ async def list_all_collection_run_ids_for_user_currently_running_computations(
         product_name: str,
         user_id: UserID,
     ) -> list[CollectionRunID]:
-
         list_query = (
             sa.select(
                 comp_runs.c.collection_run_id,
@@ -493,17 +493,17 @@ async def list_group_by_collection_run_id(
             total_count = await conn.scalar(count_query)
             items = []
             async for row in await conn.stream(list_query):
-                db_states = [DB_TO_RUNNING_STATE[s] for s in row["states"]]
+                db_states = [DB_TO_RUNNING_STATE[s] for s in row.states]
                 resolved_state = _resolve_grouped_state(db_states)
                 items.append(
                     ComputationCollectionRunRpcGet(
-                        collection_run_id=row["collection_run_id"],
-                        project_ids=row["project_ids"],
+                        collection_run_id=row.collection_run_id,
+                        project_ids=row.project_ids,
                         state=resolved_state,
-                        info={} if row["info"] is None else row["info"],
-                        submitted_at=row["submitted_at"],
-                        started_at=row["started_at"],
-                        ended_at=row["ended_at"],
+                        info={} if row.info is None else row.info,
+                        submitted_at=row.submitted_at,
+                        started_at=row.started_at,
+                        ended_at=row.ended_at,
                     )
                 )
             return cast(int, total_count), items

services/director-v2/src/simcore_service_director_v2/utils/db.py

@@ -16,11 +16,10 @@
     StateType.ABORTED: RunningState.ABORTED,
     StateType.WAITING_FOR_RESOURCES: RunningState.WAITING_FOR_RESOURCES,
     StateType.WAITING_FOR_CLUSTER: RunningState.WAITING_FOR_CLUSTER,
+    StateType.UNKNOWN: RunningState.UNKNOWN,
 }
 
-RUNNING_STATE_TO_DB = {v: k for k, v in DB_TO_RUNNING_STATE.items()} | {
-    RunningState.UNKNOWN: StateType.FAILED
-}
+RUNNING_STATE_TO_DB = {v: k for k, v in DB_TO_RUNNING_STATE.items()}
 
 _logger = logging.getLogger(__name__)
 

services/director-v2/tests/unit/test_utils_computation.py

@@ -296,9 +296,9 @@ def test_get_pipeline_state_from_task_states(
     ]
 
     pipeline_state: RunningState = get_pipeline_state_from_task_states(tasks)
-    assert (
-        pipeline_state == exp_pipeline_state
-    ), f"task states are: {task_states}, got {pipeline_state} instead of {exp_pipeline_state}"
+    assert pipeline_state == exp_pipeline_state, (
+        f"task states are: {task_states}, got {pipeline_state} instead of {exp_pipeline_state}"
+    )
 
 
 @pytest.mark.parametrize(
@@ -315,7 +315,7 @@ def test_get_pipeline_state_from_task_states(
     ],
 )
 def test_is_pipeline_running(state, exp: bool):
-    assert (
-        is_pipeline_running(state) is exp
-    ), f"pipeline in {state}, i.e. running state should be {exp}"
+    assert is_pipeline_running(state) is exp, (
+        f"pipeline in {state}, i.e. running state should be {exp}"
+    )
     assert is_pipeline_stopped is not exp

services/director-v2/tests/unit/with_dbs/comp_scheduler/test_db_repositories_comp_runs.py

@@ -956,8 +956,9 @@ async def test_list_group_by_collection_run_id_with_unknown_returns_unknown(
     )
 
     # Test the function
+    assert "product_name" in run_metadata
     total_count, items = await repo.list_group_by_collection_run_id(
-        product_name=run_metadata.get("product_name"),
+        product_name=run_metadata["product_name"],
         user_id=published_project_1.user["id"],
         offset=0,
         limit=10,
@@ -967,7 +968,7 @@ async def test_list_group_by_collection_run_id_with_unknown_returns_unknown(
     assert total_count == 1
     assert len(items) == 1
     collection_item = items[0]
-    assert collection_item.state == RunningState.FAILED
+    assert collection_item.state == RunningState.UNKNOWN
 
 
 async def test_list_group_by_collection_run_id_with_project_filter(