Feature get_execution_result argument of execute and subscribe (#257)

Author: leszekhanusz

docs/usage/extensions.rst

@@ -0,0 +1,36 @@
+.. _extensions:
+
+Extensions
+----------
+
+When you execute (or subscribe) GraphQL requests, the server will send
+responses which may have 3 fields:
+
+- data: the serialized response from the backend
+- errors: a list of potential errors
+- extensions: an optional field for additional data
+
+If there are errors in the response, then the
+:code:`execute` or :code:`subscribe` methods will
+raise a :code:`TransportQueryError`.
+
+If no errors are present, then only the data from the response is returned by default.
+
+.. code-block:: python
+
+    result = client.execute(query)
+    # result is here the content of the data field
+
+If you need to receive the extensions data too, then you can run the
+:code:`execute` or :code:`subscribe` methods with :code:`get_execution_result=True`.
+
+In that case, the full execution result is returned and you can have access
+to the extensions field
+
+.. code-block:: python
+
+    result = client.execute(query, get_execution_result=True)
+    # result is here an ExecutionResult instance
+
+    # result.data is the content of the data field
+    # result.extensions is the content of the extensions field

docs/usage/index.rst

@@ -11,3 +11,4 @@ Usage
    headers
    file_upload
    custom_scalars_and_enums
+   extensions

gql/client.py

@@ -367,8 +367,9 @@ def execute(
         operation_name: Optional[str] = None,
         serialize_variables: Optional[bool] = None,
         parse_result: Optional[bool] = None,
+        get_execution_result: bool = False,
         **kwargs,
-    ) -> Dict:
+    ) -> Union[Dict[str, Any], ExecutionResult]:
         """Execute the provided document AST synchronously using
         the sync transport.
 
@@ -382,6 +383,8 @@ def execute(
             serialized. Used for custom scalars and/or enums. Default: False.
         :param parse_result: Whether gql will unserialize the result.
             By default use the parse_results attribute of the client.
+        :param get_execution_result: return the full ExecutionResult instance instead of
+            only the "data" field. Necessary if you want to get the "extensions" field.
 
         The extra arguments are passed to the transport execute method."""
 
@@ -399,13 +402,19 @@ def execute(
         # Raise an error if an error is returned in the ExecutionResult object
         if result.errors:
             raise TransportQueryError(
-                str(result.errors[0]), errors=result.errors, data=result.data
+                str(result.errors[0]),
+                errors=result.errors,
+                data=result.data,
+                extensions=result.extensions,
             )
 
         assert (
             result.data is not None
         ), "Transport returned an ExecutionResult without data or errors"
 
+        if get_execution_result:
+            return result
+
         return result.data
 
     def fetch_schema(self) -> None:
@@ -519,8 +528,9 @@ async def subscribe(
         operation_name: Optional[str] = None,
         serialize_variables: Optional[bool] = None,
         parse_result: Optional[bool] = None,
+        get_execution_result: bool = False,
         **kwargs,
-    ) -> AsyncGenerator[Dict, None]:
+    ) -> AsyncGenerator[Union[Dict[str, Any], ExecutionResult], None]:
         """Coroutine to subscribe asynchronously to the provided document AST
         asynchronously using the async transport.
 
@@ -534,6 +544,8 @@ async def subscribe(
             serialized. Used for custom scalars and/or enums. Default: False.
         :param parse_result: Whether gql will unserialize the result.
             By default use the parse_results attribute of the client.
+        :param get_execution_result: yield the full ExecutionResult instance instead of
+            only the "data" field. Necessary if you want to get the "extensions" field.
 
         The extra arguments are passed to the transport subscribe method."""
 
@@ -554,11 +566,17 @@ async def subscribe(
                 # Raise an error if an error is returned in the ExecutionResult object
                 if result.errors:
                     raise TransportQueryError(
-                        str(result.errors[0]), errors=result.errors, data=result.data
+                        str(result.errors[0]),
+                        errors=result.errors,
+                        data=result.data,
+                        extensions=result.extensions,
                     )
 
                 elif result.data is not None:
-                    yield result.data
+                    if get_execution_result:
+                        yield result
+                    else:
+                        yield result.data
         finally:
             await inner_generator.aclose()
 
@@ -636,8 +654,9 @@ async def execute(
         operation_name: Optional[str] = None,
         serialize_variables: Optional[bool] = None,
         parse_result: Optional[bool] = None,
+        get_execution_result: bool = False,
         **kwargs,
-    ) -> Dict:
+    ) -> Union[Dict[str, Any], ExecutionResult]:
         """Coroutine to execute the provided document AST asynchronously using
         the async transport.
 
@@ -651,6 +670,8 @@ async def execute(
             serialized. Used for custom scalars and/or enums. Default: False.
         :param parse_result: Whether gql will unserialize the result.
             By default use the parse_results attribute of the client.
+        :param get_execution_result: return the full ExecutionResult instance instead of
+            only the "data" field. Necessary if you want to get the "extensions" field.
 
         The extra arguments are passed to the transport execute method."""
 
@@ -668,13 +689,19 @@ async def execute(
         # Raise an error if an error is returned in the ExecutionResult object
         if result.errors:
             raise TransportQueryError(
-                str(result.errors[0]), errors=result.errors, data=result.data
+                str(result.errors[0]),
+                errors=result.errors,
+                data=result.data,
+                extensions=result.extensions,
             )
 
         assert (
             result.data is not None
         ), "Transport returned an ExecutionResult without data or errors"
 
+        if get_execution_result:
+            return result
+
         return result.data
 
     async def fetch_schema(self) -> None:

gql/transport/exceptions.py

@@ -35,11 +35,13 @@ def __init__(
         query_id: Optional[int] = None,
         errors: Optional[List[Any]] = None,
         data: Optional[Any] = None,
+        extensions: Optional[Any] = None,
     ):
         super().__init__(msg)
         self.query_id = query_id
         self.errors = errors
         self.data = data
+        self.extensions = extensions
 
 
 class TransportClosed(TransportError):

tests/test_aiohttp.py

@@ -1070,6 +1070,6 @@ async def handler(request):
 
         query = gql(query1_str)
 
-        execution_result = await session._execute(query)
+        execution_result = await session.execute(query, get_execution_result=True)
 
         assert execution_result.extensions["key1"] == "val1"

tests/test_requests.py

@@ -328,7 +328,7 @@ def test_code():
 
             query = gql(query1_str)
 
-            execution_result = session._execute(query)
+            execution_result = session.execute(query, get_execution_result=True)
 
             assert execution_result.extensions["key1"] == "val1"
 

tests/test_websocket_query.py

@@ -596,6 +596,6 @@ async def test_websocket_simple_query_with_extensions(
 
     query = gql(query_str)
 
-    execution_result = await session._execute(query)
+    execution_result = await session.execute(query, get_execution_result=True)
 
     assert execution_result.extensions["key1"] == "val1"

tests/test_websocket_subscription.py

@@ -4,6 +4,7 @@
 from typing import List
 
 import pytest
+from graphql import ExecutionResult
 from parse import search
 
 from gql import Client, gql
@@ -142,6 +143,31 @@ async def test_websocket_subscription(event_loop, client_and_server, subscriptio
     assert count == -1
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("server", [server_countdown], indirect=True)
+@pytest.mark.parametrize("subscription_str", [countdown_subscription_str])
+async def test_websocket_subscription_get_execution_result(
+    event_loop, client_and_server, subscription_str
+):
+
+    session, server = client_and_server
+
+    count = 10
+    subscription = gql(subscription_str.format(count=count))
+
+    async for result in session.subscribe(subscription, get_execution_result=True):
+
+        assert isinstance(result, ExecutionResult)
+
+        number = result.data["number"]
+        print(f"Number received: {number}")
+
+        assert number == count
+        count -= 1
+
+    assert count == -1
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("server", [server_countdown], indirect=True)
 @pytest.mark.parametrize("subscription_str", [countdown_subscription_str])