auth-check-middleware

resolvers/schema.py

@@ -2,12 +2,16 @@ from typing import List
 from sqlalchemy import and_, select
 from sqlalchemy.orm import aliased
 from sqlalchemy.exc import SQLAlchemyError
+
+from orm.author import Author
 from orm.notification import Notification as NotificationMessage, NotificationSeen
-from services.auth import login_required
+from services.auth import check_auth
+from aiohttp.web import HTTPUnauthorized
 from services.db import local_session
-from strawberry_sqlalchemy_mapper import StrawberrySQLAlchemyMapper
 import strawberry
+from strawberry_sqlalchemy_mapper import StrawberrySQLAlchemyMapper
 from strawberry.schema.config import StrawberryConfig
+from strawberry.extensions import Extension
 import logging
 
 strawberry_sqlalchemy_mapper = StrawberrySQLAlchemyMapper()
@@ -70,7 +74,6 @@ def get_notifications(author_id: int, session, limit: int, offset: int) -> List[
 
 @strawberry.type
 class Query:
-    @login_required
     @strawberry.field
     async def load_notifications(self, info, limit: int = 50, offset: int = 0) -> NotificationsResult:
         author_id = info.context.get("author_id")
@@ -95,7 +98,6 @@ class Query:
 
 @strawberry.type
 class Mutation:
-    @login_required
     @strawberry.mutation
     async def mark_notification_as_read(self, info, notification_id: int) -> NotificationSeenResult:
         author_id = info.context.get("author_id")
@@ -113,7 +115,6 @@ class Mutation:
                     return NotificationSeenResult(error="cant mark as read")
         return NotificationSeenResult()
 
-    @login_required
     @strawberry.mutation
     async def mark_all_notifications_as_read(self, info) -> NotificationSeenResult:
         author_id = info.context.get("author_id")
@@ -135,4 +136,23 @@ class Mutation:
         return NotificationSeenResult()
 
 
-schema = strawberry.Schema(query=Query, mutation=Mutation, config=StrawberryConfig(auto_camel_case=False))
+class LoginRequiredMiddleware(Extension):
+    async def on_request_start(self):
+        context = self.execution_context.context
+        req = context.get("request")
+        is_authenticated, user_id = await check_auth(req)
+        if not is_authenticated:
+            raise HTTPUnauthorized(text="Please, login first")
+        else:
+            with local_session() as session:
+                author = session.query(Author).filter(Author.user == user_id).first()
+                if author:
+                    context["author_id"] = author.id
+                if user_id:
+                    context["user_id"] = user_id
+        context["user_id"] = user_id
+
+
+schema = strawberry.Schema(
+    query=Query, mutation=Mutation, config=StrawberryConfig(auto_camel_case=False), extensions=[LoginRequiredMiddleware]
+)