From 803aec9a4410dd7e5124c9ec0005278f83ca05b7 Mon Sep 17 00:00:00 2001 From: Robert Date: Wed, 28 Feb 2024 18:08:45 +0700 Subject: [PATCH] refactored --- .../database/models/progress_chain.py | 46 +++++++++++++++---- 1 file changed, 36 insertions(+), 10 deletions(-) diff --git a/chain_service/database/models/progress_chain.py b/chain_service/database/models/progress_chain.py index a6355ae..14f9925 100644 --- a/chain_service/database/models/progress_chain.py +++ b/chain_service/database/models/progress_chain.py @@ -1,35 +1,49 @@ from .chain import Chain from .base import BaseConfig, BaseMongoModel +from enum import Enum from datetime import datetime from pydantic import BaseModel, Field from typing import Literal, Union, Annotated, Optional, List -class WaitAction(BaseModel): - action_type: Literal["wait"] - wait_for: int - done: Annotated[bool, Field(default=False)] +class ProgressActionStatusEnum(str, Enum): + PENDING = "pending" + DONE = "done" + FAILED = "failed" + + +class BaseProgressAction(BaseModel): + action_type: Literal["wait", "comment"] + + status: Annotated[ + ProgressActionStatusEnum, Field(default=ProgressActionStatusEnum.PENDING) + ] + started_at: Annotated[Optional[datetime], Field(default=None)] finished_at: Annotated[Optional[datetime], Field(default=None)] + error_text: Annotated[Optional[str], Field(default=None)] + + +class WaitProgressAction(BaseProgressAction): + action_type: Literal["wait"] + wait_for: int class Config(BaseConfig): pass -class TextCommentAction(BaseModel): - action_type: Literal["text_comment"] +class CommentProgressAction(BaseProgressAction): + action_type: Literal["comment"] text: str - done: Annotated[bool, Field(default=False)] - started_at: Annotated[Optional[datetime], Field(default=None)] - finished_at: Annotated[Optional[datetime], Field(default=None)] class Config(BaseConfig): pass Action = Annotated[ - Union[WaitAction, TextCommentAction], Field(description="action_type") + Union[WaitProgressAction, CommentProgressAction], + Field(description="action_type"), ] @@ -50,3 +64,15 @@ class ProgressChain(BaseMongoModel): name=chain.name, actions=map(Chain.model_dump, chain.actions), ) + + @property + def is_finished(self): + return all( + action.status is ProgressActionStatusEnum.DONE for action in self.actions + ) + + @property + def has_failed(self): + return any( + action.status is ProgressActionStatusEnum.FAILED for action in self.actions + )