From d68c2b787998a03d1276405f8f2d6cbb1bc694c5 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 22 Nov 2023 08:22:30 -0300 Subject: [PATCH] Refactor Flow model serialization and validation --- .../services/database/models/flow/flow.py | 21 ++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/src/backend/langflow/services/database/models/flow/flow.py b/src/backend/langflow/services/database/models/flow/flow.py index 754dcfdf4..1bdd13195 100644 --- a/src/backend/langflow/services/database/models/flow/flow.py +++ b/src/backend/langflow/services/database/models/flow/flow.py @@ -4,10 +4,9 @@ from datetime import datetime from typing import TYPE_CHECKING, Dict, Optional from uuid import UUID, uuid4 -from pydantic import field_validator -from sqlmodel import JSON, Column, Field, Relationship - from langflow.services.database.models.base import SQLModelSerializable +from pydantic import field_serializer, field_validator +from sqlmodel import JSON, Column, Field, Relationship if TYPE_CHECKING: from langflow.services.database.models.user import User @@ -36,6 +35,22 @@ class FlowBase(SQLModelSerializable): return v + # updated_at can be serialized to JSON + @field_serializer("updated_at") + def serialize_dt(self, dt: datetime, _info): + if dt is None: + return None + return dt.isoformat() + + @field_validator("updated_at") + def validate_dt(cls, v): + if v is None: + return datetime.utcnow() + elif isinstance(v, datetime): + return v + + return datetime.fromisoformat(v) + class Flow(FlowBase, table=True): id: UUID = Field(default_factory=uuid4, primary_key=True, unique=True)