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

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

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

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

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "strawberry-graphql"
 packages = [ { include = "strawberry" } ]
-version = "0.253.1"
+version = "0.254.1"
 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"
@@ -100,6 +100,10 @@ poetry-plugin-export = "^1.6.0"
 urllib3 = "<2"
 graphlib_backport = {version = "*", python = "<3.9", optional = false}
 inline-snapshot = "^0.10.1"
+types-deprecated = "^1.2.15.20241117"
+types-six = "^1.17.0.20241205"
+types-pyyaml = "^6.0.12.20240917"
+mypy = "^1.13.0"
 
 [tool.poetry.group.integrations]
 optional = true

{strawberry_graphql-0.253.1 → strawberry_graphql-0.254.1}/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.1}/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.1}/strawberry/http/async_base_view.py

@@ -13,6 +13,7 @@ from typing import (
     Mapping,
     Optional,
     Tuple,
+    Type,
     Union,
     cast,
     overload,
@@ -21,7 +22,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 +39,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
@@ -116,11 +117,19 @@ class AsyncBaseHTTPView(
     connection_init_wait_timeout: timedelta = timedelta(minutes=1)
     request_adapter_class: Callable[[Request], AsyncHTTPRequestAdapter]
     websocket_adapter_class: Callable[
-        ["AsyncBaseHTTPView", WebSocketRequest, WebSocketResponse],
+        [
+            "AsyncBaseHTTPView[Any, Any, Any, Any, Any, Context, RootValue]",
+            WebSocketRequest,
+            WebSocketResponse,
+        ],
         AsyncWebSocketAdapter,
     ]
-    graphql_transport_ws_handler_class = BaseGraphQLTransportWSHandler
-    graphql_ws_handler_class = BaseGraphQLWSHandler
+    graphql_transport_ws_handler_class: Type[
+        BaseGraphQLTransportWSHandler[Context, RootValue]
+    ] = BaseGraphQLTransportWSHandler[Context, RootValue]
+    graphql_ws_handler_class: Type[BaseGraphQLWSHandler[Context, RootValue]] = (
+        BaseGraphQLWSHandler[Context, RootValue]
+    )
 
     @property
     @abc.abstractmethod
@@ -279,18 +288,20 @@ 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,
+                    context=context,  # type: ignore
+                    root_value=root_value,  # type: ignore
                     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(
+                    view=self,
                     websocket=websocket,
-                    context=context,
-                    root_value=root_value,
+                    context=context,  # type: ignore
+                    root_value=root_value,  # type: ignore
                     schema=self.schema,
                     debug=self.debug,
                     keep_alive=self.keep_alive,
@@ -476,5 +487,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.1}/strawberry/http/typevars.py

@@ -1,20 +1,20 @@
-from typing import TypeVar
+from typing_extensions 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")
+Context = TypeVar("Context", default=None)
+RootValue = TypeVar("RootValue", default=None)
 
 
 __all__ = [
+    "Context",
     "Request",
     "Response",
+    "RootValue",
     "SubResponse",
     "WebSocketRequest",
     "WebSocketResponse",
-    "Context",
-    "RootValue",
 ]

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

@@ -5,8 +5,10 @@ import logging
 from contextlib import suppress
 from typing import (
     TYPE_CHECKING,
+    Any,
     Awaitable,
     Dict,
+    Generic,
     List,
     Optional,
     cast,
@@ -14,11 +16,13 @@ from typing import (
 
 from graphql import GraphQLError, GraphQLSyntaxError, parse
 
+from strawberry.exceptions import ConnectionRejectionError
 from strawberry.http.exceptions import (
     NonJsonMessageReceived,
     NonTextMessageReceived,
     WebSocketDisconnected,
 )
+from strawberry.http.typevars import Context, RootValue
 from strawberry.subscriptions.protocols.graphql_transport_ws.types import (
     CompleteMessage,
     ConnectionInitMessage,
@@ -31,29 +35,32 @@ 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
 
 
-class BaseGraphQLTransportWSHandler:
+class BaseGraphQLTransportWSHandler(Generic[Context, RootValue]):
     task_logger: logging.Logger = logging.getLogger("strawberry.ws.task")
 
     def __init__(
         self,
+        view: AsyncBaseHTTPView[Any, Any, Any, Any, Any, Context, RootValue],
         websocket: AsyncWebSocketAdapter,
-        context: object,
-        root_value: object,
+        context: Context,
+        root_value: RootValue,
         schema: BaseSchema,
         debug: bool,
         connection_init_wait_timeout: timedelta,
     ) -> None:
+        self.view = view
         self.websocket = websocket
         self.context = context
         self.root_value = root_value
@@ -64,9 +71,8 @@ class BaseGraphQLTransportWSHandler:
         self.connection_init_received = False
         self.connection_acknowledged = False
         self.connection_timed_out = False
-        self.operations: Dict[str, Operation] = {}
+        self.operations: Dict[str, Operation[Context, RootValue]] = {}
         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 +175,33 @@ 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
+
+        self.context = cast(Context, self.context)
+
+        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 +243,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"],
@@ -236,7 +255,7 @@ class BaseGraphQLTransportWSHandler:
         operation.task = asyncio.create_task(self.run_operation(operation))
         self.operations[message["id"]] = operation
 
-    async def run_operation(self, operation: Operation) -> None:
+    async def run_operation(self, operation: Operation[Context, RootValue]) -> None:
         """The operation task's top level method. Cleans-up and de-registers the operation once it is done."""
         # TODO: Handle errors in this method using self.handle_task_exception()
 
@@ -320,7 +339,7 @@ class BaseGraphQLTransportWSHandler:
                 await task
 
 
-class Operation:
+class Operation(Generic[Context, RootValue]):
     """A class encapsulating a single operation with its id. Helps enforce protocol state transition."""
 
     __slots__ = [
@@ -336,7 +355,7 @@ class Operation:
 
     def __init__(
         self,
-        handler: BaseGraphQLTransportWSHandler,
+        handler: BaseGraphQLTransportWSHandler[Context, RootValue],
         id: str,
         operation_type: OperationType,
         query: str,

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

@@ -4,13 +4,17 @@ import asyncio
 from contextlib import suppress
 from typing import (
     TYPE_CHECKING,
+    Any,
     AsyncGenerator,
     Dict,
+    Generic,
     Optional,
     cast,
 )
 
+from strawberry.exceptions import ConnectionRejectionError
 from strawberry.http.exceptions import NonTextMessageReceived, WebSocketDisconnected
+from strawberry.http.typevars import Context, RootValue
 from strawberry.subscriptions.protocols.graphql_ws.types import (
     ConnectionInitMessage,
     ConnectionTerminateMessage,
@@ -20,24 +24,30 @@ 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:
+class BaseGraphQLWSHandler(Generic[Context, RootValue]):
+    context: Context
+    root_value: RootValue
+
     def __init__(
         self,
+        view: AsyncBaseHTTPView[Any, Any, Any, Any, Any, Context, RootValue],
         websocket: AsyncWebSocketAdapter,
-        context: object,
-        root_value: object,
+        context: Context,
+        root_value: RootValue,
         schema: BaseSchema,
         debug: bool,
         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 +58,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 +101,29 @@ 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
+
+        self.context = cast(Context, self.context)
 
-        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 +141,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.1}/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):