GraphQL Subscriptions Support

README.md

@@ -16,6 +16,8 @@ pip install python-graphql-client
 
 ## Usage
 
+- Query/Mutation
+
 ```python
 from python_graphql_client import GraphqlClient
 
@@ -45,6 +47,31 @@ data = asyncio.run(client.execute_async(query=query, variables=variables))
 print(data)  # => {'data': {'country': {'code': 'CA', 'name': 'Canada'}}}
 ```
 
+- Subscription
+
+```python
+from python_graphql_client import GraphqlClient
+
+# Instantiate the client with a websocket endpoint.
+client = GraphqlClient(endpoint="wss://www.your-api.com/graphql")
+
+# Create the query string and variables required for the request.
+query = """
+    subscription onMessageAdded {
+        messageAdded
+    }
+"""
+
+# Asynchronous request
+import asyncio
+
+asyncio.run(client.subscribe(query=query, handle=print))
+# => {'data': {'messageAdded': 'Error omnis quis.'}}
+# => {'data': {'messageAdded': 'Enim asperiores omnis.'}}
+# => {'data': {'messageAdded': 'Unde ullam consequatur quam eius vel.'}}
+# ...
+```
+
 ## Roadmap
 
 To start we'll try and use a Github project board for listing current work and updating priorities of upcoming features.

python_graphql_client/graphql_client.py

@@ -1,7 +1,13 @@
 """Module containing graphQL client."""
+import json
+import logging
+from typing import Callable
 
 import aiohttp
 import requests
+import websockets
+
+logging.basicConfig(level=logging.INFO, format="%(levelname)s:%(message)s")
 
 
 class GraphqlClient:
@@ -66,3 +72,33 @@ async def execute_async(
                 headers=self.__request_headers(headers),
             ) as response:
                 return await response.json()
+
+    async def subscribe(
+        self,
+        query: str,
+        handle: Callable,
+        variables: dict = None,
+        operation_name: str = None,
+        headers: dict = None,
+    ):
+        """Make asynchronous request for GraphQL subscription."""
+        connection_init_message = json.dumps({"type": "connection_init", "payload": {}})
+
+        request_body = self.__request_body(
+            query=query, variables=variables, operation_name=operation_name
+        )
+        request_message = json.dumps(
+            {"type": "start", "id": "1", "payload": request_body}
+        )
+
+        async with websockets.connect(
+            self.endpoint, subprotocols=["graphql-ws"]
+        ) as websocket:
+            await websocket.send(connection_init_message)
+            await websocket.send(request_message)
+            async for response_message in websocket:
+                response_body = json.loads(response_message)
+                if response_body["type"] == "connection_ack":
+                    logging.info("the server accepted the connection")
+                else:
+                    handle(response_body["payload"])

setup.py

@@ -29,7 +29,7 @@
     author_email="opensource@prodigygame.com",
     license="MIT",
     packages=["python_graphql_client"],
-    install_requires=["aiohttp==3.6.2", "requests==2.22.0"],
+    install_requires=["aiohttp==3.6.2", "requests==2.22.0", "websockets==8.1"],
     extras_require={
         "dev": [
             "pre-commit",

tests/test_graphql_client.py

@@ -1,7 +1,7 @@
 """Tests for main graphql client module."""
 
 from unittest import IsolatedAsyncioTestCase, TestCase
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, call, patch
 
 from aiohttp import web
 from requests.exceptions import HTTPError
@@ -216,3 +216,35 @@ async def test_execute_query_with_operation_name(self, mock_post):
             json={"query": query, "operationName": operation_name},
             headers={},
         )
+
+
+class TestGraphqlClientSubscriptions(IsolatedAsyncioTestCase):
+    """Test cases for subscribing GraphQL subscriptions."""
+
+    @patch("websockets.connect")
+    async def test_subscribe(self, mock_connect):
+        """Subsribe a GraphQL subscription."""
+        mock_websocket = mock_connect.return_value.__aenter__.return_value
+        mock_websocket.send = AsyncMock()
+        mock_websocket.__aiter__.return_value = [
+            '{"type": "data", "id": "1", "payload": {"data": {"messageAdded": "one"}}}',
+            '{"type": "data", "id": "1", "payload": {"data": {"messageAdded": "two"}}}',
+        ]
+
+        client = GraphqlClient(endpoint="ws://www.test-api.com/graphql")
+        query = """
+            subscription onMessageAdded {
+                messageAdded
+            }
+        """
+
+        mock_handle = MagicMock()
+
+        await client.subscribe(query=query, handle=mock_handle)
+
+        mock_handle.assert_has_calls(
+            [
+                call({"data": {"messageAdded": "one"}}),
+                call({"data": {"messageAdded": "two"}}),
+            ]
+        )