GeneralManager 0.16.0__py3-none-any.whl → 0.17.0__py3-none-any.whl

general_manager/api/graphql_subscription_consumer.py

@@ -0,0 +1,429 @@
+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
+
+
+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 when the WebSocket disconnects by cancelling and awaiting any active subscription tasks and clearing the subscription registry.
+        
+        Parameters:
+            code (int): The WebSocket close code received from the connection.
+        """
+        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:
+        """
+        Dispatch incoming graphql-transport-ws protocol messages to the appropriate handler.
+        
+        Parameters:
+            content (dict[str, Any]): The received JSON message; expected to include a "type" key
+                whose value is one of "connection_init", "ping", "subscribe", or "complete".
+                Messages with an unrecognized or missing "type" cause the connection to be closed
+                with code 4400.
+        """
+        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" message and start or send the corresponding subscription results.
+        
+        Validates the connection state and the message shape, checks that a subscription-capable schema is available, parses and executes the GraphQL operation, and either streams resulting events to the client or sends a single execution result. On protocol or execution errors the method sends appropriate "error" and "complete" protocol messages. If an active subscription exists for the same operation id it is stopped before the new subscription is started.
+        
+        Parameters:
+            content (dict[str, Any]): The incoming protocol message. Expected keys:
+                - "id" (str): The 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 Exception 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 results from an async iterator to the client for a subscription operation.
+        
+        Iterates the provided async_iterator and sends each yielded execution result to the client for the given operation_id. If an unexpected exception 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_id, and removes the operation from active_subscriptions.
+        
+        Parameters:
+            operation_id (str): The subscription operation identifier to use in protocol messages.
+            async_iterator (Any): An asynchronous iterator that yields execution result objects to be sent to the client.
+        
+        Raises:
+            asyncio.CancelledError: Propagated if 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 Exception 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 into a GraphQL-compatible error dictionary.
+        
+        Parameters:
+            error (Exception): The exception to format; may be a GraphQLError.
+        
+        Returns:
+            dict[str, Any]: If `error` is a `GraphQLError`, returns its `formatted` representation.
+            Otherwise returns a dictionary with a single `"message"` key containing the exception's string.
+        """
+        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 if it exposes an `aclose` coroutine.
+        
+        Parameters:
+            async_iterator (Any): An asynchronous iterator; if it has an `aclose` coroutine method, that coroutine will be awaited to close the iterator.
+        """
+        close = getattr(async_iterator, "aclose", None)
+        if close is None:
+            return
+        await close()

general_manager/apps.py

@@ -3,15 +3,17 @@ from django.apps import AppConfig
 import graphene  # type: ignore[import]
 import os
 from django.conf import settings
-from django.urls import path
+from django.urls import path, re_path
 from graphene_django.views import GraphQLView  # type: ignore[import]
-from importlib import import_module
+from importlib import import_module, util
+import importlib.abc
+import sys
 from general_manager.manager.generalManager import GeneralManager
 from general_manager.manager.meta import GeneralManagerMeta
 from general_manager.manager.input import Input
 from general_manager.api.property import graphQlProperty
 from general_manager.api.graphql import GraphQL
-from typing import TYPE_CHECKING, Type, cast
+from typing import TYPE_CHECKING, Any, Type, cast
 from django.core.checks import register
 import logging
 from django.core.management.base import BaseCommand
@@ -156,7 +158,10 @@ class GeneralmanagerConfig(AppConfig):
         pending_graphql_interfaces: list[Type[GeneralManager]],
     ) -> None:
         """
-        Creates GraphQL interfaces and mutations for the provided general manager classes, builds the GraphQL schema, and registers the GraphQL endpoint in the Django URL configuration.
+        Create GraphQL interfaces and mutations for the given manager classes, build the GraphQL schema, and add the GraphQL endpoint to the URL configuration.
+        
+        Parameters:
+            pending_graphql_interfaces (list[Type[GeneralManager]]): GeneralManager classes that require GraphQL interface and mutation generation.
         """
         logger.debug("Starting to create GraphQL interfaces and mutations...")
         for general_manager_class in pending_graphql_interfaces:
@@ -176,13 +181,26 @@ class GeneralmanagerConfig(AppConfig):
                 },
             )
             GraphQL._mutation_class = mutation_class
-            schema = graphene.Schema(
-                query=GraphQL._query_class,
-                mutation=mutation_class,
-            )
         else:
             GraphQL._mutation_class = None
-            schema = graphene.Schema(query=GraphQL._query_class)
+
+        if GraphQL._subscription_fields:
+            subscription_class = type(
+                "Subscription",
+                (graphene.ObjectType,),
+                GraphQL._subscription_fields,
+            )
+            GraphQL._subscription_class = subscription_class
+        else:
+            GraphQL._subscription_class = None
+
+        schema_kwargs: dict[str, Any] = {"query": GraphQL._query_class}
+        if GraphQL._mutation_class is not None:
+            schema_kwargs["mutation"] = GraphQL._mutation_class
+        if GraphQL._subscription_class is not None:
+            schema_kwargs["subscription"] = GraphQL._subscription_class
+        schema = graphene.Schema(**schema_kwargs)
+        GraphQL._schema = schema
         GeneralmanagerConfig.addGraphqlUrl(schema)
 
     @staticmethod
@@ -208,6 +226,143 @@ class GeneralmanagerConfig(AppConfig):
                 GraphQLView.as_view(graphiql=True, schema=schema),
             )
         )
+        GeneralmanagerConfig._ensure_asgi_subscription_route(graph_ql_url)
+
+    @staticmethod
+    def _ensure_asgi_subscription_route(graphql_url: str) -> None:
+        asgi_path = getattr(settings, "ASGI_APPLICATION", None)
+        if not asgi_path:
+            logger.debug("ASGI_APPLICATION not configured; skipping websocket setup.")
+            return
+
+        try:
+            module_path, attr_name = asgi_path.rsplit(".", 1)
+        except ValueError:
+            logger.warning(
+                "ASGI_APPLICATION '%s' is not a valid module path; skipping websocket setup.",
+                asgi_path,
+            )
+            return
+
+        try:
+            asgi_module = import_module(module_path)
+        except RuntimeError as exc:
+            if "populate() isn't reentrant" not in str(exc):
+                logger.warning(
+                    "Unable to import ASGI module '%s': %s", module_path, exc, exc_info=True
+                )
+                return
+
+            spec = util.find_spec(module_path)
+            if spec is None or spec.loader is None:
+                logger.warning(
+                    "Could not locate loader for ASGI module '%s'; skipping websocket setup.",
+                    module_path,
+                )
+                return
+
+            def finalize(module: Any) -> None:
+                GeneralmanagerConfig._finalize_asgi_module(module, attr_name, graphql_url)
+
+            class _Loader(importlib.abc.Loader):
+                def __init__(self, original_loader: importlib.abc.Loader) -> None:
+                    self._original_loader = original_loader
+
+                def create_module(self, spec):  # type: ignore[override]
+                    if hasattr(self._original_loader, "create_module"):
+                        return self._original_loader.create_module(spec)  # type: ignore[attr-defined]
+                    return None
+
+                def exec_module(self, module):  # type: ignore[override]
+                    self._original_loader.exec_module(module)
+                    finalize(module)
+
+            wrapped_loader = _Loader(spec.loader)
+
+            class _Finder(importlib.abc.MetaPathFinder):
+                def __init__(self) -> None:
+                    self._processed = False
+
+                def find_spec(self, fullname, path, target=None):  # type: ignore[override]
+                    if fullname != module_path or self._processed:
+                        return None
+                    self._processed = True
+                    new_spec = util.spec_from_loader(fullname, wrapped_loader)
+                    if new_spec is not None:
+                        new_spec.submodule_search_locations = spec.submodule_search_locations
+                    return new_spec
+
+            finder = _Finder()
+            sys.meta_path.insert(0, finder)
+            return
+        except Exception as exc:  # pragma: no cover - defensive
+            logger.warning(
+                "Unable to import ASGI module '%s': %s", module_path, exc, exc_info=True
+            )
+            return
+
+        GeneralmanagerConfig._finalize_asgi_module(asgi_module, attr_name, graphql_url)
+
+    @staticmethod
+    def _finalize_asgi_module(asgi_module: Any, attr_name: str, graphql_url: str) -> None:
+        try:
+            from channels.auth import AuthMiddlewareStack  # type: ignore[import-untyped]
+            from channels.routing import ProtocolTypeRouter, URLRouter  # type: ignore[import-untyped]
+            from general_manager.api.graphql_subscription_consumer import (
+                GraphQLSubscriptionConsumer,
+            )
+        except Exception as exc:  # pragma: no cover - optional dependency
+            logger.debug(
+                "Channels or GraphQL subscription consumer unavailable (%s); skipping websocket setup.",
+                exc,
+            )
+            return
+
+        websocket_patterns = getattr(asgi_module, "websocket_urlpatterns", None)
+        if websocket_patterns is None:
+            websocket_patterns = []
+            setattr(asgi_module, "websocket_urlpatterns", websocket_patterns)
+
+        if not hasattr(websocket_patterns, "append"):
+            logger.warning(
+                "websocket_urlpatterns in '%s' does not support appending; skipping websocket setup.",
+                asgi_module.__name__,
+            )
+            return
+
+        normalized = graphql_url.strip("/")
+        pattern = rf"^{normalized}/?$" if normalized else r"^$"
+
+        route_exists = any(
+            getattr(route, "_general_manager_graphql_ws", False)
+            for route in websocket_patterns
+        )
+        if not route_exists:
+            websocket_route = re_path(pattern, GraphQLSubscriptionConsumer.as_asgi()) # type: ignore[arg-type]
+            setattr(websocket_route, "_general_manager_graphql_ws", True)
+            websocket_patterns.append(websocket_route)
+
+        application = getattr(asgi_module, attr_name, None)
+        if application is None:
+            return
+
+        if (
+            hasattr(application, "application_mapping")
+            and isinstance(application.application_mapping, dict)
+        ):
+            application.application_mapping["websocket"] = AuthMiddlewareStack(
+                URLRouter(list(websocket_patterns))
+            )
+        else:
+            wrapped_application = ProtocolTypeRouter(
+                {
+                    "http": application,
+                    "websocket": AuthMiddlewareStack(
+                        URLRouter(list(websocket_patterns))
+                    ),
+                }
+            )
+            setattr(asgi_module, attr_name, wrapped_application)
 
     @staticmethod
     def checkPermissionClass(general_manager_class: Type[GeneralManager]) -> None: