GeneralManager 0.16.1__py3-none-any.whl → 0.18.0__py3-none-any.whl

general_manager/api/graphql_subscription_consumer.py

@@ -0,0 +1,432 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+from types import SimpleNamespace
+from typing import Any, cast
+
+from channels.generic.websocket import AsyncJsonWebsocketConsumer  # type: ignore[import-untyped]
+from graphql import (
+    ExecutionResult,
+    GraphQLError,
+    GraphQLSchema,
+    parse,
+    subscribe,
+)
+
+from general_manager.api.graphql import GraphQL
+
+
+RECOVERABLE_SUBSCRIPTION_ERRORS: tuple[type[Exception], ...] = (
+    RuntimeError,
+    ValueError,
+    TypeError,
+    LookupError,
+    ConnectionError,
+    KeyError,
+    asyncio.TimeoutError,
+)
+
+
+class GraphQLSubscriptionConsumer(AsyncJsonWebsocketConsumer):
+    """
+    Websocket consumer implementing the ``graphql-transport-ws`` protocol for GraphQL subscriptions.
+
+    The consumer streams results produced by the dynamically generated GeneralManager GraphQL schema so
+    clients such as GraphiQL can subscribe to live updates.
+    """
+
+    connection_acknowledged: bool
+    connection_params: dict[str, Any]
+
+    async def connect(self) -> None:
+        """
+        Initialize connection state and accept the WebSocket, preferring the "graphql-transport-ws" subprotocol when offered.
+
+        Sets up initial flags and containers used for subscription management (connection_acknowledged, connection_params, active_subscriptions) and accepts the WebSocket connection with the selected subprotocol.
+        """
+        self.connection_acknowledged = False
+        self.connection_params = {}
+        self.active_subscriptions: dict[str, asyncio.Task[None]] = {}
+        subprotocols = self.scope.get("subprotocols", [])
+        selected_subprotocol = (
+            "graphql-transport-ws" if "graphql-transport-ws" in subprotocols else None
+        )
+        await self.accept(subprotocol=selected_subprotocol)
+
+    async def disconnect(self, code: int) -> None:
+        """
+        Perform cleanup on WebSocket disconnect by cancelling and awaiting active subscription tasks and clearing the subscription registry.
+
+        Parameters:
+            code (int): WebSocket close code received from the connection.
+
+        Notes:
+            Awaiting cancelled tasks suppresses asyncio.CancelledError so task cancellation completes silently.
+        """
+        tasks = list(self.active_subscriptions.values())
+        for task in tasks:
+            task.cancel()
+        for task in tasks:
+            with contextlib.suppress(asyncio.CancelledError):
+                await task
+        self.active_subscriptions.clear()
+
+    async def receive_json(self, content: dict[str, Any], **_: Any) -> None:
+        """
+        Route an incoming graphql-transport-ws protocol message to the corresponding handler based on its "type" field.
+
+        Valid message types: "connection_init", "ping", "subscribe", and "complete". Messages with an unrecognized or missing "type" cause the connection to be closed with code 4400.
+
+        Parameters:
+            content (dict[str, Any]): The received JSON message; expected to include a "type" key indicating the protocol action.
+        """
+        message_type = content.get("type")
+        if message_type == "connection_init":
+            await self._handle_connection_init(content)
+        elif message_type == "ping":
+            await self._handle_ping(content)
+        elif message_type == "subscribe":
+            await self._handle_subscribe(content)
+        elif message_type == "complete":
+            await self._handle_complete(content)
+        else:
+            await self.close(code=4400)
+
+    async def _handle_connection_init(self, content: dict[str, Any]) -> None:
+        """
+        Handle a client's "connection_init" message and send a protocol acknowledgment.
+
+        If the connection has already been acknowledged, closes the WebSocket with code 4429.
+        If the incoming message contains a "payload" that is a dict, stores it on
+        self.connection_params; otherwise clears connection_params. Marks the connection
+        as acknowledged and sends a "connection_ack" protocol message.
+
+        Parameters:
+            content (dict[str, Any]): The received WebSocket message for "connection_init".
+        """
+        if self.connection_acknowledged:
+            await self.close(code=4429)
+            return
+        payload = content.get("payload")
+        if isinstance(payload, dict):
+            self.connection_params = payload
+        else:
+            self.connection_params = {}
+        self.connection_acknowledged = True
+        await self._send_protocol_message({"type": "connection_ack"})
+
+    async def _handle_ping(self, content: dict[str, Any]) -> None:
+        """
+        Responds to an incoming ping by sending a pong protocol message.
+
+        If the incoming `content` contains a `"payload"` key, its value is included in the sent pong message under the same key.
+
+        Parameters:
+            content (dict[str, Any]): The received message object; may include an optional `"payload"` to echo.
+        """
+        payload = content.get("payload")
+        response: dict[str, Any] = {"type": "pong"}
+        if payload is not None:
+            response["payload"] = payload
+        await self._send_protocol_message(response)
+
+    async def _handle_subscribe(self, content: dict[str, Any]) -> None:
+        """
+        Handle an incoming GraphQL "subscribe" protocol message and initiate or deliver the corresponding subscription results.
+
+        Parameters:
+            content (dict[str, Any]): The incoming protocol message. Expected keys:
+                - "id" (str): Operation identifier.
+                - "payload" (dict): Operation payload containing:
+                    - "query" (str): GraphQL query string (required).
+                    - "variables" (dict, optional): Operation variables.
+                    - "operationName" (str, optional): Named operation to execute.
+        """
+        if not self.connection_acknowledged:
+            await self.close(code=4401)
+            return
+
+        operation_id = content.get("id")
+        payload = content.get("payload", {})
+        if not isinstance(operation_id, str) or not isinstance(payload, dict):
+            await self.close(code=4403)
+            return
+
+        schema = GraphQL.get_schema()
+        if schema is None or self._schema_has_no_subscription(schema.graphql_schema):
+            await self._send_protocol_message(
+                {
+                    "type": "error",
+                    "id": operation_id,
+                    "payload": [
+                        {"message": "GraphQL subscriptions are not configured."}
+                    ],
+                }
+            )
+            await self._send_protocol_message({"type": "complete", "id": operation_id})
+            return
+
+        query = payload.get("query")
+        if not isinstance(query, str):
+            await self._send_protocol_message(
+                {
+                    "type": "error",
+                    "id": operation_id,
+                    "payload": [{"message": "A GraphQL query string is required."}],
+                }
+            )
+            await self._send_protocol_message({"type": "complete", "id": operation_id})
+            return
+
+        variables = payload.get("variables")
+        if variables is not None and not isinstance(variables, dict):
+            await self._send_protocol_message(
+                {
+                    "type": "error",
+                    "id": operation_id,
+                    "payload": [
+                        {"message": "Variables must be provided as an object."}
+                    ],
+                }
+            )
+            await self._send_protocol_message({"type": "complete", "id": operation_id})
+            return
+
+        operation_name = payload.get("operationName")
+        if operation_name is not None and not isinstance(operation_name, str):
+            await self._send_protocol_message(
+                {
+                    "type": "error",
+                    "id": operation_id,
+                    "payload": [
+                        {
+                            "message": "The operation name must be a string when provided."
+                        }
+                    ],
+                }
+            )
+            await self._send_protocol_message({"type": "complete", "id": operation_id})
+            return
+
+        try:
+            document = parse(query)
+        except GraphQLError as error:
+            await self._send_protocol_message(
+                {
+                    "type": "error",
+                    "id": operation_id,
+                    "payload": [error.formatted],
+                }
+            )
+            await self._send_protocol_message({"type": "complete", "id": operation_id})
+            return
+
+        context = self._build_context()
+
+        try:
+            subscription = await subscribe(
+                schema.graphql_schema,
+                document,
+                variable_values=variables,
+                operation_name=operation_name,
+                context_value=context,
+            )
+        except GraphQLError as error:
+            await self._send_protocol_message(
+                {
+                    "type": "error",
+                    "id": operation_id,
+                    "payload": [error.formatted],
+                }
+            )
+            await self._send_protocol_message({"type": "complete", "id": operation_id})
+            return
+        except (
+            RECOVERABLE_SUBSCRIPTION_ERRORS
+        ) as error:  # pragma: no cover - defensive safeguard
+            await self._send_protocol_message(
+                {
+                    "type": "error",
+                    "id": operation_id,
+                    "payload": [{"message": str(error)}],
+                }
+            )
+            await self._send_protocol_message({"type": "complete", "id": operation_id})
+            return
+
+        if isinstance(subscription, ExecutionResult):
+            await self._send_execution_result(operation_id, subscription)
+            await self._send_protocol_message({"type": "complete", "id": operation_id})
+            return
+
+        if operation_id in self.active_subscriptions:
+            await self._stop_subscription(operation_id)
+
+        self.active_subscriptions[operation_id] = asyncio.create_task(
+            self._stream_subscription(operation_id, subscription)
+        )
+
+    async def _handle_complete(self, content: dict[str, Any]) -> None:
+        """
+        Handle an incoming "complete" protocol message by stopping the subscription for the specified operation.
+
+        If the message payload contains an "id" field that is a string, the corresponding active subscription task is cancelled and cleaned up; otherwise the message is ignored.
+
+        Parameters:
+                content (dict[str, Any]): The received protocol message payload. Expected to contain an "id" key with the operation identifier.
+        """
+        operation_id = content.get("id")
+        if isinstance(operation_id, str):
+            await self._stop_subscription(operation_id)
+
+    async def _stream_subscription(
+        self, operation_id: str, async_iterator: Any
+    ) -> None:
+        """
+        Stream execution results from an async iterator to the client for a subscription operation.
+
+        Sends each yielded execution result for the given operation_id to the client. If a recoverable error occurs while iterating, sends an error payload for the operation. In all cases, attempts to close the iterator, sends a completion message for the operation, and removes the operation from active_subscriptions.
+
+        Parameters:
+            operation_id (str): The subscription operation identifier used in protocol messages.
+            async_iterator (Any): An asynchronous iterator that yields execution result objects to be sent to the client.
+
+        Raises:
+            asyncio.CancelledError: Propagated when the surrounding subscription task is cancelled.
+        """
+        try:
+            async for result in async_iterator:
+                await self._send_execution_result(operation_id, result)
+        except asyncio.CancelledError:
+            raise
+        except (
+            RECOVERABLE_SUBSCRIPTION_ERRORS
+        ) as error:  # pragma: no cover - defensive safeguard
+            await self._send_protocol_message(
+                {
+                    "type": "error",
+                    "id": operation_id,
+                    "payload": [{"message": str(error)}],
+                }
+            )
+        finally:
+            await self._close_iterator(async_iterator)
+            await self._send_protocol_message({"type": "complete", "id": operation_id})
+            self.active_subscriptions.pop(operation_id, None)
+
+    async def _stop_subscription(self, operation_id: str) -> None:
+        """
+        Cancel and await the active subscription task for the given operation id, if one exists.
+
+        If a task is found for operation_id it is cancelled and awaited; CancelledError raised during awaiting is suppressed. If no task exists this is a no-op.
+        """
+        task = self.active_subscriptions.pop(operation_id, None)
+        if task is None:
+            return
+        task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await task
+
+    async def _send_execution_result(
+        self, operation_id: str, result: ExecutionResult
+    ) -> None:
+        """
+        Send a GraphQL execution result to the client as a "next" protocol message.
+
+        The message payload includes a "data" field when result.data is present and an "errors" field when result.errors is non-empty; errors are converted to serializable dictionaries.
+
+        Parameters:
+            operation_id (str): The operation identifier to include as the message `id`.
+            result (ExecutionResult): The GraphQL execution result to serialize and send.
+        """
+        payload: dict[str, Any] = {}
+        if result.data is not None:
+            payload["data"] = result.data
+        if result.errors:
+            payload["errors"] = [self._format_error(error) for error in result.errors]
+        await self._send_protocol_message(
+            {"type": "next", "id": operation_id, "payload": payload}
+        )
+
+    async def _send_protocol_message(self, message: dict[str, Any]) -> None:
+        """
+        Send a JSON-serializable GraphQL transport protocol message over the WebSocket.
+
+        Parameters:
+            message (dict[str, Any]): The protocol message to send. If the connection is already closed, the message is discarded silently.
+        """
+        try:
+            await self.send_json(message)
+        except RuntimeError:
+            # The connection has already been closed. There is nothing else to send.
+            pass
+
+    def _build_context(self) -> Any:
+        """
+        Builds a request context object for GraphQL execution containing the current user, decoded headers, scope, and connection parameters.
+
+        Returns:
+            context (SimpleNamespace): An object with attributes:
+                - `user`: the value of `scope["user"]` (may be None).
+                - `headers`: a dict mapping header names to decoded string values.
+                - `scope`: the consumer's `scope`.
+                - `connection_params`: the connection parameters provided during `connection_init`.
+        """
+        user = self.scope.get("user")
+        raw_headers = self.scope.get("headers") or []
+        headers = {
+            (key.decode("latin1") if isinstance(key, (bytes, bytearray)) else key): (
+                value.decode("latin1")
+                if isinstance(value, (bytes, bytearray))
+                else value
+            )
+            for key, value in raw_headers
+        }
+        return SimpleNamespace(
+            user=user,
+            headers=headers,
+            scope=self.scope,
+            connection_params=self.connection_params,
+        )
+
+    @staticmethod
+    def _schema_has_no_subscription(schema: GraphQLSchema) -> bool:
+        """
+        Check whether the provided GraphQL schema defines no subscription root type.
+
+        Parameters:
+            schema (GraphQLSchema): The schema to inspect.
+
+        Returns:
+            bool: `True` if the schema has no subscription type, `False` otherwise.
+        """
+        return schema.subscription_type is None
+
+    @staticmethod
+    def _format_error(error: Exception) -> dict[str, Any]:
+        """
+        Format an exception as a GraphQL-compatible error dictionary.
+
+        Parameters:
+            error (Exception): The exception to format; if a `GraphQLError`, its `.formatted` representation is used.
+
+        Returns:
+            dict[str, Any]: The error payload: the `GraphQLError.formatted` mapping for GraphQLError instances, otherwise `{"message": str(error)}`.
+        """
+        if isinstance(error, GraphQLError):
+            return cast(dict[str, Any], error.formatted)
+        return {"message": str(error)}
+
+    @staticmethod
+    async def _close_iterator(async_iterator: Any) -> None:
+        """
+        Close an asynchronous iterator by awaiting its `aclose` coroutine if present.
+
+        Parameters:
+            async_iterator (Any): The iterator to close; if it defines an `aclose` coroutine method, that coroutine will be awaited.
+        """
+        close = getattr(async_iterator, "aclose", None)
+        if close is None:
+            return
+        await close()

general_manager/api/mutation.py

@@ -2,6 +2,7 @@
 
 import inspect
 from typing import (
+    Any,
     Callable,
     Optional,
     TypeVar,
@@ -13,21 +14,83 @@ from typing import (
     Type,
     get_type_hints,
     cast,
+    TypeAliasType,
 )
 import graphene  # type: ignore[import]
 from graphql import GraphQLResolveInfo
 
-from general_manager.api.graphql import GraphQL
+from general_manager.api.graphql import GraphQL, HANDLED_MANAGER_ERRORS
 from general_manager.manager.generalManager import GeneralManager
 
 from general_manager.utils.formatString import snake_to_camel
-from typing import TypeAliasType
 from general_manager.permission.mutationPermission import MutationPermission
+from types import UnionType
 
 
 FuncT = TypeVar("FuncT", bound=Callable[..., object])
 
 
+class MissingParameterTypeHintError(TypeError):
+    """Raised when a mutation resolver parameter lacks a type hint."""
+
+    def __init__(self, parameter_name: str, function_name: str) -> None:
+        """
+        Initialize the exception indicating a missing type hint for a function parameter.
+
+        Parameters:
+            parameter_name (str): Name of the parameter that lacks a type hint.
+            function_name (str): Name of the function containing the parameter.
+        """
+        super().__init__(
+            f"Missing type hint for parameter {parameter_name} in {function_name}."
+        )
+
+
+class MissingMutationReturnAnnotationError(TypeError):
+    """Raised when a mutation resolver does not specify a return annotation."""
+
+    def __init__(self, function_name: str) -> None:
+        """
+        Initialize the exception indicating a mutation is missing a return annotation.
+
+        Parameters:
+            function_name (str): Name of the mutation function that lacks a return annotation.
+        """
+        super().__init__(f"Mutation {function_name} missing return annotation.")
+
+
+class InvalidMutationReturnTypeError(TypeError):
+    """Raised when a mutation resolver declares a non-type return value."""
+
+    def __init__(self, function_name: str, return_type: object) -> None:
+        """
+        Initialize an InvalidMutationReturnTypeError for a mutation whose return annotation is not a valid type.
+
+        Parameters:
+            function_name (str): Name of the mutation function that provided the invalid return annotation.
+            return_type (object): The invalid return annotation value that triggered the error.
+        """
+        super().__init__(
+            f"Mutation {function_name} return type {return_type} is not a type."
+        )
+
+
+class DuplicateMutationOutputNameError(ValueError):
+    """Raised when a mutation resolver would expose duplicate output field names."""
+
+    def __init__(self, function_name: str, field_name: str) -> None:
+        """
+        Initialize the exception indicating duplicate output field names.
+
+        Parameters:
+            function_name (str): Name of the mutation function that produced duplicates.
+            field_name (str): The conflicting output field name.
+        """
+        super().__init__(
+            f"Mutation {function_name} produces duplicate output field name '{field_name}'."
+        )
+
+
 def graphQlMutation(
     _func: FuncT | type[MutationPermission] | None = None,
     permission: Optional[Type[MutationPermission]] = None,
@@ -68,21 +131,19 @@ def graphQlMutation(
         mutation_name = snake_to_camel(fn.__name__)
 
         # Build Arguments inner class dynamically
-        arg_fields = {}
+        arg_fields: dict[str, Any] = {}
         for name, param in sig.parameters.items():
             if name == "info":
                 continue
             ann = hints.get(name)
             if ann is None:
-                raise TypeError(
-                    f"Missing type hint for parameter {name} in {fn.__name__}"
-                )
+                raise MissingParameterTypeHintError(name, fn.__name__)
             required = True
             default = param.default
             has_default = default is not inspect._empty
 
             # Prepare kwargs
-            kwargs = {}
+            kwargs: dict[str, Any] = {}
             if required:
                 kwargs["required"] = True
             if has_default:
@@ -90,13 +151,14 @@ def graphQlMutation(
 
             # Handle Optional[...] → not required
             origin = get_origin(ann)
-            if origin is Union and type(None) in get_args(ann):
+            if (origin is Union or origin is UnionType) and type(None) in get_args(ann):
                 required = False
                 # extract inner type
-                ann = [a for a in get_args(ann) if a is not type(None)][0]
+                ann = next(a for a in get_args(ann) if a is not type(None))
                 kwargs["required"] = False
 
             # Resolve list types to List scalar
+            field: Any
             if get_origin(ann) is list or get_origin(ann) is List:
                 inner = get_args(ann)[0]
                 field = graphene.List(
@@ -114,12 +176,12 @@ def graphQlMutation(
         Arguments = type("Arguments", (), arg_fields)
 
         # Build output fields: success + fn return types
-        outputs = {
+        outputs: dict[str, Any] = {
             "success": graphene.Boolean(required=True),
         }
         return_ann: type | tuple[type] | None = hints.get("return")
         if return_ann is None:
-            raise TypeError(f"Mutation {fn.__name__} missing return annotation")
+            raise MissingMutationReturnAnnotationError(fn.__name__)
 
         # Unpack tuple return or single
         out_types = (
@@ -131,11 +193,11 @@ def graphQlMutation(
             is_named_type = isinstance(out, TypeAliasType)
             is_type = isinstance(out, type)
             if not is_type and not is_named_type:
-                raise TypeError(
-                    f"Mutation {fn.__name__} return type {out} is not a type"
-                )
+                raise InvalidMutationReturnTypeError(fn.__name__, out)
             name = out.__name__
             field_name = name[0].lower() + name[1:]
+            if field_name in outputs:
+                raise DuplicateMutationOutputNameError(fn.__name__, field_name)
 
             basis_type = out.__value__ if is_named_type else out
 
@@ -150,26 +212,27 @@ def graphQlMutation(
             **kwargs: object,
         ) -> graphene.Mutation:
             """
-            Execute the mutation resolver, enforcing permissions and formatting output.
+            Execute the mutation resolver, enforce an optional permission check, and convert the resolver result into the mutation's output fields.
 
             Parameters:
                 root: Graphene root object (unused).
-                info: GraphQL execution info passed by Graphene.
+                info: GraphQL execution info provided by Graphene.
                 **kwargs: Mutation arguments provided by the client.
 
             Returns:
-                mutation_class: Instance populated with resolver results and a success flag.
+                mutation_class: Instance of the mutation with output fields populated; `success` is `True` on successful execution and `False` if a handled manager error occurred (after being forwarded to GraphQL._handleGraphQLError).
             """
             if permission:
                 permission.check(kwargs, info.context.user)
             try:
                 result = fn(info, **kwargs)
-                data = {}
+                data: dict[str, Any] = {}
                 if isinstance(result, tuple):
                     # unpack according to outputs ordering after success
                     for (field, _), val in zip(
                         outputs.items(),
-                        [None, *list(result)],  # None for success field to be set later
+                        [None, *list(result)],
+                        strict=False,  # None for success field to be set later
                     ):
                         # skip success
                         if field == "success":
@@ -180,12 +243,11 @@ def graphQlMutation(
                     data[only] = result
                 data["success"] = True
                 return mutation_class(**data)
-            except Exception as e:
-                GraphQL._handleGraphQLError(e)
-                return mutation_class(**{"success": False})
+            except HANDLED_MANAGER_ERRORS as error:
+                raise GraphQL._handleGraphQLError(error) from error
 
         # Assemble class dict
-        class_dict = {
+        class_dict: dict[str, Any] = {
             "Arguments": Arguments,
             "__doc__": fn.__doc__,
             "mutate": staticmethod(_mutate),