Merge branch 'master' into custom-scalar-types

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,
@@ -204,19 +203,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
@@ -261,45 +263,71 @@ 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:
                 if self.client.type_adapter:
                     result.data = self.client.type_adapter.convert_scalars(result.data)
                 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"
 
         if self.client.type_adapter:
             result.data = self.client.type_adapter.convert_scalars(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

@@ -134,6 +134,8 @@ def __init__(
         self._no_more_listeners: asyncio.Event = asyncio.Event()
         self._no_more_listeners.set()
 
+        self._connecting: bool = False
+
         self.close_exception: Optional[Exception] = None
 
     async def _send(self, message: str) -> None:
@@ -291,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
@@ -333,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:
@@ -467,7 +474,11 @@ async def connect(self) -> None:
 
         GRAPHQLWS_SUBPROTOCOL: Subprotocol = cast(Subprotocol, "graphql-ws")
 
-        if self.websocket is None:
+        if self.websocket is None and not self._connecting:
+
+            # Set connecting to True to avoid a race condition if user is trying
+            # to connect twice using the same client at the same time
+            self._connecting = True
 
             # If the ssl parameter is not provided,
             # generate the ssl value depending on the url
@@ -489,9 +500,13 @@ async def connect(self) -> None:
 
             # Connection to the specified url
             # Generate a TimeoutError if taking more than connect_timeout seconds
-            self.websocket = await asyncio.wait_for(
-                websockets.connect(self.url, **connect_args,), self.connect_timeout,
-            )
+            # Set the _connecting flag to False after in all cases
+            try:
+                self.websocket = await asyncio.wait_for(
+                    websockets.connect(self.url, **connect_args,), self.connect_timeout,
+                )
+            finally:
+                self._connecting = False
 
             self.next_query_id = 1
             self.close_exception = None

tests/test_websocket_exceptions.py

@@ -1,12 +1,14 @@
 import asyncio
 import json
 import types
+from typing import List
 
 import pytest
 import websockets
 
 from gql import Client, gql
 from gql.transport.exceptions import (
+    TransportAlreadyConnected,
     TransportClosed,
     TransportProtocolError,
     TransportQueryError,
@@ -44,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 {
@@ -75,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,'
@@ -170,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 = ["{}"]
@@ -294,3 +320,31 @@ async def test_websocket_server_sending_invalid_query_errors(event_loop, server)
     # Invalid server message is ignored
     async with Client(transport=sample_transport):
         await asyncio.sleep(2 * MS)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("server", [server_sending_invalid_query_errors], indirect=True)
+async def test_websocket_non_regression_bug_105(event_loop, server):
+
+    # This test will check a fix to a race condition which happens if the user is trying
+    # to connect using the same client twice at the same time
+    # See bug #105
+
+    url = f"ws://{server.hostname}:{server.port}/graphql"
+    print(f"url = {url}")
+
+    sample_transport = WebsocketsTransport(url=url)
+
+    client = Client(transport=sample_transport)
+
+    # Create a coroutine which start the connection with the transport but does nothing
+    async def client_connect(client):
+        async with client:
+            await asyncio.sleep(2 * MS)
+
+    # Create two tasks which will try to connect using the same client (not allowed)
+    connect_task1 = asyncio.ensure_future(client_connect(client))
+    connect_task2 = asyncio.ensure_future(client_connect(client))
+
+    with pytest.raises(TransportAlreadyConnected):
+        await asyncio.gather(connect_task1, connect_task2)