Merge pull request #144 from Tauffer-Consulting/fix/workflows

Fix/workflows

frontend/src/features/workflowEditor/context/types/workflowPieceData.ts

@@ -62,6 +62,7 @@ export interface TasksDataModel {
   task_id: string;
   piece: {
     name: string;
+    source_image: string;
   };
   piece_input_kwargs: Record<string, any>;
   dependencies?: string[];

frontend/src/features/workflowEditor/context/workflowsEditor.tsx

@@ -165,6 +165,7 @@ const WorkflowsEditorProvider: FC<{ children?: React.ReactNode }> = ({
 
   const generateWorkflowsEditorBodyParams = useCallback(
     async ({
+      workflowPieces,
       workflowPiecesData,
       workflowSettingsData,
       workflowNodes,
@@ -187,8 +188,8 @@ const WorkflowsEditorProvider: FC<{ children?: React.ReactNode }> = ({
 
       for (const element of workflowNodes) {
         const elementData = workflowPiecesData[element.id];
+        const pieceData = workflowPieces[element.id];
 
-        const numberId = getIdSlice(element.id);
         const taskName = generateTaskName(element.data.name, element.id);
 
         ui_schema.nodes[taskName] = element;
@@ -265,6 +266,7 @@ const WorkflowsEditorProvider: FC<{ children?: React.ReactNode }> = ({
           task_id: taskName,
           piece: {
             name: element.data.name,
+            source_image: pieceData.source_image,
           },
           dependencies,
           piece_input_kwargs: pieceInputKwargs,

rest/repository/piece_repository.py

@@ -22,7 +22,7 @@ def find_by_ids(self, ids: list):
                 session.expunge_all()
         return result
 
-    def find_repositories_by_piece_name_and_workspace_id(self, pieces_names: list, workspace_id):
+    def find_repositories_by_piece_name_and_workspace_id(self, pieces_names: list, sources_images: list, workspace_id):
         # Find pieces repositories by pieces names and workspace_id
         with session_scope() as session:
             query = session.query(
@@ -32,7 +32,8 @@ def find_repositories_by_piece_name_and_workspace_id(self, pieces_names: list, w
             )\
                 .filter(PieceRepositoryDatabaseModel.workspace_id == workspace_id)\
                     .join(Piece, Piece.repository_id == PieceRepositoryDatabaseModel.id)\
-                        .filter(Piece.name.in_(pieces_names))
+                        .filter(Piece.name.in_(pieces_names))\
+                            .filter(Piece.source_image.in_(sources_images))
 
             result = query.all()
             session.flush()

rest/schemas/requests/workflow.py

@@ -1,6 +1,6 @@
 from typing import Dict, List, Optional
 from enum import Enum
-from pydantic import BaseModel, field_validator, Field
+from pydantic import BaseModel, field_validator, Field, ConfigDict
 from pydantic_core.core_schema import FieldValidationInfo
 from datetime import datetime
 from constants.default_pieces.storage import AWSS3StoragePiece
@@ -82,8 +82,8 @@ def end_date_validator(cls, v, info: FieldValidationInfo):
         except ValueError:
             raise ValueError(f"Invalid end date: {v}")
 
-    class Config:
-        populate_by_name = True
+
+    model_config = ConfigDict(populate_by_name=True)
 
 
 storage_default_piece_model_map = {
@@ -107,11 +107,12 @@ class WorkflowSharedStorageDataModel(BaseModel):
     mode: Optional[WorkflowSharedStorageModeEnum] = None
     provider_options: Optional[Dict] = None
 
-    class Config:
-        use_enum_values = True
+
+    model_config = ConfigDict(use_enum_values=True)
 
 class TaskPieceDataModel(BaseModel):
     name: str
+    source_image: str
 
 class SystemRequirementsModel(BaseModel):
     cpu: float

rest/schemas/responses/workflow.py

@@ -1,5 +1,5 @@
 from schemas.responses.base import PaginationSet
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, Field, field_validator, ConfigDict
 from datetime import datetime, timezone
 from typing import Dict, Optional, List, Union
 from enum import Enum
@@ -154,8 +154,7 @@ def set_state(cls, state):
         return state or WorkflowRunState.none
 
 
-    class Config:
-        populate_by_name = True
+    model_config = ConfigDict(populate_by_name=True)
 
 
 class GetWorkflowRunsResponse(BaseModel):
@@ -179,8 +178,8 @@ class GetWorkflowRunTasksResponseData(BaseModel):
     def set_state(cls, state):
         return state or WorkflowRunTaskState.none
 
-    class Config:
-        populate_by_name = True
+
+    model_config = ConfigDict(populate_by_name=True)
 
 
 class GetWorkflowRunTasksResponse(BaseModel):

rest/services/workflow_service.py

@@ -304,9 +304,11 @@ def _validate_workflow_tasks(self, tasks_dict: dict, workspace_id: int):
         # get all workspaces ids necessary for repositories referenced by tasks
         # check if user has access to all necessary workspaces
         pieces_names = set()
+        pieces_source_images = set()
         shared_storage_sources = []
         for v in tasks_dict.values():
             pieces_names.add(v['piece']['name'])
+            pieces_source_images.add(v['piece']['source_image'])
             shared_storage_sources.append(v['workflow_shared_storage']['source'])
 
         shared_storage_source = list(set(shared_storage_sources))
@@ -336,6 +338,7 @@ def _validate_workflow_tasks(self, tasks_dict: dict, workspace_id: int):
         # Find necessary repositories from pieces names and workspace_id
         necessary_repositories_and_pieces = self.piece_repository.find_repositories_by_piece_name_and_workspace_id(
             pieces_names=pieces_names,
+            sources_images=pieces_source_images,
             workspace_id=workspace_id
         )
 

rest/tests/auth/test_auth_router.py

@@ -20,7 +20,7 @@ def test_register(user: User, register: Response, authorization_token: Dict):
         )
         response = register
         content = response.json()
-        mock_response_content = json.loads(mock_response.json())
+        mock_response_content = json.loads(mock_response.model_dump_json())
 
         assert response.status_code == 201
         assert content["email"] == mock_response_content["email"]
@@ -36,7 +36,7 @@ def test_login(user: User, register: Response, login: Response, authorization_to
         )
         response = login
         content = response.json()
-        mock_response_content = json.loads(mock_response.json())
+        mock_response_content = json.loads(mock_response.model_dump_json())
 
         assert response.status_code == 200
         assert content.keys() == mock_response_content.keys()

rest/tests/piece_repository/fixtures.py

@@ -34,7 +34,7 @@ def add_piece_repository(client: ApiTestClient, authorization_token: Dict, piece
         version=piece_repository.version,
         url=piece_repository.url
     )
-    body = json.loads(add_piece_repository_request.json())
+    body = json.loads(add_piece_repository_request.model_dump_json())
     response = client.post(
         "/pieces-repositories",
         headers={"Authorization": authorization_token["header"]},
@@ -79,7 +79,7 @@ def patch_piece_repository(client: ApiTestClient, authorization_token: Dict, pie
     patch_piece_repository_request = PatchRepositoryRequest(
         version=piece_repository.version
     )
-    body = json.loads(patch_piece_repository_request.json())
+    body = json.loads(patch_piece_repository_request.model_dump_json())
     response = client.patch(
         f"/pieces-repositories/{piece_repository.id}",
         headers={"Authorization": authorization_token["header"]},

rest/tests/piece_repository/test_piece_repository_router.py

@@ -21,7 +21,7 @@ def test_get_repository_releases(get_repository_releases: Response):
         ]
         response = get_repository_releases
         content = response.json()
-        mock_response_content = json.loads(mock_response[0].json())
+        mock_response_content = json.loads(mock_response[0].model_dump_json())
 
         assert response.status_code == 200
         assert content[0].keys() == mock_response_content.keys()
@@ -41,7 +41,7 @@ def test_add_piece_repository(add_piece_repository: Response, piece_repository:
         response = add_piece_repository
         content = response.json()
         if content.get("id") != None:
-            mock_response_content = json.loads(mock_response.json())
+            mock_response_content = json.loads(mock_response.model_dump_json())
 
             assert response.status_code == 200       
             assert content.keys() == mock_response_content.keys()
@@ -72,7 +72,7 @@ def test_get_pieces_repositories(add_piece_repository: Response, get_pieces_repo
         )
         response = get_pieces_repositories
         content = response.json()
-        mock_response_content = json.loads(mock_response.json())
+        mock_response_content = json.loads(mock_response.model_dump_json())
 
         assert response.status_code == 200
         assert content.keys() == mock_response_content.keys()
@@ -92,7 +92,7 @@ def test_get_piece_repository_release_data(add_piece_repository: Response, get_p
         )
         response = get_piece_repository_release_data
         content = response.json()
-        mock_response_content = json.loads(mock_response.json())
+        mock_response_content = json.loads(mock_response.model_dump_json())
 
         assert response.status_code == 200
         assert content.keys() == mock_response_content.keys()
@@ -112,7 +112,7 @@ def test_get_piece_repository_by_id(add_piece_repository: Response, piece_reposi
         )
         response = get_piece_repository_by_id
         content = response.json()
-        mock_response_content = json.loads(mock_response.json())
+        mock_response_content = json.loads(mock_response.model_dump_json())
 
         assert response.status_code == 200
         assert content.keys() == mock_response_content.keys()
@@ -133,7 +133,7 @@ def test_patch_piece_repository(add_piece_repository: Response, patch_piece_repo
         )
         response = patch_piece_repository
         content = response.json()
-        mock_response_content = json.loads(mock_response.json())
+        mock_response_content = json.loads(mock_response.model_dump_json())
 
         assert response.status_code == 200
         assert content.keys() == mock_response_content.keys()

rest/tests/secret/test_secret_router.py

@@ -49,7 +49,7 @@ def test_get_secrets_by_piece_name(get_secrets_by_piece_name: Response):
         ]
         response = get_secrets_by_piece_name
         content = response.json()
-        mock_response_content = json.loads(mock_response[0].json())
+        mock_response_content = json.loads(mock_response[0].model_dump_json())
 
         assert response.status_code == 200
         assert content[0] == mock_response_content

rest/tests/workflow/create_workflow_request_model.py

@@ -11,8 +11,8 @@
         "SimpleLogP_0298c1669d404e08b631ebe1490e1c45": {
             "task_id": "SimpleLogP_0298c1669d404e08b631ebe1490e1c45",
             "piece": {
-                "id": 988,
-                "name": "SimpleLogPiece"
+                "name": "SimpleLogPiece",
+                "source_image": "ghcr.io/tauffer-consulting/default_domino_pieces_tests:0.0.2-group0"
             },
             "workflow_shared_storage": {
                 "source": "None",

rest/tests/workflow/test_workflow_router.py

@@ -46,7 +46,7 @@ def test_create_workflow(patch_piece_secret: Response, create_workflow: Response
         )
         response=create_workflow
         content=response.json()
-        mock_response_content=json.loads(mock_response.json(by_alias=True))
+        mock_response_content=json.loads(mock_response.model_dump_json(by_alias=True))
 
         assert response.status_code == 201
         assert content.keys() == mock_response_content.keys()
@@ -75,19 +75,19 @@ def test_create_workflow(patch_piece_secret: Response, create_workflow: Response
     @pytest.mark.skip(reason="Requires a workflow to be created")
     def test_get_workflow(get_workflow: Response, workflow: Workflow, default_workspace: Workspace, user: User):
         mock_response = GetWorkflowResponse(
-            id = workflow.id,
-            name = workflow.name,
-            created_at = workflow.created_at,
-            schema = BaseWorkflowModel(
-                workflow = WorkflowConfigResponse(
-                    name = workflow.name,
-                    start_date = str(datetime.utcnow())
+            id=workflow.id,
+            name=workflow.name,
+            created_at=workflow.created_at,
+            schema=BaseWorkflowModel(
+                workflow=WorkflowConfigResponse(
+                    name=workflow.name,
+                    start_date=str(datetime.utcnow())
                 ),
-                tasks = dict()
+                tasks=dict()
             ),
-            ui_schema = BaseUiSchema(
-                nodes = dict(),
-                edges = list(dict())
+            ui_schema=BaseUiSchema(
+                nodes=dict(),
+                edges=list(dict())
             ),
             last_changed_at = datetime.utcnow(),
             last_changed_by = user.id,

rest/tests/workspace/fixtures.py

@@ -69,7 +69,7 @@ def create_workspace(client: ApiTestClient, authorization_token: Dict, workspace
     create_workspace_request = CreateWorkspaceRequest(
         name = workspace.name
     )
-    body = json.loads(create_workspace_request.json())
+    body = json.loads(create_workspace_request.model_dump_json())
     response = client.post(
         "/workspaces",
         headers = {"Authorization": authorization_token["header"]},

rest/tests/workspace/test_workspace_router.py

@@ -23,7 +23,7 @@ def test_create_workspace(workspace: Workspace, create_workspace: Response):
         )
         response = create_workspace
         content = response.json()
-        mock_response_content = json.loads(mock_response.json())
+        mock_response_content = json.loads(mock_response.model_dump_json())
 
         assert response.status_code == 200
         assert content.keys() == mock_response_content.keys()
@@ -43,7 +43,7 @@ def test_get_workspaces(create_workspace: Response, workspace: Workspace, get_wo
         ]
         response = get_workspaces
         content = response.json()
-        mock_response_content = json.loads(mock_response[0].json())
+        mock_response_content = json.loads(mock_response[0].model_dump_json())
 
         assert response.status_code == 200
         assert content[0]["workspace_name"].startswith("Default")
@@ -63,7 +63,7 @@ def test_get_workspace(create_workspace: Response, get_workspace: Response, work
         )
         response = get_workspace
         content = response.json()
-        mock_response_content = json.loads(mock_response.json())
+        mock_response_content = json.loads(mock_response.model_dump_json())
 
         assert response.status_code == 200
         assert content.keys() == mock_response_content.keys()
@@ -97,7 +97,7 @@ def test_reject_invite(
         )
         response = reject_invite
         content = response.json()
-        mock_response_content = json.loads(mock_response.json())
+        mock_response_content = json.loads(mock_response.model_dump_json())
 
         assert response.status_code == 200
         assert content.keys() == mock_response_content.keys()
@@ -121,7 +121,7 @@ def test_accept_invite(
         )
         response = accept_invite
         content = response.json()
-        mock_response_content = json.loads(mock_response.json())
+        mock_response_content = json.loads(mock_response.model_dump_json())
 
         assert response.status_code == 200
         assert content.keys() == mock_response_content.keys()
@@ -162,7 +162,7 @@ def test_list_workspace_users(
         )
         response = list_workspace_users
         content = response.json()
-        mock_response_content = json.loads(mock_response.json())
+        mock_response_content = json.loads(mock_response.model_dump_json())
 
         assert response.status_code == 200
         assert content.keys() == mock_response_content.keys()
@@ -189,7 +189,7 @@ def test_patch_workspace(create_workspace: Response, patch_workspace: Response,
         )
         response = patch_workspace
         content = response.json()
-        mock_response_content = json.loads(mock_response.json())
+        mock_response_content = json.loads(mock_response.model_dump_json())
         assert response.status_code == 200
         assert content.keys() == mock_response_content.keys()
         for key in content.keys():