Merge branch 'master' into fix_race_condition_in_websockets_connect

Author: KingDarBoja

gql/client.py

@@ -1,6 +1,5 @@
 import asyncio
-from inspect import isawaitable
-from typing import Any, AsyncGenerator, Dict, Generator, Optional, Union, cast
+from typing import Any, AsyncGenerator, Dict, Generator, Optional, Union
 
 from graphql import (
     DocumentNode,
@@ -196,19 +195,22 @@ class SyncClientSession:
     def __init__(self, client: Client):
         self.client = client
 
-    def execute(self, document: DocumentNode, *args, **kwargs) -> Dict:
+    def _execute(self, document: DocumentNode, *args, **kwargs) -> ExecutionResult:
 
         # Validate document
         if self.client.schema:
             self.client.validate(document)
 
-        result = self.transport.execute(document, *args, **kwargs)
+        return self.transport.execute(document, *args, **kwargs)
+
+    def execute(self, document: DocumentNode, *args, **kwargs) -> Dict:
 
-        assert not isawaitable(result), "Transport returned an awaitable result."
-        result = cast(ExecutionResult, result)
+        # Validate and execute on the transport
+        result = self._execute(document, *args, **kwargs)
 
+        # Raise an error if an error is returned in the ExecutionResult object
         if result.errors:
-            raise TransportQueryError(str(result.errors[0]))
+            raise TransportQueryError(str(result.errors[0]), errors=result.errors)
 
         assert (
             result.data is not None
@@ -250,43 +252,69 @@ async def fetch_and_validate(self, document: DocumentNode):
         if self.client.schema:
             self.client.validate(document)
 
-    async def subscribe(
+    async def _subscribe(
         self, document: DocumentNode, *args, **kwargs
-    ) -> AsyncGenerator[Dict, None]:
+    ) -> AsyncGenerator[ExecutionResult, None]:
 
         # Fetch schema from transport if needed and validate document if possible
         await self.fetch_and_validate(document)
 
-        # Subscribe to the transport and yield data or raise error
-        self._generator: AsyncGenerator[
+        # Subscribe to the transport
+        inner_generator: AsyncGenerator[
             ExecutionResult, None
         ] = self.transport.subscribe(document, *args, **kwargs)
 
-        async for result in self._generator:
+        # Keep a reference to the inner generator to allow the user to call aclose()
+        # before a break if python version is too old (pypy3 py 3.6.1)
+        self._generator = inner_generator
+
+        async for result in inner_generator:
             if result.errors:
                 # Note: we need to run generator.aclose() here or the finally block in
                 # transport.subscribe will not be reached in pypy3 (py 3.6.1)
-                await self._generator.aclose()
+                await inner_generator.aclose()
+
+            yield result
+
+    async def subscribe(
+        self, document: DocumentNode, *args, **kwargs
+    ) -> AsyncGenerator[Dict, None]:
 
-                raise TransportQueryError(str(result.errors[0]))
+        # Validate and subscribe on the transport
+        async for result in self._subscribe(document, *args, **kwargs):
+
+            # Raise an error if an error is returned in the ExecutionResult object
+            if result.errors:
+                raise TransportQueryError(str(result.errors[0]), errors=result.errors)
 
             elif result.data is not None:
                 yield result.data
 
-    async def execute(self, document: DocumentNode, *args, **kwargs) -> Dict:
+    async def _execute(
+        self, document: DocumentNode, *args, **kwargs
+    ) -> ExecutionResult:
 
         # Fetch schema from transport if needed and validate document if possible
         await self.fetch_and_validate(document)
 
         # Execute the query with the transport with a timeout
-        result = await asyncio.wait_for(
+        return await asyncio.wait_for(
             self.transport.execute(document, *args, **kwargs),
             self.client.execute_timeout,
         )
 
+    async def execute(self, document: DocumentNode, *args, **kwargs) -> Dict:
+
+        # Validate and execute on the transport
+        result = await self._execute(document, *args, **kwargs)
+
         # Raise an error if an error is returned in the ExecutionResult object
         if result.errors:
-            raise TransportQueryError(str(result.errors[0]))
+            raise TransportQueryError(str(result.errors[0]), errors=result.errors)
+
+        assert (
+            result.data is not None
+        ), "Transport returned an ExecutionResult without data or errors"
 
         return result.data
 

gql/transport/exceptions.py

@@ -1,3 +1,6 @@
+from typing import Any, List, Optional
+
+
 class TransportError(Exception):
     pass
 
@@ -22,9 +25,15 @@ class TransportQueryError(Exception):
     This exception should not close the transport connection.
     """
 
-    def __init__(self, msg, query_id=None):
+    def __init__(
+        self,
+        msg: str,
+        query_id: Optional[int] = None,
+        errors: Optional[List[Any]] = None,
+    ):
         super().__init__(msg)
         self.query_id = query_id
+        self.errors = errors
 
 
 class TransportClosed(TransportError):

gql/transport/websockets.py

@@ -293,7 +293,9 @@ def _parse_answer(
 
                     elif answer_type == "error":
 
-                        raise TransportQueryError(str(payload), query_id=answer_id)
+                        raise TransportQueryError(
+                            str(payload), query_id=answer_id, errors=[payload]
+                        )
 
             elif answer_type == "ka":
                 # KeepAlive message
@@ -335,6 +337,9 @@ async def _receive_data_loop(self) -> None:
                     # ==> Add an exception to this query queue
                     # The exception is raised for this specific query,
                     # but the transport is not closed.
+                    assert isinstance(
+                        e.query_id, int
+                    ), "TransportQueryError should have a query_id defined here"
                     try:
                         await self.listeners[e.query_id].set_exception(e)
                     except KeyError:

tests/test_websocket_exceptions.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import types
+from typing import List
 
 import pytest
 import websockets
@@ -45,9 +46,17 @@ async def test_websocket_invalid_query(event_loop, client_and_server, query_str)
 
     query = gql(query_str)
 
-    with pytest.raises(TransportQueryError):
+    with pytest.raises(TransportQueryError) as exc_info:
         await session.execute(query)
 
+    exception = exc_info.value
+
+    assert isinstance(exception.errors, List)
+
+    error = exception.errors[0]
+
+    assert error["extensions"]["code"] == "INTERNAL_SERVER_ERROR"
+
 
 invalid_subscription_str = """
     subscription getContinents {
@@ -76,10 +85,18 @@ async def test_websocket_invalid_subscription(event_loop, client_and_server, que
 
     query = gql(query_str)
 
-    with pytest.raises(TransportQueryError):
+    with pytest.raises(TransportQueryError) as exc_info:
         async for result in session.subscribe(query):
             pass
 
+    exception = exc_info.value
+
+    assert isinstance(exception.errors, List)
+
+    error = exception.errors[0]
+
+    assert error["extensions"]["code"] == "INTERNAL_SERVER_ERROR"
+
 
 connection_error_server_answer = (
     '{"type":"connection_error","id":null,'
@@ -171,9 +188,17 @@ async def monkey_patch_send_query(
 
     query = gql(query_str)
 
-    with pytest.raises(TransportQueryError):
+    with pytest.raises(TransportQueryError) as exc_info:
         await session.execute(query)
 
+    exception = exc_info.value
+
+    assert isinstance(exception.errors, List)
+
+    error = exception.errors[0]
+
+    assert error["message"] == "Must provide document"
+
 
 not_json_answer = ["BLAHBLAH"]
 missing_type_answer = ["{}"]