From 4dd90348c672a2f87ccf069a1687e37c61f8a5e1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20Chm=C3=BArny?= Date: Thu, 13 Feb 2025 00:27:58 +0100 Subject: [PATCH] feat: Add endpoint_body_model arg to pipeline models --- forloop_common_structures/core/pipeline.py | 9 ++++++--- forloop_common_structures/database_models.py | 9 +++++---- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/forloop_common_structures/core/pipeline.py b/forloop_common_structures/core/pipeline.py index 0a73972..d542942 100644 --- a/forloop_common_structures/core/pipeline.py +++ b/forloop_common_structures/core/pipeline.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, field -from typing import Optional +from typing import Any, Optional @dataclass @@ -12,6 +12,7 @@ class Pipeline: remaining_nodes_uids: list = field(default_factory=list) project_uid: str = "0" + endpoint_body_model: dict[str, Any] = field(default_factory=dict) uid: Optional[str] = None def update(self, **kwargs) -> None: @@ -19,10 +20,12 @@ def update(self, **kwargs) -> None: if key in vars(self).keys(): setattr(self, key, value) else: - raise AttributeError(f"Attribute '{key}' cannot be updated, as it does not exist") + raise AttributeError( + f"Attribute '{key}' cannot be updated, as it does not exist" + ) def run(self): self.is_active = True - def stop(self): #todo: pause + def stop(self): # todo: pause self.is_active = False diff --git a/forloop_common_structures/database_models.py b/forloop_common_structures/database_models.py index 40f9fe7..bb45cbc 100644 --- a/forloop_common_structures/database_models.py +++ b/forloop_common_structures/database_models.py @@ -106,6 +106,7 @@ class DBPipeline(dh.AbstractModel): project_uid: int # Foreign Key Many-to-1 system_reactivation_status: str = None + endpoint_body_model: dict[Any] = {} def cast_pipeline_types_to_app(pipelines_df: pd.DataFrame) -> pd.DataFrame: @@ -113,8 +114,8 @@ def cast_pipeline_types_to_app(pipelines_df: pd.DataFrame) -> pd.DataFrame: pipelines_df = pipelines_df.astype( {"uid": str, "project_uid": str, "is_active": bool} ) - pipelines_df[["active_nodes_uids", "remaining_nodes_uids"] - ] = pipelines_df[["active_nodes_uids", "remaining_nodes_uids"]].map(json.loads) + pipelines_df[["active_nodes_uids", "remaining_nodes_uids", "endpoint_body_model"] + ] = pipelines_df[["active_nodes_uids", "remaining_nodes_uids", "endpoint_body_model"]].map(json.loads) return pipelines_df @@ -122,8 +123,8 @@ def cast_pipeline_types_to_db(pipelines_df: pd.DataFrame) -> pd.DataFrame: """Cast in-app python datatypes to DB datatypes.""" pipelines_df = pipelines_df.drop("uid", axis=1) pipelines_df = pipelines_df.astype({"project_uid": int, "is_active": int}) - pipelines_df[["active_nodes_uids", "remaining_nodes_uids"] - ] = pipelines_df[["active_nodes_uids", "remaining_nodes_uids"]].map(json.dumps) + pipelines_df[["active_nodes_uids", "remaining_nodes_uids", "endpoint_body_model"] + ] = pipelines_df[["active_nodes_uids", "remaining_nodes_uids", "endpoint_body_model"]].map(json.dumps) pipelines_df = pipelines_df.map(escape_if_string) return pipelines_df