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

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,

strawberry/quart/views.py

@@ -1,6 +1,7 @@
 import warnings
 from collections.abc import Mapping
 from typing import TYPE_CHECKING, AsyncGenerator, Callable, Dict, Optional, cast
+from typing_extensions import TypeGuard
 
 from quart import Request, Response, request
 from quart.views import View
@@ -46,7 +47,9 @@ class QuartHTTPRequestAdapter(AsyncHTTPRequestAdapter):
 
 
 class GraphQLView(
-    AsyncBaseHTTPView[Request, Response, Response, Context, RootValue],
+    AsyncBaseHTTPView[
+        Request, Response, Response, Request, Response, Context, RootValue
+    ],
     View,
 ):
     _ide_subscription_enabled = False
@@ -121,5 +124,16 @@ class GraphQLView(
             },
         )
 
+    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/sanic/views.py

@@ -13,6 +13,7 @@ from typing import (
     Type,
     cast,
 )
+from typing_extensions import TypeGuard
 
 from sanic.request import Request
 from sanic.response import HTTPResponse, html
@@ -71,7 +72,15 @@ class SanicHTTPRequestAdapter(AsyncHTTPRequestAdapter):
 
 
 class GraphQLView(
-    AsyncBaseHTTPView[Request, HTTPResponse, TemporalResponse, Context, RootValue],
+    AsyncBaseHTTPView[
+        Request,
+        HTTPResponse,
+        TemporalResponse,
+        Request,
+        TemporalResponse,
+        Context,
+        RootValue,
+    ],
     HTTPMethodView,
 ):
     """Class based view to handle GraphQL HTTP Requests.
@@ -206,5 +215,16 @@ class GraphQLView(
         # corner case
         return None  # 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]
+    ) -> TemporalResponse:
+        raise NotImplementedError
+
 
 __all__ = ["GraphQLView"]

strawberry/subscriptions/protocols/graphql_transport_ws/handlers.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import asyncio
 import logging
-from abc import ABC, abstractmethod
 from contextlib import suppress
 from typing import (
     TYPE_CHECKING,
@@ -16,6 +15,7 @@ from typing import (
 
 from graphql import GraphQLError, GraphQLSyntaxError, parse
 
+from strawberry.http.exceptions import NonJsonMessageReceived
 from strawberry.subscriptions.protocols.graphql_transport_ws.types import (
     CompleteMessage,
     ConnectionAckMessage,
@@ -38,6 +38,7 @@ from strawberry.utils.operation import get_operation_type
 if TYPE_CHECKING:
     from datetime import timedelta
 
+    from strawberry.http.async_base_view import AsyncWebSocketAdapter
     from strawberry.schema import BaseSchema
     from strawberry.schema.subscribe import SubscriptionResult
     from strawberry.subscriptions.protocols.graphql_transport_ws.types import (
@@ -45,15 +46,21 @@ if TYPE_CHECKING:
     )
 
 
-class BaseGraphQLTransportWSHandler(ABC):
+class BaseGraphQLTransportWSHandler:
     task_logger: logging.Logger = logging.getLogger("strawberry.ws.task")
 
     def __init__(
         self,
+        websocket: AsyncWebSocketAdapter,
+        context: object,
+        root_value: object,
         schema: BaseSchema,
         debug: bool,
         connection_init_wait_timeout: timedelta,
     ) -> None:
+        self.websocket = websocket
+        self.context = context
+        self.root_value = root_value
         self.schema = schema
         self.debug = debug
         self.connection_init_wait_timeout = connection_init_wait_timeout
@@ -65,28 +72,16 @@ class BaseGraphQLTransportWSHandler(ABC):
         self.completed_tasks: List[asyncio.Task] = []
         self.connection_params: Optional[Dict[str, Any]] = None
 
-    @abstractmethod
-    async def get_context(self) -> Any:
-        """Return the operations context."""
-
-    @abstractmethod
-    async def get_root_value(self) -> Any:
-        """Return the schemas root value."""
-
-    @abstractmethod
-    async def send_json(self, data: dict) -> None:
-        """Send the data JSON encoded to the WebSocket client."""
-
-    @abstractmethod
-    async def close(self, code: int, reason: str) -> None:
-        """Close the WebSocket with the passed code and reason."""
-
-    @abstractmethod
-    async def handle_request(self) -> Any:
-        """Handle the request this instance was created for."""
-
     async def handle(self) -> Any:
-        return await self.handle_request()
+        self.on_request_accepted()
+
+        try:
+            async for message in self.websocket.iter_json():
+                await self.handle_message(message)
+        except NonJsonMessageReceived:
+            await self.handle_invalid_message("WebSocket message type must be text")
+        finally:
+            await self.shutdown()
 
     async def shutdown(self) -> None:
         if self.connection_init_timeout_task:
@@ -118,7 +113,7 @@ class BaseGraphQLTransportWSHandler(ABC):
 
             self.connection_timed_out = True
             reason = "Connection initialisation timeout"
-            await self.close(code=4408, reason=reason)
+            await self.websocket.close(code=4408, reason=reason)
         except Exception as error:
             await self.handle_task_exception(error)  # pragma: no cover
         finally:
@@ -189,14 +184,16 @@ class BaseGraphQLTransportWSHandler(ABC):
         )
 
         if not isinstance(payload, dict):
-            await self.close(code=4400, reason="Invalid connection init payload")
+            await self.websocket.close(
+                code=4400, reason="Invalid connection init payload"
+            )
             return
 
         self.connection_params = payload
 
         if self.connection_init_received:
             reason = "Too many initialisation requests"
-            await self.close(code=4429, reason=reason)
+            await self.websocket.close(code=4429, reason=reason)
             return
 
         self.connection_init_received = True
@@ -211,13 +208,13 @@ class BaseGraphQLTransportWSHandler(ABC):
 
     async def handle_subscribe(self, message: SubscribeMessage) -> None:
         if not self.connection_acknowledged:
-            await self.close(code=4401, reason="Unauthorized")
+            await self.websocket.close(code=4401, reason="Unauthorized")
             return
 
         try:
             graphql_document = parse(message.payload.query)
         except GraphQLSyntaxError as exc:
-            await self.close(code=4400, reason=exc.message)
+            await self.websocket.close(code=4400, reason=exc.message)
             return
 
         try:
@@ -225,12 +222,14 @@ class BaseGraphQLTransportWSHandler(ABC):
                 graphql_document, message.payload.operationName
             )
         except RuntimeError:
-            await self.close(code=4400, reason="Can't get GraphQL operation type")
+            await self.websocket.close(
+                code=4400, reason="Can't get GraphQL operation type"
+            )
             return
 
         if message.id in self.operations:
             reason = f"Subscriber for {message.id} already exists"
-            await self.close(code=4409, reason=reason)
+            await self.websocket.close(code=4409, reason=reason)
             return
 
         if self.debug:  # pragma: no cover
@@ -240,26 +239,28 @@ class BaseGraphQLTransportWSHandler(ABC):
                 message.payload.variables,
             )
 
-        context = await self.get_context()
-        if isinstance(context, dict):
-            context["connection_params"] = self.connection_params
-        root_value = await self.get_root_value()
+        if isinstance(self.context, dict):
+            self.context["connection_params"] = self.connection_params
+        elif hasattr(self.context, "connection_params"):
+            self.context.connection_params = self.connection_params
+
         result_source: Awaitable[ExecutionResult] | Awaitable[SubscriptionResult]
+
         # Get an AsyncGenerator yielding the results
         if operation_type == OperationType.SUBSCRIPTION:
             result_source = self.schema.subscribe(
                 query=message.payload.query,
                 variable_values=message.payload.variables,
                 operation_name=message.payload.operationName,
-                context_value=context,
-                root_value=root_value,
+                context_value=self.context,
+                root_value=self.root_value,
             )
         else:
             result_source = self.schema.execute(
                 query=message.payload.query,
                 variable_values=message.payload.variables,
-                context_value=context,
-                root_value=root_value,
+                context_value=self.context,
+                root_value=self.root_value,
                 operation_name=message.payload.operationName,
             )
 
@@ -312,11 +313,11 @@ class BaseGraphQLTransportWSHandler(ABC):
         await self.cleanup_operation(operation_id=message.id)
 
     async def handle_invalid_message(self, error_message: str) -> None:
-        await self.close(code=4400, reason=error_message)
+        await self.websocket.close(code=4400, reason=error_message)
 
     async def send_message(self, message: GraphQLTransportMessage) -> None:
         data = message.as_dict()
-        await self.send_json(data)
+        await self.websocket.send_json(data)
 
     async def cleanup_operation(self, operation_id: str) -> None:
         if operation_id not in self.operations: