fix: run&run_step time field usage timestamp type

MLT-OSS · Feb 21, 2024 · c301922 · c301922
1 parent 6a9f378
commit c301922
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 11 deletions.
diff --git a/app/models/run.py b/app/models/run.py
@@ -1,9 +1,9 @@
-from datetime import datetime
 from typing import Optional
 
 from sqlalchemy import Column, Enum
 from sqlmodel import Field, JSON, TEXT
 
+from app.libs.types import Timestamp
 from app.models.base_model import BaseModel, TimeStampMixin, PrimaryKeyMixin
 
 
@@ -35,11 +35,11 @@ class Run(BaseModel, PrimaryKeyMixin, TimeStampMixin, table=True):
     last_error: Optional[dict] = Field(default=None, sa_column=Column(JSON))
     required_action: Optional[dict] = Field(default=None, sa_column=Column(JSON))
     tools: Optional[list] = Field(default=[], sa_column=Column(JSON))
-    started_at: Optional[datetime] = Field(default=None)
-    completed_at: Optional[datetime] = Field(default=None)
-    cancelled_at: Optional[datetime] = Field(default=None)
-    expires_at: Optional[datetime] = Field(default=None)
-    failed_at: Optional[datetime] = Field(default=None)
+    started_at: Optional[Timestamp] = Field(default=None)
+    completed_at: Optional[Timestamp] = Field(default=None)
+    cancelled_at: Optional[Timestamp] = Field(default=None)
+    expires_at: Optional[Timestamp] = Field(default=None)
+    failed_at: Optional[Timestamp] = Field(default=None)
 
 
 class RunRead(Run):

diff --git a/app/models/run_step.py b/app/models/run_step.py
@@ -1,9 +1,9 @@
-from datetime import datetime
 from typing import Optional
 
 from sqlalchemy import Index, Column, Enum
 from sqlmodel import Field, JSON
 
+from app.libs.types import Timestamp
 from app.models.base_model import BaseModel, TimeStampMixin, PrimaryKeyMixin
 
 
@@ -24,10 +24,10 @@ class RunStep(BaseModel, PrimaryKeyMixin, TimeStampMixin, table=True):
     metadata_: Optional[dict] = Field(default=None, sa_column=Column("metadata", JSON))
     last_error: Optional[dict] = Field(default=None, sa_column=Column(JSON))
     step_details: Optional[dict] = Field(default=None, sa_column=Column(JSON))
-    completed_at: Optional[datetime] = Field(default=None)
-    cancelled_at: Optional[datetime] = Field(default=None)
-    expires_at: Optional[datetime] = Field(default=None)
-    failed_at: Optional[datetime] = Field(default=None)
+    completed_at: Optional[Timestamp] = Field(default=None)
+    cancelled_at: Optional[Timestamp] = Field(default=None)
+    expires_at: Optional[Timestamp] = Field(default=None)
+    failed_at: Optional[Timestamp] = Field(default=None)
     message_id: Optional[str] = Field(default=None)