diff --git a/README.md b/README.md index 83e059d..fecd4e5 100644 --- a/README.md +++ b/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. diff --git a/python_graphql_client/graphql_client.py b/python_graphql_client/graphql_client.py index 9a14db6..2787971 100644 --- a/python_graphql_client/graphql_client.py +++ b/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"]) diff --git a/setup.py b/setup.py index 8710511..f2a73de 100644 --- a/setup.py +++ b/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", diff --git a/tests/test_graphql_client.py b/tests/test_graphql_client.py index ff2a0dd..a35980e 100644 --- a/tests/test_graphql_client.py +++ b/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"}}), + ] + )