strawberry-graphql 0.243.1__py3-none-any.whl → 0.244.1__py3-none-any.whl

strawberry/django/views.py

@@ -13,6 +13,7 @@ from typing import (
     Union,
     cast,
 )
+from typing_extensions import TypeGuard
 
 from asgiref.sync import markcoroutinefunction
 from django.core.serializers.json import DjangoJSONEncoder
@@ -258,7 +259,13 @@ class GraphQLView(
 class AsyncGraphQLView(
     BaseView,
     AsyncBaseHTTPView[
-        HttpRequest, HttpResponseBase, TemporalHttpResponse, Context, RootValue
+        HttpRequest,
+        HttpResponseBase,
+        TemporalHttpResponse,
+        HttpRequest,
+        TemporalHttpResponse,
+        Context,
+        RootValue,
     ],
     View,
 ):
@@ -312,5 +319,16 @@ class AsyncGraphQLView(
 
         return response
 
+    def is_websocket_request(self, request: HttpRequest) -> TypeGuard[HttpRequest]:
+        return False
+
+    async def pick_websocket_subprotocol(self, request: HttpRequest) -> Optional[str]:
+        raise NotImplementedError
+
+    async def create_websocket_response(
+        self, request: HttpRequest, subprotocol: Optional[str]
+    ) -> TemporalHttpResponse:
+        raise NotImplementedError
+
 
 __all__ = ["GraphQLView", "AsyncGraphQLView"]

strawberry/fastapi/router.py

@@ -17,6 +17,7 @@ from typing import (
     Union,
     cast,
 )
+from typing_extensions import TypeGuard
 
 from starlette import status
 from starlette.background import BackgroundTasks  # noqa: TCH002
@@ -34,10 +35,9 @@ from fastapi import APIRouter, Depends, params
 from fastapi.datastructures import Default
 from fastapi.routing import APIRoute
 from fastapi.utils import generate_unique_id
-from strawberry.asgi import ASGIRequestAdapter
+from strawberry.asgi import ASGIRequestAdapter, ASGIWebSocketAdapter
 from strawberry.exceptions import InvalidCustomContext
 from strawberry.fastapi.context import BaseContext, CustomContext
-from strawberry.fastapi.handlers import GraphQLTransportWSHandler, GraphQLWSHandler
 from strawberry.http import process_result
 from strawberry.http.async_base_view import AsyncBaseHTTPView
 from strawberry.http.exceptions import HTTPException
@@ -58,12 +58,14 @@ if TYPE_CHECKING:
 
 
 class GraphQLRouter(
-    AsyncBaseHTTPView[Request, Response, Response, Context, RootValue], APIRouter
+    AsyncBaseHTTPView[
+        Request, Response, Response, WebSocket, WebSocket, Context, RootValue
+    ],
+    APIRouter,
 ):
-    graphql_ws_handler_class = GraphQLWSHandler
-    graphql_transport_ws_handler_class = GraphQLTransportWSHandler
     allow_queries_via_get = True
     request_adapter_class = ASGIRequestAdapter
+    websocket_adapter_class = ASGIWebSocketAdapter
 
     @staticmethod
     async def __get_root_value() -> None:
@@ -261,44 +263,7 @@ class GraphQLRouter(
             context: Context = Depends(self.context_getter),
             root_value: RootValue = Depends(self.root_value_getter),
         ) -> None:
-            async def _get_context() -> Context:
-                return context
-
-            async def _get_root_value() -> RootValue:
-                return root_value
-
-            preferred_protocol = self.pick_preferred_protocol(websocket)
-            if preferred_protocol == GRAPHQL_TRANSPORT_WS_PROTOCOL:
-                await self.graphql_transport_ws_handler_class(
-                    schema=self.schema,
-                    debug=self.debug,
-                    connection_init_wait_timeout=self.connection_init_wait_timeout,
-                    get_context=_get_context,
-                    get_root_value=_get_root_value,
-                    ws=websocket,
-                ).handle()
-            elif preferred_protocol == GRAPHQL_WS_PROTOCOL:
-                await self.graphql_ws_handler_class(
-                    schema=self.schema,
-                    debug=self.debug,
-                    keep_alive=self.keep_alive,
-                    keep_alive_interval=self.keep_alive_interval,
-                    get_context=_get_context,
-                    get_root_value=_get_root_value,
-                    ws=websocket,
-                ).handle()
-            else:
-                # Code 4406 is "Subprotocol not acceptable"
-                await websocket.close(code=4406)
-
-    def pick_preferred_protocol(self, ws: WebSocket) -> Optional[str]:
-        protocols = ws["subprotocols"]
-        intersection = set(protocols) & set(self.protocols)
-        return min(
-            intersection,
-            key=lambda i: protocols.index(i),
-            default=None,
-        )
+            await self.run(request=websocket, context=context, root_value=root_value)
 
     async def render_graphql_ide(self, request: Request) -> HTMLResponse:
         return HTMLResponse(self.graphql_ide_html)
@@ -309,12 +274,12 @@ class GraphQLRouter(
         return process_result(result)
 
     async def get_context(
-        self, request: Request, response: Response
+        self, request: Union[Request, WebSocket], response: Union[Response, WebSocket]
     ) -> Context:  # pragma: no cover
         raise ValueError("`get_context` is not used by FastAPI GraphQL Router")
 
     async def get_root_value(
-        self, request: Request
+        self, request: Union[Request, WebSocket]
     ) -> Optional[RootValue]:  # pragma: no cover
         raise ValueError("`get_root_value` is not used by FastAPI GraphQL Router")
 
@@ -350,5 +315,22 @@ class GraphQLRouter(
             },
         )
 
+    def is_websocket_request(
+        self, request: Union[Request, WebSocket]
+    ) -> TypeGuard[WebSocket]:
+        return request.scope["type"] == "websocket"
+
+    async def pick_websocket_subprotocol(self, request: WebSocket) -> Optional[str]:
+        protocols = request["subprotocols"]
+        intersection = set(protocols) & set(self.protocols)
+        sorted_intersection = sorted(intersection, key=protocols.index)
+        return next(iter(sorted_intersection), None)
+
+    async def create_websocket_response(
+        self, request: WebSocket, subprotocol: Optional[str]
+    ) -> WebSocket:
+        await request.accept(subprotocol=subprotocol)
+        return request
+
 
 __all__ = ["GraphQLRouter"]

strawberry/flask/views.py

@@ -9,6 +9,7 @@ from typing import (
     Union,
     cast,
 )
+from typing_extensions import TypeGuard
 
 from flask import Request, Response, render_template_string, request
 from flask.views import View
@@ -159,7 +160,9 @@ class AsyncFlaskHTTPRequestAdapter(AsyncHTTPRequestAdapter):
 
 class AsyncGraphQLView(
     BaseGraphQLView,
-    AsyncBaseHTTPView[Request, Response, Response, Context, RootValue],
+    AsyncBaseHTTPView[
+        Request, Response, Response, Request, Response, Context, RootValue
+    ],
     View,
 ):
     methods = ["GET", "POST"]
@@ -187,6 +190,17 @@ class AsyncGraphQLView(
     async def render_graphql_ide(self, request: Request) -> Response:
         return render_template_string(self.graphql_ide_html)  # type: ignore
 
+    def is_websocket_request(self, request: Request) -> TypeGuard[Request]:
+        return False
+
+    async def pick_websocket_subprotocol(self, request: Request) -> Optional[str]:
+        raise NotImplementedError
+
+    async def create_websocket_response(
+        self, request: Request, subprotocol: Optional[str]
+    ) -> Response:
+        raise NotImplementedError
+
 
 __all__ = [
     "GraphQLView",

strawberry/http/async_base_view.py

@@ -2,6 +2,7 @@ import abc
 import asyncio
 import contextlib
 import json
+from datetime import timedelta
 from typing import (
     Any,
     AsyncGenerator,
@@ -13,8 +14,10 @@ from typing import (
     Optional,
     Tuple,
     Union,
+    cast,
+    overload,
 )
-from typing_extensions import Literal
+from typing_extensions import Literal, TypeGuard
 
 from graphql import GraphQLError
 
@@ -29,6 +32,11 @@ from strawberry.http import (
 from strawberry.http.ides import GraphQL_IDE
 from strawberry.schema.base import BaseSchema
 from strawberry.schema.exceptions import InvalidOperationTypeError
+from strawberry.subscriptions import GRAPHQL_TRANSPORT_WS_PROTOCOL, GRAPHQL_WS_PROTOCOL
+from strawberry.subscriptions.protocols.graphql_transport_ws.handlers import (
+    BaseGraphQLTransportWSHandler,
+)
+from strawberry.subscriptions.protocols.graphql_ws.handlers import BaseGraphQLWSHandler
 from strawberry.types import ExecutionResult, SubscriptionExecutionResult
 from strawberry.types.graphql import OperationType
 
@@ -36,7 +44,15 @@ from .base import BaseView
 from .exceptions import HTTPException
 from .parse_content_type import parse_content_type
 from .types import FormData, HTTPMethod, QueryParams
-from .typevars import Context, Request, Response, RootValue, SubResponse
+from .typevars import (
+    Context,
+    Request,
+    Response,
+    RootValue,
+    SubResponse,
+    WebSocketRequest,
+    WebSocketResponse,
+)
 
 
 class AsyncHTTPRequestAdapter(abc.ABC):
@@ -63,14 +79,42 @@ class AsyncHTTPRequestAdapter(abc.ABC):
     async def get_form_data(self) -> FormData: ...
 
 
+class AsyncWebSocketAdapter(abc.ABC):
+    @abc.abstractmethod
+    def iter_json(self) -> AsyncGenerator[Dict[str, object], None]: ...
+
+    @abc.abstractmethod
+    async def send_json(self, message: Mapping[str, object]) -> None: ...
+
+    @abc.abstractmethod
+    async def close(self, code: int, reason: str) -> None: ...
+
+
 class AsyncBaseHTTPView(
     abc.ABC,
     BaseView[Request],
-    Generic[Request, Response, SubResponse, Context, RootValue],
+    Generic[
+        Request,
+        Response,
+        SubResponse,
+        WebSocketRequest,
+        WebSocketResponse,
+        Context,
+        RootValue,
+    ],
 ):
     schema: BaseSchema
     graphql_ide: Optional[GraphQL_IDE]
+    debug: bool
+    keep_alive = False
+    keep_alive_interval: Optional[float] = None
+    connection_init_wait_timeout: timedelta = timedelta(minutes=1)
     request_adapter_class: Callable[[Request], AsyncHTTPRequestAdapter]
+    websocket_adapter_class: Callable[
+        [WebSocketRequest, WebSocketResponse], AsyncWebSocketAdapter
+    ]
+    graphql_transport_ws_handler_class = BaseGraphQLTransportWSHandler
+    graphql_ws_handler_class = BaseGraphQLWSHandler
 
     @property
     @abc.abstractmethod
@@ -80,10 +124,16 @@ class AsyncBaseHTTPView(
     async def get_sub_response(self, request: Request) -> SubResponse: ...
 
     @abc.abstractmethod
-    async def get_context(self, request: Request, response: SubResponse) -> Context: ...
+    async def get_context(
+        self,
+        request: Union[Request, WebSocketRequest],
+        response: Union[SubResponse, WebSocketResponse],
+    ) -> Context: ...
 
     @abc.abstractmethod
-    async def get_root_value(self, request: Request) -> Optional[RootValue]: ...
+    async def get_root_value(
+        self, request: Union[Request, WebSocketRequest]
+    ) -> Optional[RootValue]: ...
 
     @abc.abstractmethod
     def create_response(
@@ -102,6 +152,21 @@ class AsyncBaseHTTPView(
     ) -> Response:
         raise ValueError("Multipart responses are not supported")
 
+    @abc.abstractmethod
+    def is_websocket_request(
+        self, request: Union[Request, WebSocketRequest]
+    ) -> TypeGuard[WebSocketRequest]: ...
+
+    @abc.abstractmethod
+    async def pick_websocket_subprotocol(
+        self, request: WebSocketRequest
+    ) -> Optional[str]: ...
+
+    @abc.abstractmethod
+    async def create_websocket_response(
+        self, request: WebSocketRequest, subprotocol: Optional[str]
+    ) -> WebSocketResponse: ...
+
     async def execute_operation(
         self, request: Request, context: Context, root_value: Optional[RootValue]
     ) -> Union[ExecutionResult, SubscriptionExecutionResult]:
@@ -167,35 +232,90 @@ class AsyncBaseHTTPView(
     ) -> None:
         """Hook to allow custom handling of errors, used by the Sentry Integration."""
 
+    @overload
     async def run(
         self,
         request: Request,
         context: Optional[Context] = UNSET,
         root_value: Optional[RootValue] = UNSET,
-    ) -> Response:
-        request_adapter = self.request_adapter_class(request)
+    ) -> Response: ...
 
-        if not self.is_request_allowed(request_adapter):
-            raise HTTPException(405, "GraphQL only supports GET and POST requests.")
+    @overload
+    async def run(
+        self,
+        request: WebSocketRequest,
+        context: Optional[Context] = UNSET,
+        root_value: Optional[RootValue] = UNSET,
+    ) -> WebSocketResponse: ...
 
-        if self.should_render_graphql_ide(request_adapter):
-            if self.graphql_ide:
-                return await self.render_graphql_ide(request)
+    async def run(
+        self,
+        request: Union[Request, WebSocketRequest],
+        context: Optional[Context] = UNSET,
+        root_value: Optional[RootValue] = UNSET,
+    ) -> Union[Response, WebSocketResponse]:
+        root_value = (
+            await self.get_root_value(request) if root_value is UNSET else root_value
+        )
+
+        if self.is_websocket_request(request):
+            websocket_subprotocol = await self.pick_websocket_subprotocol(request)
+            websocket_response = await self.create_websocket_response(
+                request, websocket_subprotocol
+            )
+            websocket = self.websocket_adapter_class(request, websocket_response)
+
+            context = (
+                await self.get_context(request, response=websocket_response)
+                if context is UNSET
+                else context
+            )
+
+            if websocket_subprotocol == GRAPHQL_TRANSPORT_WS_PROTOCOL:
+                await self.graphql_transport_ws_handler_class(
+                    websocket=websocket,
+                    context=context,
+                    root_value=root_value,
+                    schema=self.schema,
+                    debug=self.debug,
+                    connection_init_wait_timeout=self.connection_init_wait_timeout,
+                ).handle()
+            elif websocket_subprotocol == GRAPHQL_WS_PROTOCOL:
+                await self.graphql_ws_handler_class(
+                    websocket=websocket,
+                    context=context,
+                    root_value=root_value,
+                    schema=self.schema,
+                    debug=self.debug,
+                    keep_alive=self.keep_alive,
+                    keep_alive_interval=self.keep_alive_interval,
+                ).handle()
             else:
-                raise HTTPException(404, "Not Found")
+                await websocket.close(4406, "Subprotocol not acceptable")
+
+            return websocket_response
+        else:
+            request = cast(Request, request)
 
+        request_adapter = self.request_adapter_class(request)
         sub_response = await self.get_sub_response(request)
         context = (
             await self.get_context(request, response=sub_response)
             if context is UNSET
             else context
         )
-        root_value = (
-            await self.get_root_value(request) if root_value is UNSET else root_value
-        )
 
         assert context
 
+        if not self.is_request_allowed(request_adapter):
+            raise HTTPException(405, "GraphQL only supports GET and POST requests.")
+
+        if self.should_render_graphql_ide(request_adapter):
+            if self.graphql_ide:
+                return await self.render_graphql_ide(request)
+            else:
+                raise HTTPException(404, "Not Found")
+
         try:
             result = await self.execute_operation(
                 request=request, context=context, root_value=root_value

strawberry/http/exceptions.py

@@ -4,4 +4,8 @@ class HTTPException(Exception):
         self.reason = reason
 
 
+class NonJsonMessageReceived(Exception):
+    pass
+
+
 __all__ = ["HTTPException"]

strawberry/http/typevars.py

@@ -3,8 +3,18 @@ from typing import TypeVar
 Request = TypeVar("Request", contravariant=True)
 Response = TypeVar("Response")
 SubResponse = TypeVar("SubResponse")
+WebSocketRequest = TypeVar("WebSocketRequest")
+WebSocketResponse = TypeVar("WebSocketResponse")
 Context = TypeVar("Context")
 RootValue = TypeVar("RootValue")
 
 
-__all__ = ["Request", "Response", "SubResponse", "Context", "RootValue"]
+__all__ = [
+    "Request",
+    "Response",
+    "SubResponse",
+    "WebSocketRequest",
+    "WebSocketResponse",
+    "Context",
+    "RootValue",
+]