From 81d4c73a1f516f0a03294065ab5ec27c86b31db9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Michal=20Kv=C3=A1=C4=8Dek?= <michal.kvacek@performax.cz>
Date: Tue, 18 Jul 2023 10:20:36 +0200
Subject: [PATCH] Uprava vycitani a nastavovani copilota

---
 src/graphql_schema/entities/copilot.py |  4 ++--
 src/graphql_schema/entities/flight.py  | 32 ++++++++++++++------------
 src/main.py                            |  5 +++-
 3 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/src/graphql_schema/entities/copilot.py b/src/graphql_schema/entities/copilot.py
index b126aa5..9bb1b9d 100644
--- a/src/graphql_schema/entities/copilot.py
+++ b/src/graphql_schema/entities/copilot.py
@@ -13,5 +13,5 @@ class CopilotType:
 @strawberry.type
 class CopilotQueries:
     @strawberry.field
-    async def pilots(root, info) -> List[CopilotType]:
-        return (await info.context['db'].scalars(select(Copilot))).all()
+    async def copilots(root, info) -> List[CopilotType]:
+        return (await info.context.db.scalars(select(Copilot))).all()
diff --git a/src/graphql_schema/entities/flight.py b/src/graphql_schema/entities/flight.py
index f42af52..f055793 100644
--- a/src/graphql_schema/entities/flight.py
+++ b/src/graphql_schema/entities/flight.py
@@ -161,8 +161,19 @@ async def handle_track_edit(db: AsyncSession, flight: models.Flight, track: List
         order += 1
 
 
-def handle_copilot_edit():
-    pass
+async def handle_copilot_edit(db: AsyncSession, copilot: CopilotInput, user_id: int) -> int:
+    if copilot.id:
+        return copilot.id
+    else:
+        copilot = await models.Copilot.create(
+            db,
+            data={
+                "name": copilot.name,
+                "created_by_id": user_id,
+            }
+        )
+        await db.flush()
+        return copilot.id
 
 
 @strawberry.type
@@ -179,18 +190,9 @@ class EditFlightMutation:
         if input.track is not None:
             await handle_track_edit(db=info.context.db, flight=flight, track=input.track, user_id=info.context.user_id)
 
-        if input.copilot is not None and not flight.solo:
-            if input.copilot.id:
-                flight.copilot_id = input.copilot.id
-            else:
-                copilot = await models.Copilot.create(
-                    info.context.db,
-                    data={
-                        "created_by_id": info.context.user_id,
-                        "name": input.copilot.name
-                    }
-                )
-                info.context.db.flush()
-                flight.copilot_id = copilot.id
+        if flight.solo:
+            flight.copilot_id = None
+        elif input.copilot is not None:
+            flight.copilot_id = await handle_copilot_edit(info.context.db, input.copilot, info.context.user_id)
 
         return flight
diff --git a/src/main.py b/src/main.py
index 8f153b2..512892e 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,9 +1,10 @@
 from datetime import timedelta
-from fastapi import FastAPI, APIRouter, Depends, Security
+from fastapi import FastAPI, APIRouter, Depends, Security, HTTPException
 from fastapi_jwt import JwtAuthorizationCredentials, JwtAccessBearerCookie, JwtRefreshBearerCookie
 from sqlalchemy.ext.asyncio import AsyncSession
 from starlette.middleware.cors import CORSMiddleware
 from starlette.responses import RedirectResponse, Response
+from starlette.status import HTTP_401_UNAUTHORIZED
 from strawberry.fastapi import GraphQLRouter
 from config import APP_SECRET_KEY, GRAPHIQL, APP_DEBUG
 from dependencies.db import db_session
@@ -63,6 +64,8 @@ class App:
                 credentials: JwtAuthorizationCredentials = Security(self.access_security),
                 db: AsyncSession = Depends(db_session),
         ):
+            if not credentials:
+                raise HTTPException(status_code=HTTP_401_UNAUTHORIZED)
 
             return GraphQLContext(
                 jwt_auth_credentials=credentials,