Unprotoype Makros (#216)

macros

Author: JWittmeyer

alembic/versions/63364f1fe61d_macro_tables.py

@@ -0,0 +1,138 @@
+"""Macro tables
+
+Revision ID: 194838aa0431
+Revises: a14f1a9b12b0
+Create Date: 2024-06-05 11:42:56.258816
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = '194838aa0431'
+down_revision = 'a14f1a9b12b0'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('macro',
+    sa.Column('id', postgresql.UUID(as_uuid=True), nullable=False),
+    sa.Column('macro_type', sa.String(), nullable=True),
+    sa.Column('scope', sa.String(), nullable=True),
+    sa.Column('organization_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('project_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('created_by', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('created_at', sa.DateTime(), nullable=True),
+    sa.Column('state', sa.String(), nullable=True),
+    sa.Column('name', sa.String(), nullable=True),
+    sa.Column('description', sa.String(), nullable=True),
+    sa.ForeignKeyConstraint(['created_by'], ['user.id'], ondelete='SET NULL'),
+    sa.ForeignKeyConstraint(['organization_id'], ['organization.id'], ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['project_id'], ['cognition.project.id'], ondelete='CASCADE'),
+    sa.PrimaryKeyConstraint('id'),
+    schema='cognition'
+    )
+    op.create_index(op.f('ix_cognition_macro_created_by'), 'macro', ['created_by'], unique=False, schema='cognition')
+    op.create_index(op.f('ix_cognition_macro_organization_id'), 'macro', ['organization_id'], unique=False, schema='cognition')
+    op.create_index(op.f('ix_cognition_macro_project_id'), 'macro', ['project_id'], unique=False, schema='cognition')
+    op.create_index(op.f('ix_cognition_macro_scope'), 'macro', ['scope'], unique=False, schema='cognition')
+    op.create_table('macro_execution',
+    sa.Column('id', postgresql.UUID(as_uuid=True), nullable=False),
+    sa.Column('organization_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('macro_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('created_by', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('created_at', sa.DateTime(), nullable=True),
+    sa.Column('state', sa.String(), nullable=True),
+    sa.Column('execution_group_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('meta_info', sa.JSON(), nullable=True),
+    sa.ForeignKeyConstraint(['created_by'], ['user.id'], ondelete='SET NULL'),
+    sa.ForeignKeyConstraint(['organization_id'], ['organization.id'], ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['macro_id'], ['cognition.macro.id'], ondelete='CASCADE'),
+    sa.PrimaryKeyConstraint('id'),
+    schema='cognition'
+    )
+    op.create_index(op.f('ix_cognition_macro_execution_created_by'), 'macro_execution', ['created_by'], unique=False, schema='cognition')
+    op.create_index(op.f('ix_cognition_macro_execution_execution_group_id'), 'macro_execution', ['execution_group_id'], unique=False, schema='cognition')
+    op.create_index(op.f('ix_cognition_macro_execution_macro_id'), 'macro_execution', ['macro_id'], unique=False, schema='cognition')
+    op.create_table('macro_node',
+    sa.Column('id', postgresql.UUID(as_uuid=True), nullable=False),
+    sa.Column('macro_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('created_by', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('created_at', sa.DateTime(), nullable=True),
+    sa.Column('is_root', sa.Boolean(), nullable=True),
+    sa.Column('config', sa.JSON(), nullable=True),
+    sa.ForeignKeyConstraint(['created_by'], ['user.id'], ondelete='SET NULL'),
+    sa.ForeignKeyConstraint(['macro_id'], ['cognition.macro.id'], ondelete='CASCADE'),
+    sa.PrimaryKeyConstraint('id'),
+    schema='cognition'
+    )
+    op.create_index(op.f('ix_cognition_macro_node_created_by'), 'macro_node', ['created_by'], unique=False, schema='cognition')
+    op.create_index(op.f('ix_cognition_macro_node_macro_id'), 'macro_node', ['macro_id'], unique=False, schema='cognition')
+    op.create_table('macro_edge',
+    sa.Column('id', postgresql.UUID(as_uuid=True), nullable=False),
+    sa.Column('macro_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('from_node_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('to_node_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('created_by', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('created_at', sa.DateTime(), nullable=True),
+    sa.Column('config', sa.JSON(), nullable=True),
+    sa.ForeignKeyConstraint(['created_by'], ['user.id'], ondelete='SET NULL'),
+    sa.ForeignKeyConstraint(['from_node_id'], ['cognition.macro_node.id'], ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['macro_id'], ['cognition.macro.id'], ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['to_node_id'], ['cognition.macro_node.id'], ondelete='CASCADE'),
+    sa.PrimaryKeyConstraint('id'),
+    schema='cognition'
+    )
+    op.create_index(op.f('ix_cognition_macro_edge_created_by'), 'macro_edge', ['created_by'], unique=False, schema='cognition')
+    op.create_index(op.f('ix_cognition_macro_edge_from_node_id'), 'macro_edge', ['from_node_id'], unique=False, schema='cognition')
+    op.create_index(op.f('ix_cognition_macro_edge_macro_id'), 'macro_edge', ['macro_id'], unique=False, schema='cognition')
+    op.create_index(op.f('ix_cognition_macro_edge_to_node_id'), 'macro_edge', ['to_node_id'], unique=False, schema='cognition')
+    op.create_table('macro_execution_link',
+    sa.Column('id', postgresql.UUID(as_uuid=True), nullable=False),
+    sa.Column('organization_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('execution_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('execution_node_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('action', sa.String(), nullable=True),
+    sa.Column('other_id_target', sa.String(), nullable=True),
+    sa.Column('other_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.ForeignKeyConstraint(['execution_id'], ['cognition.macro_execution.id'], ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['organization_id'], ['organization.id'], ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['execution_node_id'], ['cognition.macro_node.id'], ondelete='CASCADE'),
+    sa.PrimaryKeyConstraint('id'),
+    schema='cognition'
+    )
+    op.create_index(op.f('ix_cognition_macro_execution_link_execution_id'), 'macro_execution_link', ['execution_id'], unique=False, schema='cognition')
+    op.create_index(op.f('ix_cognition_macro_execution_link_execution_node_id'), 'macro_execution_link', ['execution_node_id'], unique=False, schema='cognition')
+    op.create_index(op.f('ix_cognition_macro_execution_link_other_id'), 'macro_execution_link', ['other_id'], unique=False, schema='cognition')
+    op.add_column('project', sa.Column('macro_config', sa.JSON(), nullable=True), schema='cognition')
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('project', 'macro_config', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_execution_link_other_id'), table_name='macro_execution_link', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_execution_link_execution_node_id'), table_name='macro_execution_link', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_execution_link_execution_id'), table_name='macro_execution_link', schema='cognition')
+    op.drop_table('macro_execution_link', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_edge_to_node_id'), table_name='macro_edge', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_edge_macro_id'), table_name='macro_edge', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_edge_from_node_id'), table_name='macro_edge', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_edge_created_by'), table_name='macro_edge', schema='cognition')
+    op.drop_table('macro_edge', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_node_macro_id'), table_name='macro_node', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_node_created_by'), table_name='macro_node', schema='cognition')
+    op.drop_table('macro_node', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_execution_macro_id'), table_name='macro_execution', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_execution_execution_group_id'), table_name='macro_execution', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_execution_created_by'), table_name='macro_execution', schema='cognition')
+    op.drop_table('macro_execution', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_scope'), table_name='macro', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_project_id'), table_name='macro', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_organization_id'), table_name='macro', schema='cognition')
+    op.drop_index(op.f('ix_cognition_macro_created_by'), table_name='macro', schema='cognition')
+    op.drop_table('macro', schema='cognition')
+    # ### end Alembic commands ###

api/transfer.py

@@ -2,7 +2,7 @@
 import traceback
 import time
 from typing import Optional, Dict
-
+from time import sleep
 from starlette.endpoints import HTTPEndpoint
 from starlette.responses import PlainTextResponse, JSONResponse
 from controller.embedding.manager import recreate_embeddings
@@ -22,7 +22,10 @@
     project as refinery_project,
 )
 
-from submodules.model.cognition_objects import project as cognition_project
+from submodules.model.cognition_objects import (
+    project as cognition_project,
+    macro as macro_db_bo,
+)
 
 from controller.transfer import manager as transfer_manager
 from controller.upload_task import manager as upload_task_manager
@@ -279,6 +282,45 @@ def __add_parse_markdown_file_thread(
             general.remove_and_refresh_session(ctx_token, False)
 
 
+class CognitionStartMacroExecutionGroup(HTTPEndpoint):
+    def put(self, request) -> PlainTextResponse:
+        macro_id = request.path_params["macro_id"]
+        group_id = request.path_params["group_id"]
+
+        execution_entries = macro_db_bo.get_all_macro_executions(macro_id, group_id)
+
+        if len(execution_entries) == 0:
+            return PlainTextResponse("No executions found", status_code=400)
+        if not (cognition_prj_id := execution_entries[0].meta_info.get("project_id")):
+            return PlainTextResponse("No project id found", status_code=400)
+        cognition_prj = cognition_project.get(cognition_prj_id)
+        refinery_prj_id = str(
+            refinery_project.get_or_create_queue_project(
+                cognition_prj.organization_id, cognition_prj.created_by, True
+            ).id
+        )
+
+        def queue_tasks():
+            token = general.get_ctx_token()
+            for entry in execution_entries:
+                task_queue_manager.add_task(
+                    refinery_prj_id,
+                    TaskType.RUN_COGNITION_MACRO,
+                    entry.created_by,
+                    {
+                        "macro_id": macro_id,
+                        "execution_id": str(entry.id),
+                        "execution_group_id": group_id,
+                    },
+                )
+            general.commit()
+            general.remove_and_refresh_session(token, False)
+
+        daemon.run(queue_tasks)
+
+        return PlainTextResponse("OK")
+
+
 class AssociationsImport(HTTPEndpoint):
     async def post(self, request) -> JSONResponse:
         project_id = request.path_params["project_id"]

app.py

@@ -15,6 +15,7 @@
     CognitionImport,
     CognitionPrepareProject,
     CognitionParseMarkdownFile,
+    CognitionStartMacroExecutionGroup,
 )
 from fast_api.routes.organization import router as org_router
 from fast_api.routes.project import router as project_router
@@ -138,6 +139,10 @@
     ),
     Route("/project/{project_id:str}/import/task/{task_id:str}", UploadTaskInfo),
     Route("/project", ProjectCreationFromWorkflow),
+    Route(
+        "/macro/{macro_id:str}/execution-group/{group_id:str}/queue",
+        CognitionStartMacroExecutionGroup,
+    ),
     Route("/is_managed", IsManagedRest),
     Route("/is_demo", IsDemoRest),
     Mount("/api", app=fastapi_app, name="REST API"),

controller/organization/manager.py

@@ -8,6 +8,7 @@
 from util import notification
 from controller.auth import kratos
 from submodules.model.util import sql_alchemy_to_dict
+from submodules.s3 import controller as s3
 
 USER_INFO_WHITELIST = {"id", "role"}
 ORGANIZATION_WHITELIST = {
@@ -94,6 +95,7 @@ def create_organization(name: str) -> Organization:
             f"Organization with name {name} already exists"
         )
     organization_item = organization.create(name, with_commit=True)
+    s3.create_bucket(str(organization_item.id))
     return organization_item
 
 

controller/task_queue/handler/macro.py

@@ -0,0 +1,44 @@
+from typing import Any, Dict, Tuple, Callable
+import os
+
+import requests
+from submodules.model.business_objects import (
+    task_queue as task_queue_db_bo,
+    general,
+)
+
+from submodules.model.cognition_objects import macro as macro_db_bo
+
+BASE_URI = os.getenv("COGNITION_GATEWAY")
+
+
+def get_task_functions() -> Tuple[Callable, Callable, int]:
+    return __start_task, __check_finished, 1
+
+
+def __start_task(task: Dict[str, Any]) -> bool:
+    # check task still relevant
+    task_db_obj = task_queue_db_bo.get(task["id"])
+    if task_db_obj is None or task_db_obj.is_active:
+        return False
+
+    task_db_obj.is_active = True
+    general.commit()
+
+    action = task["task_info"]
+
+    macro_id = action["macro_id"]
+    execution_id = action["execution_id"]
+    group_id = action.get("execution_group_id")
+    requests.put(
+        f"{BASE_URI}/api/v1/converters/internal/macros/{macro_id}/execution/{execution_id}/start?group_execution_id={group_id}"
+    )
+    return True
+
+
+def __check_finished(task: Dict[str, Any]) -> bool:
+
+    action = task["task_info"]
+    return macro_db_bo.macro_execution_finished(
+        action["macro_id"], action["execution_id"], action["execution_group_id"]
+    )

controller/task_queue/manager.py

@@ -17,7 +17,8 @@
     attribute_calculation as attribute_calculation_handler,
     task_queue as task_queue_handler,
     markdown_file as markdown_file_handler,
-    parse_cognition_tmp_file as cognition_tmp_file,
+    parse_cognition_tmp_file as cognition_tmp_file_handler,
+    macro as macro_handler,
 )
 from .util import if_task_queue_send_websocket
 
@@ -92,14 +93,20 @@ def get_task_function_by_type(task_type: str) -> Tuple[Callable, Callable, int]:
     if task_type == enums.TaskType.PARSE_MARKDOWN_FILE.value:
         return markdown_file_handler.get_task_functions()
     if task_type == enums.TaskType.PARSE_COGNITION_TMP_FILE.value:
-        return cognition_tmp_file.get_task_functions()
+        return cognition_tmp_file_handler.get_task_functions()
+    if task_type == enums.TaskType.RUN_COGNITION_MACRO.value:
+        return macro_handler.get_task_functions()
     raise ValueError(f"Task type {task_type} not supported yet")
 
 
 def add_task_to_task_queue(task: TaskQueueDBObj) -> int:
     start_func, check_func, check_every = get_task_function_by_type(task.task_type)
     queue = None
-    if task.task_type == enums.TaskType.TASK_QUEUE.value:
+    if (
+        task.task_type == enums.TaskType.TASK_QUEUE.value
+        or task.task_type == enums.TaskType.RUN_COGNITION_MACRO.value
+    ):
+        # macros have tasks (e.g. etl parsing) inside so the execution shouldn't block own queue items
         queue = task_queue.get_task_queue_queue()
     else:
         queue = task_queue.get_task_queue()

controller/task_queue/task_queue.py

@@ -5,8 +5,9 @@
 from submodules.model.models import TaskQueue as TaskQueueDBObj
 from . import manager
 from submodules.model.business_objects import general, task_queue as task_queue_db_bo
+from submodules.model.cognition_objects import macro as macro_db_bo
 import traceback
-from submodules.model.enums import TaskType
+from submodules.model.enums import TaskType, MacroExecutionState
 
 
 # custom class wrapping a list in order to make it thread safe
@@ -207,8 +208,21 @@ def init_task_queues() -> CustomTaskQueue:
         )
         if task.task_type == TaskType.TASK_QUEUE.value:
             task_queue_queue.add_task(task, start_func, check_func, check_every)
+        elif task.task_type == TaskType.RUN_COGNITION_MACRO.value:
+            # macros that were running when the server was restarted are set to failed since we dont have pointers to the running process
+            item = macro_db_bo.get_macro_execution(
+                task.task_info["execution_id"],
+                task.task_info["execution_group_id"],
+                MacroExecutionState.RUNNING,
+            )
+            if item is not None:
+                item.state = MacroExecutionState.FAILED.value
+                task_queue_db_bo.remove_task_from_queue(task.project_id, task.id)
+                continue
+            task_queue_queue.add_task(task, start_func, check_func, check_every)
         else:
             task_queue.add_task(task, start_func, check_func, check_every)
+    general.commit()
     return task_queue