strawberry-graphql 0.253.1__tar.gz → 0.254.0__tar.gz

{strawberry_graphql-0.253.1 → strawberry_graphql-0.254.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: strawberry-graphql
-Version: 0.253.1
+Version: 0.254.0
 Summary: A library for creating GraphQL APIs
 Home-page: https://strawberry.rocks/
 License: MIT

{strawberry_graphql-0.253.1 → strawberry_graphql-0.254.0}/pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "strawberry-graphql"
 packages = [ { include = "strawberry" } ]
-version = "0.253.1"
+version = "0.254.0"
 description = "A library for creating GraphQL APIs"
 authors = ["Patrick Arminio <patrick.arminio@gmail.com>"]
 license = "MIT"
@@ -76,7 +76,7 @@ pygments = "^2.3"
 pyinstrument = {version = ">=4.0.0", optional = false}
 pytest = "^7.2"
 pytest-asyncio = ">=0.20.3"
-pytest-codspeed = "^2.0.1"
+pytest-codspeed = { version = ">=3.0.0", python = ">=3.9" }
 pytest-cov = "^4.0.0"
 pytest-emoji = "^0.2.0"
 pytest-mock = "^3.10"

{strawberry_graphql-0.253.1 → strawberry_graphql-0.254.0}/strawberry/exceptions/__init__.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from functools import cached_property
-from typing import TYPE_CHECKING, Optional, Set, Union
+from typing import TYPE_CHECKING, Dict, Optional, Set, Union
 
 from graphql import GraphQLError
 
@@ -157,6 +157,13 @@ class StrawberryGraphQLError(GraphQLError):
     """Use it when you want to override the graphql.GraphQLError in custom extensions."""
 
 
+class ConnectionRejectionError(Exception):
+    """Use it when you want to reject a WebSocket connection."""
+
+    def __init__(self, payload: Dict[str, object] = {}) -> None:
+        self.payload = payload
+
+
 __all__ = [
     "StrawberryException",
     "UnableToFindExceptionSource",

{strawberry_graphql-0.253.1 → strawberry_graphql-0.254.0}/strawberry/experimental/pydantic/_compat.py

@@ -134,7 +134,7 @@ class PydanticV2Compat:
                 type_=field.annotation,
                 outer_type_=field.annotation,
                 default=field.default,
-                default_factory=field.default_factory,
+                default_factory=field.default_factory,  # type: ignore
                 required=field.is_required(),
                 alias=field.alias,
                 # v2 doesn't have allow_none

{strawberry_graphql-0.253.1 → strawberry_graphql-0.254.0}/strawberry/http/async_base_view.py

@@ -21,7 +21,6 @@ from typing_extensions import Literal, TypeGuard
 
 from graphql import GraphQLError
 
-from strawberry import UNSET
 from strawberry.exceptions import MissingQueryError
 from strawberry.file_uploads.utils import replace_placeholders_with_files
 from strawberry.http import (
@@ -39,6 +38,7 @@ from strawberry.subscriptions.protocols.graphql_transport_ws.handlers import (
 from strawberry.subscriptions.protocols.graphql_ws.handlers import BaseGraphQLWSHandler
 from strawberry.types import ExecutionResult, SubscriptionExecutionResult
 from strawberry.types.graphql import OperationType
+from strawberry.types.unset import UNSET, UnsetType
 
 from .base import BaseView
 from .exceptions import HTTPException
@@ -279,6 +279,7 @@ class AsyncBaseHTTPView(
 
             if websocket_subprotocol == GRAPHQL_TRANSPORT_WS_PROTOCOL:
                 await self.graphql_transport_ws_handler_class(
+                    view=self,
                     websocket=websocket,
                     context=context,
                     root_value=root_value,
@@ -288,6 +289,7 @@ class AsyncBaseHTTPView(
                 ).handle()
             elif websocket_subprotocol == GRAPHQL_WS_PROTOCOL:
                 await self.graphql_ws_handler_class(
+                    view=self,
                     websocket=websocket,
                     context=context,
                     root_value=root_value,
@@ -476,5 +478,10 @@ class AsyncBaseHTTPView(
     ) -> GraphQLHTTPResponse:
         return process_result(result)
 
+    async def on_ws_connect(
+        self, context: Context
+    ) -> Union[UnsetType, None, Dict[str, object]]:
+        return UNSET
+
 
 __all__ = ["AsyncBaseHTTPView"]

{strawberry_graphql-0.253.1 → strawberry_graphql-0.254.0}/strawberry/subscriptions/protocols/graphql_transport_ws/handlers.py

@@ -14,6 +14,7 @@ from typing import (
 
 from graphql import GraphQLError, GraphQLSyntaxError, parse
 
+from strawberry.exceptions import ConnectionRejectionError
 from strawberry.http.exceptions import (
     NonJsonMessageReceived,
     NonTextMessageReceived,
@@ -31,13 +32,14 @@ from strawberry.subscriptions.protocols.graphql_transport_ws.types import (
 from strawberry.types import ExecutionResult
 from strawberry.types.execution import PreExecutionError
 from strawberry.types.graphql import OperationType
+from strawberry.types.unset import UnsetType
 from strawberry.utils.debug import pretty_print_graphql_operation
 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.http.async_base_view import AsyncBaseHTTPView, AsyncWebSocketAdapter
     from strawberry.schema import BaseSchema
     from strawberry.schema.subscribe import SubscriptionResult
 
@@ -47,6 +49,7 @@ class BaseGraphQLTransportWSHandler:
 
     def __init__(
         self,
+        view: AsyncBaseHTTPView,
         websocket: AsyncWebSocketAdapter,
         context: object,
         root_value: object,
@@ -54,6 +57,7 @@ class BaseGraphQLTransportWSHandler:
         debug: bool,
         connection_init_wait_timeout: timedelta,
     ) -> None:
+        self.view = view
         self.websocket = websocket
         self.context = context
         self.root_value = root_value
@@ -66,7 +70,6 @@ class BaseGraphQLTransportWSHandler:
         self.connection_timed_out = False
         self.operations: Dict[str, Operation] = {}
         self.completed_tasks: List[asyncio.Task] = []
-        self.connection_params: Optional[Dict[str, object]] = None
 
     async def handle(self) -> None:
         self.on_request_accepted()
@@ -169,15 +172,31 @@ class BaseGraphQLTransportWSHandler:
             )
             return
 
-        self.connection_params = payload
-
         if self.connection_init_received:
             reason = "Too many initialisation requests"
             await self.websocket.close(code=4429, reason=reason)
             return
 
         self.connection_init_received = True
-        await self.send_message({"type": "connection_ack"})
+
+        if isinstance(self.context, dict):
+            self.context["connection_params"] = payload
+        elif hasattr(self.context, "connection_params"):
+            self.context.connection_params = payload
+
+        try:
+            connection_ack_payload = await self.view.on_ws_connect(self.context)
+        except ConnectionRejectionError:
+            await self.websocket.close(code=4403, reason="Forbidden")
+            return
+
+        if isinstance(connection_ack_payload, UnsetType):
+            await self.send_message({"type": "connection_ack"})
+        else:
+            await self.send_message(
+                {"type": "connection_ack", "payload": connection_ack_payload}
+            )
+
         self.connection_acknowledged = True
 
     async def handle_ping(self, message: PingMessage) -> None:
@@ -219,11 +238,6 @@ class BaseGraphQLTransportWSHandler:
                 message["payload"].get("variables"),
             )
 
-        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
-
         operation = Operation(
             self,
             message["id"],

{strawberry_graphql-0.253.1 → strawberry_graphql-0.254.0}/strawberry/subscriptions/protocols/graphql_ws/handlers.py

@@ -10,6 +10,7 @@ from typing import (
     cast,
 )
 
+from strawberry.exceptions import ConnectionRejectionError
 from strawberry.http.exceptions import NonTextMessageReceived, WebSocketDisconnected
 from strawberry.subscriptions.protocols.graphql_ws.types import (
     ConnectionInitMessage,
@@ -20,16 +21,18 @@ from strawberry.subscriptions.protocols.graphql_ws.types import (
     StopMessage,
 )
 from strawberry.types.execution import ExecutionResult, PreExecutionError
+from strawberry.types.unset import UnsetType
 from strawberry.utils.debug import pretty_print_graphql_operation
 
 if TYPE_CHECKING:
-    from strawberry.http.async_base_view import AsyncWebSocketAdapter
+    from strawberry.http.async_base_view import AsyncBaseHTTPView, AsyncWebSocketAdapter
     from strawberry.schema import BaseSchema
 
 
 class BaseGraphQLWSHandler:
     def __init__(
         self,
+        view: AsyncBaseHTTPView,
         websocket: AsyncWebSocketAdapter,
         context: object,
         root_value: object,
@@ -38,6 +41,7 @@ class BaseGraphQLWSHandler:
         keep_alive: bool,
         keep_alive_interval: Optional[float],
     ) -> None:
+        self.view = view
         self.websocket = websocket
         self.context = context
         self.root_value = root_value
@@ -48,7 +52,6 @@ class BaseGraphQLWSHandler:
         self.keep_alive_task: Optional[asyncio.Task] = None
         self.subscriptions: Dict[str, AsyncGenerator] = {}
         self.tasks: Dict[str, asyncio.Task] = {}
-        self.connection_params: Optional[Dict[str, object]] = None
 
     async def handle(self) -> None:
         try:
@@ -92,9 +95,27 @@ class BaseGraphQLWSHandler:
             await self.websocket.close(code=1000, reason="")
             return
 
-        self.connection_params = payload
+        if isinstance(self.context, dict):
+            self.context["connection_params"] = payload
+        elif hasattr(self.context, "connection_params"):
+            self.context.connection_params = payload
 
-        await self.send_message({"type": "connection_ack"})
+        try:
+            connection_ack_payload = await self.view.on_ws_connect(self.context)
+        except ConnectionRejectionError as e:
+            await self.send_message({"type": "connection_error", "payload": e.payload})
+            await self.websocket.close(code=1011, reason="")
+            return
+
+        if (
+            isinstance(connection_ack_payload, UnsetType)
+            or connection_ack_payload is None
+        ):
+            await self.send_message({"type": "connection_ack"})
+        else:
+            await self.send_message(
+                {"type": "connection_ack", "payload": connection_ack_payload}
+            )
 
         if self.keep_alive:
             keep_alive_handler = self.handle_keep_alive()
@@ -112,11 +133,6 @@ class BaseGraphQLWSHandler:
         operation_name = payload.get("operationName")
         variables = payload.get("variables")
 
-        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
-
         if self.debug:
             pretty_print_graphql_operation(operation_name, query, variables)
 

{strawberry_graphql-0.253.1 → strawberry_graphql-0.254.0}/strawberry/subscriptions/protocols/graphql_ws/types.py

@@ -37,6 +37,7 @@ class ConnectionErrorMessage(TypedDict):
 
 class ConnectionAckMessage(TypedDict):
     type: Literal["connection_ack"]
+    payload: NotRequired[Dict[str, object]]
 
 
 class DataMessagePayload(TypedDict):