Fix flake

src/graphql_schema/entities/resolvers/aircraft.py

@@ -1,8 +1,14 @@
 from operator import or_
 from typing import Set, Optional
 from database import models
+from database.transaction import get_session
+from graphql_schema.entities.helpers.combobox import handle_combobox_save
 from graphql_schema.entities.resolvers.base import BaseMutationResolver, BaseQueryResolver
+from graphql_schema.entities.types.mutation_input import EditAircraftInput, CreateAircraftInput
 from graphql_schema.entities.types.types import Aircraft
+from paths import AIRCRAFT_UPLOAD_DEST_PATH
+from utils.file import delete_file
+from utils.upload import handle_file_upload
 
 
 class AircraftQueryResolver(BaseQueryResolver):
@@ -39,3 +45,39 @@ class AircraftQueryResolver(BaseQueryResolver):
 class AircraftMutationResolver(BaseMutationResolver):
     def __init__(self):
         super().__init__(graphql_type=Aircraft, model=models.Aircraft)
+
+    async def create_new(self, data: CreateAircraftInput, user_id: int) -> Aircraft:
+        input_data = data.to_dict()
+        if data.photo:
+            input_data['photo_filename'] = await handle_file_upload(data.photo, AIRCRAFT_UPLOAD_DEST_PATH)
+
+        if data.organization:
+            async with get_session() as db:
+                input_data['organization_id'] = await handle_combobox_save(
+                    db,
+                    models.Organization,
+                    input=data.organization,
+                    user_id=user_id,
+                )
+
+        return await self.create(data=input_data, user_id=user_id)
+
+    async def edit(self, id: int, user_id: int, data: EditAircraftInput) -> Aircraft:
+        update_data = data.to_dict()
+        aircraft = await AircraftQueryResolver().get_one(id, user_id)
+
+        if data.photo:
+            if aircraft.photo_filename:
+                delete_file(AIRCRAFT_UPLOAD_DEST_PATH + "/" + aircraft.photo_filename, silent=True)
+            update_data['photo_filename'] = await handle_file_upload(data.photo, AIRCRAFT_UPLOAD_DEST_PATH)
+
+        async with get_session() as db:
+            if data.organization:
+                update_data['organization_id'] = await handle_combobox_save(
+                    db,
+                    models.Organization,
+                    input=data.organization,
+                    user_id=user_id,
+                )
+            aircraft = await models.Aircraft.update(db, id=id, data=update_data)
+            return Aircraft(**aircraft.as_dict())