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

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",
+]

strawberry/litestar/controller.py

@@ -7,19 +7,19 @@ from datetime import timedelta
 from typing import (
     TYPE_CHECKING,
     Any,
+    AsyncGenerator,
     AsyncIterator,
     Callable,
     Dict,
     FrozenSet,
-    List,
     Optional,
-    Set,
     Tuple,
     Type,
     TypedDict,
     Union,
     cast,
 )
+from typing_extensions import TypeGuard
 
 from msgspec import Struct
 
@@ -35,23 +35,24 @@ from litestar import (
 )
 from litestar.background_tasks import BackgroundTasks
 from litestar.di import Provide
-from litestar.exceptions import NotFoundException, ValidationException
+from litestar.exceptions import (
+    NotFoundException,
+    SerializationException,
+    ValidationException,
+    WebSocketDisconnect,
+)
 from litestar.response.streaming import Stream
 from litestar.status_codes import HTTP_200_OK
 from strawberry.exceptions import InvalidCustomContext
-from strawberry.http.async_base_view import AsyncBaseHTTPView, AsyncHTTPRequestAdapter
-from strawberry.http.exceptions import HTTPException
+from strawberry.http.async_base_view import (
+    AsyncBaseHTTPView,
+    AsyncHTTPRequestAdapter,
+    AsyncWebSocketAdapter,
+)
+from strawberry.http.exceptions import HTTPException, NonJsonMessageReceived
 from strawberry.http.types import FormData, HTTPMethod, QueryParams
 from strawberry.http.typevars import Context, RootValue
 from strawberry.subscriptions import GRAPHQL_TRANSPORT_WS_PROTOCOL, GRAPHQL_WS_PROTOCOL
-from strawberry.subscriptions.protocols.graphql_transport_ws import (
-    WS_4406_PROTOCOL_NOT_ACCEPTABLE,
-)
-
-from .handlers.graphql_transport_ws_handler import (
-    GraphQLTransportWSHandler as BaseGraphQLTransportWSHandler,
-)
-from .handlers.graphql_ws_handler import GraphQLWSHandler as BaseGraphQLWSHandler
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
@@ -152,22 +153,6 @@ class GraphQLResource(Struct):
     extensions: Optional[dict[str, object]]
 
 
-class GraphQLWSHandler(BaseGraphQLWSHandler):
-    async def get_context(self) -> Any:
-        return await self._get_context()
-
-    async def get_root_value(self) -> Any:
-        return await self._get_root_value()
-
-
-class GraphQLTransportWSHandler(BaseGraphQLTransportWSHandler):
-    async def get_context(self) -> Any:
-        return await self._get_context()
-
-    async def get_root_value(self) -> Any:
-        return await self._get_root_value()
-
-
 class LitestarRequestAdapter(AsyncHTTPRequestAdapter):
     def __init__(self, request: Request[Any, Any, Any]) -> None:
         self.request = request
@@ -203,10 +188,37 @@ class LitestarRequestAdapter(AsyncHTTPRequestAdapter):
         return FormData(form=multipart_data, files=multipart_data)
 
 
+class LitestarWebSocketAdapter(AsyncWebSocketAdapter):
+    def __init__(self, request: WebSocket, response: WebSocket) -> None:
+        self.ws = response
+
+    async def iter_json(self) -> AsyncGenerator[Dict[str, object], None]:
+        try:
+            try:
+                while self.ws.connection_state != "disconnect":
+                    yield await self.ws.receive_json()
+            except (SerializationException, ValueError):
+                raise NonJsonMessageReceived()
+        except WebSocketDisconnect:
+            pass
+
+    async def send_json(self, message: Mapping[str, object]) -> None:
+        await self.ws.send_json(message)
+
+    async def close(self, code: int, reason: str) -> None:
+        await self.ws.close(code=code, reason=reason)
+
+
 class GraphQLController(
     Controller,
     AsyncBaseHTTPView[
-        Request[Any, Any, Any], Response[Any], Response[Any], Context, RootValue
+        Request[Any, Any, Any],
+        Response[Any],
+        Response[Any],
+        WebSocket,
+        WebSocket,
+        Context,
+        RootValue,
     ],
 ):
     path: str = ""
@@ -219,10 +231,7 @@ class GraphQLController(
     }
 
     request_adapter_class = LitestarRequestAdapter
-    graphql_ws_handler_class: Type[GraphQLWSHandler] = GraphQLWSHandler
-    graphql_transport_ws_handler_class: Type[GraphQLTransportWSHandler] = (
-        GraphQLTransportWSHandler
-    )
+    websocket_adapter_class = LitestarWebSocketAdapter
 
     allow_queries_via_get: bool = True
     graphiql_allowed_accept: FrozenSet[str] = frozenset({"text/html", "*/*"})
@@ -236,6 +245,23 @@ class GraphQLController(
     keep_alive: bool = False
     keep_alive_interval: float = 1
 
+    def is_websocket_request(
+        self, request: Union[Request, WebSocket]
+    ) -> TypeGuard[WebSocket]:
+        return isinstance(request, WebSocket)
+
+    async def pick_websocket_subprotocol(self, request: WebSocket) -> Optional[str]:
+        subprotocols = request.scope["subprotocols"]
+        intersection = set(subprotocols) & set(self.protocols)
+        sorted_intersection = sorted(intersection, key=subprotocols.index)
+        return next(iter(sorted_intersection), None)
+
+    async def create_websocket_response(
+        self, request: WebSocket, subprotocol: Optional[str]
+    ) -> WebSocket:
+        await request.accept(subprotocols=subprotocol)
+        return request
+
     async def execute_request(
         self,
         request: Request[Any, Any, Any],
@@ -245,8 +271,6 @@ class GraphQLController(
         try:
             return await self.run(
                 request,
-                # TODO: check the dependency, above, can we make it so that
-                # we don't need to type ignore here?
                 context=context,
                 root_value=root_value,
             )
@@ -328,14 +352,29 @@ class GraphQLController(
             root_value=root_value,
         )
 
+    @websocket()
+    async def websocket_endpoint(
+        self,
+        socket: WebSocket,
+        context_ws: Any,
+        root_value: Any,
+    ) -> None:
+        await self.run(
+            request=socket,
+            context=context_ws,
+            root_value=root_value,
+        )
+
     async def get_context(
-        self, request: Request[Any, Any, Any], response: Response
+        self,
+        request: Union[Request[Any, Any, Any], WebSocket],
+        response: Union[Response, WebSocket],
     ) -> Context:  # pragma: no cover
         msg = "`get_context` is not used by Litestar's controller"
         raise ValueError(msg)
 
     async def get_root_value(
-        self, request: Request[Any, Any, Any]
+        self, request: Union[Request[Any, Any, Any], WebSocket]
     ) -> RootValue | None:  # pragma: no cover
         msg = "`get_root_value` is not used by Litestar's controller"
         raise ValueError(msg)
@@ -343,54 +382,6 @@ class GraphQLController(
     async def get_sub_response(self, request: Request[Any, Any, Any]) -> Response:
         return self.temporal_response
 
-    @websocket()
-    async def websocket_endpoint(
-        self,
-        socket: WebSocket,
-        context_ws: Any,
-        root_value: Any,
-    ) -> None:
-        async def _get_context() -> Any:
-            return context_ws
-
-        async def _get_root_value() -> Any:
-            return root_value
-
-        preferred_protocol = self.pick_preferred_protocol(socket)
-        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=socket,
-            ).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=socket,
-            ).handle()
-        else:
-            await socket.close(code=WS_4406_PROTOCOL_NOT_ACCEPTABLE)
-
-    def pick_preferred_protocol(self, socket: WebSocket) -> str | None:
-        protocols: List[str] = socket.scope["subprotocols"]
-        intersection: Set[str] = set(protocols) & set(self.protocols)
-        return (
-            min(
-                intersection,
-                key=lambda i: protocols.index(i) if i else "",
-                default=None,
-            )
-            or None
-        )
-
 
 def make_graphql_controller(
     schema: BaseSchema,