Fix flake

src/graphql_schema/dataloaders/single_model.py

@@ -17,7 +17,10 @@ airport_weather_info_loader = create_dataloader(models.WeatherInfo)
 poi_dataloader = create_dataloader(models.PointOfInterest)
 poi_type_dataloader = create_dataloader(models.PointOfInterestType)
 flight_dataloader = create_dataloader(models.Flight)
-photo_adjustment_dataloader = create_dataloader(models.PhotoAdjustment, relationship_column=models.PhotoAdjustment.photo_id)
+photo_adjustment_dataloader = create_dataloader(
+    models.PhotoAdjustment,
+    relationship_column=models.PhotoAdjustment.photo_id
+)
 photo_dataloader = create_dataloader(models.Photo)
 cover_photo_loader = create_dataloader(
     models.Photo,

src/graphql_schema/entities/aircraft.py

@@ -1,12 +1,6 @@
 from typing import List
 import strawberry
-from database import models
 from decorators.endpoints import authenticated_user_only
-from database.transaction import get_session
-from graphql_schema.entities.helpers.combobox import handle_combobox_save
-from paths import AIRCRAFT_UPLOAD_DEST_PATH
-from utils.file import delete_file
-from utils.upload import handle_file_upload
 from .resolvers.aircraft import AircraftMutationResolver, AircraftQueryResolver
 from graphql_schema.entities.types.mutation_input import CreateAircraftInput, EditAircraftInput
 from graphql_schema.entities.types.types import Aircraft
@@ -35,50 +29,12 @@ class AircraftMutation:
     @strawberry.mutation
     @authenticated_user_only()
     async def create_aircraft(root, info, input: CreateAircraftInput) -> Aircraft:
-        input_data = input.to_dict()
-        if input.photo:
-            input_data['photo_filename'] = await handle_file_upload(input.photo, AIRCRAFT_UPLOAD_DEST_PATH)
-
-        if input.organization:
-            async with get_session() as db:
-                input_data['organization_id'] = await handle_combobox_save(
-                    db,
-                    models.Organization,
-                    input=input.organization,
-                    user_id=info.context.user_id,
-                )
-
-        return await AircraftMutationResolver().create(data=input_data, user_id=info.context.user_id)
+        return await AircraftMutationResolver().create_new(input, info.context.user_id)
 
     @strawberry.mutation
     @authenticated_user_only()
     async def edit_aircraft(root, info, id: int, input: EditAircraftInput) -> Aircraft:
-        update_data = input.to_dict()
-        async with get_session() as db:
-            if input.organization:
-                update_data['organization_id'] = await handle_combobox_save(
-                    db,
-                    models.Organization,
-                    input=input.organization,
-                    user_id=info.context.user_id,
-                )
-
-            query = AircraftQueryResolver().get_query(
-                user_id=info.context.user_id,
-                object_id=id,
-                organization_ids=info.context.organization_ids
-            )
-            aircraft = (await db.scalars(query)).one()
-            existing_photo_filename = aircraft.photo_filename
-
-        if input.photo:
-            if existing_photo_filename:
-                delete_file(AIRCRAFT_UPLOAD_DEST_PATH + "/" + existing_photo_filename, silent=True)
-            update_data['photo_filename'] = await handle_file_upload(input.photo, AIRCRAFT_UPLOAD_DEST_PATH)
-
-        async with get_session() as db:
-            aircraft = await models.Aircraft.update(db, id=id, data=update_data)
-            return Aircraft(**aircraft.as_dict())
+        return await AircraftMutationResolver().edit(id, user_id=info.context.user_id, data=input)
 
     @strawberry.mutation
     @authenticated_user_only()

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())

src/graphql_schema/entities/resolvers/flight.py

@@ -42,7 +42,7 @@ class FlightMutationResolver(BaseMutationResolver):
     def __init__(self):
         super().__init__(Flight, models.Flight)
 
-    async def update(self, context, id: int, input: EditFlightInput):
+    async def update(self, context, id: int, input: EditFlightInput) -> Flight:
         user_id = context.user_id
         async with get_session() as db:
             flight = await self._get_one(db, id, user_id)

src/graphql_schema/entities/resolvers/photo.py

@@ -60,7 +60,6 @@ class PhotoMutationResolver(BaseMutationResolver):
             original_filename=self._copy_original(path, filename)
         )
 
-
     async def upload(self, info, input: UploadPhotoInput) -> Photo:
         path = get_photo_basepath(input.flight_id)
         filename = await handle_file_upload(input.photo, path)
@@ -161,10 +160,13 @@ class PhotoMutationResolver(BaseMutationResolver):
 
         info.context.background_tasks.add_task(generate_thumbnail, path=photo.path, filename=photo.filename)
 
-        async with get_session() as db:
+        async with (get_session() as db):
             await db.execute(delete(models.PhotoAdjustment).filter(models.PhotoAdjustment.photo_id == id))
 
-            crop_info = {"crop_" + key: value for key, value in adjustment.crop.to_dict().items()} if adjustment.crop else {}
+            crop_info = {
+                "crop_" + key: value
+                for key, value in adjustment.crop.to_dict().items()
+            } if adjustment.crop else {}
             await models.PhotoAdjustment.create(db, {
                 "photo_id": id,
                 "contrast": adjustment.contrast,
@@ -189,4 +191,3 @@ class PhotoMutationResolver(BaseMutationResolver):
         delete_file(f"{base_path}/thumbs/{photo.filename}", silent=True)
 
         return photo
-

src/graphql_schema/entities/types/types.py

@@ -14,7 +14,8 @@ from graphql_schema.dataloaders.multi_models import (
 )
 from graphql_schema.dataloaders.single_model import (
     poi_dataloader, poi_type_dataloader, event_dataloader, aircraft_dataloader, airport_dataloader, cover_photo_loader,
-    airport_weather_info_loader, organizations_dataloader, flight_dataloader, photo_adjustment_dataloader, photo_dataloader
+    airport_weather_info_loader, organizations_dataloader, flight_dataloader, photo_adjustment_dataloader,
+    photo_dataloader
 )
 from graphql_schema.sqlalchemy_to_strawberry_type import strawberry_sqlalchemy_type
 from paths import (
@@ -85,7 +86,9 @@ class Photo:
         resolver=lambda root: poi_dataloader.load(root.point_of_interest_id)
     )
     flight: Flight = strawberry.field(resolver=lambda root: flight_dataloader.load(root.flight_id))
-    adjustment: Optional[PhotoAdjustment] = strawberry.field(resolver=lambda root: photo_adjustment_dataloader.load(root.id))
+    adjustment: Optional[PhotoAdjustment] = strawberry.field(
+        resolver=lambda root: photo_adjustment_dataloader.load(root.id)
+    )
 
 
 @strawberry_sqlalchemy_type(models.Flight)

src/main.py

@@ -9,7 +9,10 @@ from starlette.middleware.cors import CORSMiddleware
 from starlette.responses import RedirectResponse, Response
 from starlette.staticfiles import StaticFiles
 from strawberry.fastapi import GraphQLRouter
-from config import APP_SECRET_KEY, GRAPHIQL, APP_DEBUG, ALLOW_CORS_ORIGINS, SENTRY_DSN, ACCESS_TOKEN_VALIDITY_MINUTES, REFRESH_TOKEN_VALIDITY_DAYS
+from config import (
+    APP_SECRET_KEY, GRAPHIQL, APP_DEBUG, ALLOW_CORS_ORIGINS, SENTRY_DSN, ACCESS_TOKEN_VALIDITY_MINUTES,
+    REFRESH_TOKEN_VALIDITY_DAYS
+)
 from database import models, async_session
 from endpoints.login import LoginEndpoint, LoginInput, RefreshEndpoint, LogoutEndpoint
 from endpoints.photo_editor_preview import PhotoEditorEndpoint