Made a new capability to support queueing curator workflow requests

shared/workspaces/alembic/versions/5939146da7bb_add_curator_capability.py

+"""add curator capability
+
+Revision ID: 5939146da7bb
+Revises: 3eae1178cace
+Create Date: 2023-10-23 22:04:43.801226
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '5939146da7bb'
+down_revision = '3eae1178cace'
+branch_labels = None
+depends_on = None
+
+wf_name = "curator"
+
+
+def upgrade():
+    op.execute(
+        f"""
+        INSERT INTO capabilities (capability_name, max_jobs)
+        VALUES (E'{wf_name}', 10)
+        """
+    )
+
+    op.execute(
+        f"""
+        INSERT INTO capability_state_machines (capability_name, machine_type, associated_workflows)
+        VALUES (E'{wf_name}', 'simple', '{{"workflow_name": "{wf_name}"}}')
+        """
+    )
+
+    # This capability will have the same states, transitions, and actions as the download capability
+    op.execute(
+        f"""
+        INSERT INTO capability_states
+            SELECT E'{wf_name}', state
+                FROM capability_states
+                WHERE capability_name = 'download'
+        """
+    )
+
+    op.execute(
+        f"""
+        INSERT INTO capability_state_transitions (capability_name, from_state, pattern, to_state)
+            SELECT E'{wf_name}', from_state, pattern, to_state
+                FROM capability_state_transitions
+                WHERE capability_name = 'download';
+        """
+    )
+
+    # Making these explicit to grab the correct transition_ids
+    op.execute(
+        f"""
+        INSERT INTO capability_state_actions (transition_id, action_type, arguments)
+            VALUES (
+            (SELECT transition_id FROM capability_state_transitions
+                WHERE capability_name=E'{wf_name}'
+                AND from_state='Start'
+                AND to_state='Queued'),
+                'QueueWorkflow', E'{wf_name}');
+
+        INSERT INTO capability_state_actions (transition_id, action_type, arguments)
+            VALUES (
+            (SELECT transition_id FROM capability_state_transitions
+                WHERE capability_name=E'{wf_name}'
+                AND from_state='Start'
+                AND to_state='Queued'),
+                'SendNotification', '{{"template": "workflow_submitted"}}');
+    
+        INSERT INTO capability_state_actions (transition_id, action_type, arguments)
+            VALUES (
+            (SELECT transition_id FROM capability_state_transitions
+                WHERE capability_name=E'{wf_name}'
+                AND from_state='Queued'
+                AND to_state='Executing'),
+                'ExecuteWorkflow', E'{wf_name}');
+
+        INSERT INTO capability_state_actions (transition_id, action_type, arguments)
+            VALUES (
+            (SELECT transition_id FROM capability_state_transitions
+                WHERE capability_name=E'{wf_name}'
+                AND from_state='Queued'
+                AND to_state='Cancelled'),
+                'SendMessage', 'capability_cancelled');
+
+        INSERT INTO capability_state_actions (transition_id, action_type, arguments)
+            VALUES (
+            (SELECT transition_id FROM capability_state_transitions
+                WHERE capability_name=E'{wf_name}'
+                AND from_state='Executing'
+                AND to_state='Complete'),
+                'SendMessage', 'execution_complete');
+
+        INSERT INTO capability_state_actions (transition_id, action_type, arguments)
+            VALUES (
+            (SELECT transition_id FROM capability_state_transitions
+                WHERE capability_name=E'{wf_name}'
+                AND from_state='Executing'
+                AND to_state='Complete'), 
+                'SendNotification', '{{"template": "workflow_complete", "send_to_qa": false}}');
+
+        INSERT INTO capability_state_actions (transition_id, action_type, arguments)
+            VALUES (
+            (SELECT transition_id FROM capability_state_transitions
+                WHERE capability_name=E'{wf_name}'
+                AND from_state='Executing'
+                AND to_state='Error'),
+                'SendMessage', 'execution_error');
+
+        INSERT INTO capability_state_actions (transition_id, action_type, arguments)
+            VALUES (
+            (SELECT transition_id FROM capability_state_transitions
+                WHERE capability_name=E'{wf_name}'
+                AND from_state='Executing'
+                AND to_state='Error'),
+                'SendNotification', '{{"template": "workflow_failed", "send_to_qa": false}}');
+
+        INSERT INTO capability_state_actions (transition_id, action_type, arguments)
+            VALUES (
+            (SELECT transition_id FROM capability_state_transitions
+                WHERE capability_name=E'{wf_name}'
+                AND from_state='Executing'
+                AND to_state='Cancelled'),
+                'Cancel', E'{wf_name}');
+
+        INSERT INTO capability_state_actions (transition_id, action_type, arguments)
+            VALUES (
+            (SELECT transition_id FROM capability_state_transitions
+                WHERE capability_name=E'{wf_name}'
+                AND from_state='Executing'
+                AND to_state='Cancelled'),
+                'SendMessage', 'capability_cancelled');
+        """
+    )
+
+
+def downgrade():
+    op.execute(
+        f"""
+        DELETE FROM capability_state_actions WHERE transition_id IN (
+            SELECT transition_id FROM capability_state_transitions WHERE capability_name = {wf_name}
+        )
+        """
+    )
+
+    op.execute(
+        f"""
+        DELETE FROM capability_state_transitions WHERE capability_name = {wf_name}
+        """
+    )
+
+    op.execute(
+        f"""
+        DELETE FROM capability_states WHERE capability_name = {wf_name}
+        """
+    )
+
+    op.execute(
+        f"""
+        DELETE FROM capability_state_machines WHERE capability_name = {wf_name}
+        """
+    )
+
+    op.execute(
+        f"""
+        DELETE FROM capabilities WHERE capability_name = {wf_name}
+        """
+    )