Use create and get_or_none, and save model if it exists

Signed-off-by: Aaron Chong <[email protected]>
open-rmf · Mar 17, 2023 · 2a38ce1 · 2a38ce1
1 parent b494b87
commit 2a38ce1
Showing 1 changed file with 18 additions and 29 deletions.
diff --git a/packages/api-server/api_server/repositories/tasks.py b/packages/api-server/api_server/repositories/tasks.py
@@ -28,15 +28,12 @@ def __init__(self, user: User):
         self.user = user
 
     async def save_task_state(self, task_state: TaskState) -> None:
-        task_exists = await DbTaskState.exists(id_=task_state.booking.id)
-
-        if task_exists:
-            previous_task_state = await DbTaskState.get(id_=task_state.booking.id)
+        db_task_state = await DbTaskState.get_or_none(id_=task_state.booking.id)
+        if db_task_state is not None:
             task_state.booking.unix_millis_request_time = (
-                int(round(previous_task_state.unix_millis_request_time.timestamp()))
-                * 1000
+                int(round(db_task_state.unix_millis_request_time.timestamp())) * 1000
             )
-            await ttm.TaskState.update_or_create(
+            db_task_state.update_from_dict(
                 {
                     "data": task_state.json(),
                     "category": task_state.category.__root__
@@ -54,32 +51,24 @@ async def save_task_state(self, task_state: TaskState) -> None:
                         task_state.unix_millis_finish_time / 1000
                     ),
                     "status": task_state.status if task_state.status else None,
-                },
-                id_=task_state.booking.id,
+                }
             )
+            await db_task_state.save()
         else:
             task_state.booking.unix_millis_request_time = round(time.time() * 1000)
-            await ttm.TaskState.update_or_create(
-                {
-                    "data": task_state.json(),
-                    "category": task_state.category.__root__
-                    if task_state.category
-                    else None,
-                    "assigned_to": task_state.assigned_to.name
-                    if task_state.assigned_to
-                    else None,
-                    "unix_millis_start_time": task_state.unix_millis_start_time
-                    and datetime.fromtimestamp(
-                        task_state.unix_millis_start_time / 1000
-                    ),
-                    "unix_millis_finish_time": task_state.unix_millis_finish_time
-                    and datetime.fromtimestamp(
-                        task_state.unix_millis_finish_time / 1000
-                    ),
-                    "status": task_state.status if task_state.status else None,
-                    "unix_millis_request_time": datetime.now(),
-                },
+            await ttm.TaskState.create(
                 id_=task_state.booking.id,
+                data=task_state.json(),
+                category=task_state.category.__root__ if task_state.category else None,
+                assigned_to=task_state.assigned_to.name
+                if task_state.assigned_to
+                else None,
+                unix_millis_start_time=task_state.unix_millis_start_time
+                and datetime.fromtimestamp(task_state.unix_millis_start_time / 1000),
+                unix_millis_finish_time=task_state.unix_millis_finish_time
+                and datetime.fromtimestamp(task_state.unix_millis_finish_time / 1000),
+                status=task_state.status if task_state.status else None,
+                unix_millis_request_time=datetime.now(),
             )
 
     async def query_task_states(