Sprava eventu a organizaci, bugfixing a drobny refaktoring

src/graphql_schema/entities/aircraft.py

@@ -1,16 +1,20 @@
-from typing import List, Optional, Annotated, TYPE_CHECKING
+from typing import List, Optional, Annotated, TYPE_CHECKING, Set
 import strawberry
 from strawberry.file_uploads import Upload
-from sqlalchemy import select
+from sqlalchemy import select, or_
 from database import models
 from decorators.endpoints import authenticated_user_only
 from dependencies.db import get_session
 from graphql_schema.sqlalchemy_to_strawberry_type import strawberry_sqlalchemy_type, strawberry_sqlalchemy_input
 from upload_utils import handle_file_upload, delete_file, get_public_url
+from .helpers.flight import handle_combobox_save
 from ..dataloaders.flight import flights_by_aircraft_dataloader
+from ..dataloaders.organizations import organizations_dataloader
+from ..types import ComboboxInput
 
 if TYPE_CHECKING:
     from .flight import Flight
+    from .organization import Organization
 
 AIRCRAFT_UPLOAD_DEST_PATH = "/app/uploads/aircrafts/"
 
@@ -20,18 +24,29 @@ class Aircraft:
     async def load_flights(root):
         return await flights_by_aircraft_dataloader.load(root.id)
 
+    async def load_organization(root):
+        return await organizations_dataloader.load(root.organization_id)
+
     photo_url: Optional[str] = strawberry.field(
         resolver=lambda root: get_public_url(f"aircrafts/{root.photo_filename}") if root.photo_filename else None
     )
-
     flights: List[Annotated["Flight", strawberry.lazy('.flight')]] = strawberry.field(resolver=load_flights)
+    organization: Optional[Annotated["Organization", strawberry.lazy(".organization")]] = strawberry.field(
+        resolver=load_organization
+    )
 
 
-def get_base_query(user_id: int):
+def get_base_query(user_id: int, organization_ids: Set[int]):
     return (
         select(models.Aircraft)
-        .filter(models.Aircraft.created_by_id == user_id)
+        .filter(
+            or_(
+                models.Aircraft.created_by_id == user_id,
+                models.Aircraft.organization_id.in_(organization_ids)
+            )
+        )
         .filter(models.Aircraft.deleted.is_(False))
+        .order_by(models.Aircraft.id.desc())
     )
 
 
@@ -41,12 +56,10 @@ class AircraftQueries:
     @strawberry.field()
     @authenticated_user_only()
     async def aircrafts(root, info) -> List[Aircraft]:
-        query = (
-            get_base_query(info.context.user_id)
-            .order_by(models.Aircraft.id.desc())
-        )
         async with get_session() as db:
-            aircrafts = (await db.scalars(query)).all()
+            aircrafts = (await db.scalars(
+                get_base_query(info.context.user_id, info.context.organization_ids)
+            )).all()
 
             return [Aircraft(**a.as_dict()) for a in aircrafts]
 
@@ -54,7 +67,7 @@ class AircraftQueries:
     @authenticated_user_only()
     async def aircraft(root, info, id: int) -> Aircraft:
         query = (
-            get_base_query(info.context.user_id)
+            get_base_query(info.context.user_id, info.context.organization_ids)
             .filter(models.Aircraft.id == id)
         )
         async with get_session() as db:
@@ -67,6 +80,7 @@ class CreateAircraftMutation:
     @strawberry_sqlalchemy_input(models.Aircraft, exclude_fields=['id', 'photo_filename'])
     class CreateAircraftInput:
         photo: Optional[Upload]
+        organization: Optional[ComboboxInput] = None
 
     @strawberry.mutation
     @authenticated_user_only()
@@ -78,6 +92,15 @@ class CreateAircraftMutation:
             input_data['photo_filename'] = await handle_file_upload(input.photo, AIRCRAFT_UPLOAD_DEST_PATH)
 
         async with get_session() as db:
+
+            if input.organization:
+                input_data['organization_id'] = await handle_combobox_save(
+                    db,
+                    models.Organization,
+                    input=input.organization,
+                    user_id=info.context.user_id,
+                )
+
             aircraft = await models.Aircraft.create(
                 db,
                 data=dict(
@@ -94,6 +117,7 @@ class EditAircraftMutation:
     @strawberry_sqlalchemy_input(models.Aircraft, exclude_fields=['photo_filename'])
     class EditAircraftInput:
         photo: Optional[Upload]
+        organization: Optional[ComboboxInput] = None
 
     @strawberry.mutation
     @authenticated_user_only()
@@ -102,17 +126,27 @@ class EditAircraftMutation:
 
         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,
+                )
+
             aircraft = (await db.scalars(
-                get_base_query(info.context.user_id)
+                get_base_query(info.context.user_id, set())  # TODO: bude fungovat prazdny set?
                 .filter(models.Aircraft.id == id)
             )).one()
 
-        if input.photo:
-            if aircraft.photo_filename:
-                delete_file(AIRCRAFT_UPLOAD_DEST_PATH + "/" + aircraft.photo_filename, silent=True)
-            update_data['photo_filename'] = await handle_file_upload(input.photo, AIRCRAFT_UPLOAD_DEST_PATH)
+            if input.photo:
+                if aircraft.photo_filename:
+                    delete_file(AIRCRAFT_UPLOAD_DEST_PATH + "/" + aircraft.photo_filename, silent=True)
+                update_data['photo_filename'] = await handle_file_upload(input.photo, AIRCRAFT_UPLOAD_DEST_PATH)
 
-        return await models.Aircraft.update(db, obj=aircraft, data=update_data)
+            aircraft = await models.Aircraft.update(db, obj=aircraft, data=update_data)
+
+            return Aircraft(**aircraft.as_dict())
 
 
 @strawberry.type