diff --git a/chain_service/database/models/progress_chain.py b/chain_service/database/models/progress_chain.py new file mode 100644 index 0000000..54a8c43 --- /dev/null +++ b/chain_service/database/models/progress_chain.py @@ -0,0 +1,52 @@ +from .chain import Chain +from .base import BaseConfig, BaseMongoModel + +from datetime import datetime +from pydantic import BaseModel, Field, TypeAdapter +from typing import Literal, Union, Annotated, Optional, List + + +class WaitAction(BaseModel): + action_type: Literal["wait"] + wait_for: int + 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 + + +class TextCommentAction(BaseModel): + action_type: Literal["text_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") +] + + +class ProgressChain(BaseMongoModel): + task_id: int + recipients: Annotated[Optional[List[int]], Field(default=[])] + name: Annotated[Optional[str], Field(default=None)] + actions: Annotated[Optional[List[Action]], Field(default=[])] + created_at: Annotated[datetime, Field(default_factory=datetime.utcnow)] + + @classmethod + def create_from_chain( + cls, chain: Chain, task_id: int, recipients: Optional[List[int]] = [] + ): + return ProgressChain( + task_id=task_id, + recipients=recipients, + name=chain.name, + actions=TypeAdapter(List[Action]).dump_python(chain.actions), + )