Merge pull request #254 from jschlyter/graphql_async_client_session

Use GraphQL AsyncClientSession with automatic retries

custom_components/polestar_api/pypolestar/const.py

@@ -7,6 +7,9 @@
 HTTPX_TIMEOUT = 30
 TOKEN_REFRESH_WINDOW_MIN = 300
 
+GRAPHQL_CONNECT_RETRIES = 5
+GRAPHQL_EXECUTE_RETRIES = 3
+
 OIDC_PROVIDER_BASE_URL = "https://polestarid.eu.polestar.com"
 OIDC_REDIRECT_URI = "https://www.polestar.com/sign-in-callback"
 OIDC_CLIENT_ID = "l3oopkc_10"

custom_components/polestar_api/pypolestar/graphql.py

@@ -1,8 +1,11 @@
+import backoff
 import httpx
-from gql import Client, gql
+from gql import gql
+from gql.client import AsyncClientSession, Client
+from gql.transport.exceptions import TransportError, TransportQueryError
 from gql.transport.httpx import HTTPXAsyncTransport
 
-from .const import HTTPX_TIMEOUT
+from .const import GRAPHQL_CONNECT_RETRIES, GRAPHQL_EXECUTE_RETRIES, HTTPX_TIMEOUT
 
 
 class _HTTPXAsyncTransport(HTTPXAsyncTransport):
@@ -21,6 +24,7 @@ async def close(self):
 
 
 def get_gql_client(client: httpx.AsyncClient, url: str) -> Client:
+    """Get GraphQL Client using existing httpx AsyncClient"""
     transport = _HTTPXAsyncTransport(url=url, client=client)
     return Client(
         transport=transport,
@@ -29,6 +33,26 @@ def get_gql_client(client: httpx.AsyncClient, url: str) -> Client:
     )
 
 
+async def get_gql_session(client: Client) -> AsyncClientSession:
+    """Get GraphQL Session with automatic retries"""
+    retry_connect = backoff.on_exception(
+        wait_gen=backoff.expo,
+        exception=(TransportError, httpx.TransportError),
+        max_tries=GRAPHQL_CONNECT_RETRIES,
+    )
+    retry_execute = backoff.on_exception(
+        wait_gen=backoff.expo,
+        exception=(TransportError,),
+        max_tries=GRAPHQL_EXECUTE_RETRIES,
+        giveup=lambda e: isinstance(e, TransportQueryError),
+    )
+    return await client.connect_async(
+        reconnecting=True,
+        retry_connect=retry_connect,
+        retry_execute=retry_execute,
+    )
+
+
 QUERY_GET_AUTH_TOKEN = gql(
     """
     query getAuthToken($code: String!) {

custom_components/polestar_api/pypolestar/polestar.py

@@ -7,6 +7,7 @@
 from datetime import datetime, timedelta
 
 import httpx
+from gql.client import AsyncClientSession
 from gql.transport.exceptions import TransportQueryError
 from graphql import DocumentNode
 
@@ -24,6 +25,7 @@
     QUERY_GET_CONSUMER_CARS_V2_VERBOSE,
     QUERY_GET_ODOMETER_DATA,
     get_gql_client,
+    get_gql_session,
 )
 from .models import CarBatteryData, CarInformationData, CarOdometerData
 
@@ -54,16 +56,20 @@ def __init__(
         self.logger = _LOGGER.getChild(unique_id) if unique_id else _LOGGER
         self.api_url = API_MYSTAR_V2_URL
         self.gql_client = get_gql_client(url=self.api_url, client=self.client_session)
+        self.gql_session: AsyncClientSession | None = None
 
     async def async_init(self, verbose: bool = False) -> None:
         """Initialize the Polestar API."""
+
         await self.auth.async_init()
         await self.auth.get_token()
 
         if self.auth.access_token is None:
             self.logger.warning("No access token %s", self.username)
             return
 
+        self.gql_session = await get_gql_session(self.gql_client)
+
         if not (car_data := await self._get_vehicle_data(verbose=verbose)):
             self.logger.warning("No cars found for %s", self.username)
             return
@@ -276,34 +282,33 @@ async def _query_graph_ql(
         operation_name: str | None = None,
         variable_values: dict | None = None,
     ):
+        if self.gql_session is None:
+            raise RuntimeError("GraphQL not connected")
+
         self.logger.debug("GraphQL URL: %s", self.api_url)
 
-        async with self.gql_client as client:
-            try:
-                result = await client.execute(
-                    query,
-                    operation_name=operation_name,
-                    variable_values=variable_values,
-                    extra_args={
-                        "headers": {"Authorization": f"Bearer {self.auth.access_token}"}
-                    },
-                )
-            except TransportQueryError as exc:
-                self.logger.debug("GraphQL TransportQueryError: %s", str(exc))
-                if (
-                    exc.errors
-                    and exc.errors[0].get("extensions", {}).get("code")
-                    == "UNAUTHENTICATED"
-                ):
-                    self.latest_call_code = 401
-                    raise PolestarNotAuthorizedException(
-                        exc.errors[0]["message"]
-                    ) from exc
-                self.latest_call_code = 500
-                raise PolestarApiException from exc
-            except Exception as exc:
-                self.logger.debug("GraphQL Exception: %s", str(exc))
-                raise exc
+        try:
+            result = await self.gql_session.execute(
+                query,
+                operation_name=operation_name,
+                variable_values=variable_values,
+                extra_args={
+                    "headers": {"Authorization": f"Bearer {self.auth.access_token}"}
+                },
+            )
+        except TransportQueryError as exc:
+            self.logger.debug("GraphQL TransportQueryError: %s", str(exc))
+            if (
+                exc.errors
+                and exc.errors[0].get("extensions", {}).get("code") == "UNAUTHENTICATED"
+            ):
+                self.latest_call_code = 401
+                raise PolestarNotAuthorizedException(exc.errors[0]["message"]) from exc
+            self.latest_call_code = 500
+            raise PolestarApiException from exc
+        except Exception as exc:
+            self.logger.debug("GraphQL Exception: %s", str(exc))
+            raise exc
 
         self.logger.debug("GraphQL Result: %s", result)
         self.latest_call_code = 200