Handle Subscription Complete

python_graphql_client/graphql_client.py

@@ -11,12 +11,16 @@
 class GraphqlClient:
     """Class which represents the interface to make graphQL requests through."""
 
-    def __init__(self, endpoint: str, headers: dict = {}, **kwargs: Any):
+    def __init__(self, endpoint: str, headers: dict = {}, session=None, **kwargs: Any):
         """Insantiate the client."""
         self.logger = logging.getLogger(__name__)
         self.endpoint = endpoint
         self.headers = headers
         self.options = kwargs
+        self.session = session
+        if self.session is None:
+            self.session = requests.Session()
+
 
     def __request_body(
         self, query: str, variables: dict = None, operation_name: str = None
@@ -44,7 +48,7 @@ def execute(
             query=query, variables=variables, operation_name=operation_name
         )
 
-        result = requests.post(
+        result = self.session.post(
             self.endpoint,
             json=request_body,
             headers={**self.headers, **headers},
@@ -92,12 +96,12 @@ async def subscribe(
             query=query, variables=variables, operation_name=operation_name
         )
         request_message = json.dumps(
-            {"type": "start", "id": "1", "payload": request_body}
+            {"type": "subscribe", "id": "1", "payload": request_body}
         )
 
         async with websockets.connect(
             self.endpoint,
-            subprotocols=["graphql-ws"],
+            subprotocols=["graphql-transport-ws"],
             extra_headers={**self.headers, **headers},
         ) as websocket:
             await websocket.send(connection_init_message)
@@ -108,5 +112,7 @@ async def subscribe(
                     self.logger.info("the server accepted the connection")
                 elif response_body["type"] == "ka":
                     self.logger.info("the server sent a keep alive message")
+                elif response_body["type"] == "complete":
+                    return
                 else:
                     handle(response_body["payload"])

setup.py

@@ -10,7 +10,7 @@
 
 setup(
     name="python_graphql_client",
-    version="0.4.3",
+    version="0.4.4",
     description="Python GraphQL Client",
     long_description=long_description,
     long_description_content_type="text/markdown",