feat(graphql): add async client

nejch · max-wittig · commit 288f39c828eb · 2024-12-04T14:28:01.000+01:00
diff --git a/README.rst b/README.rst
@@ -25,9 +25,10 @@ python-gitlab
 .. image:: https://img.shields.io/github/license/python-gitlab/python-gitlab
    :target: https://github.com/python-gitlab/python-gitlab/blob/main/COPYING
 
-``python-gitlab`` is a Python package providing access to the GitLab server API.
+``python-gitlab`` is a Python package providing access to the GitLab APIs.
 
-It supports the v4 API of GitLab, and provides a CLI tool (``gitlab``).
+It includes a client for GitLab's v4 REST API, synchronous and asynchronous GraphQL API
+clients, as well as a CLI tool (``gitlab``) wrapping REST API endpoints.
 
 .. _features:
 
@@ -39,6 +40,7 @@ Features
 * write Pythonic code to manage your GitLab resources.
 * pass arbitrary parameters to the GitLab API. Simply follow GitLab's docs
   on what parameters are available.
+* use a synchronous or asynchronous client when using the GraphQL API.
 * access arbitrary endpoints as soon as they are available on GitLab, by using
   lower-level API methods.
 * use persistent requests sessions for authentication, proxy and certificate handling.
diff --git a/docs/api-usage-graphql.rst b/docs/api-usage-graphql.rst
@@ -2,7 +2,8 @@
 Using the GraphQL API (beta)
 ############################
 
-python-gitlab provides basic support for executing GraphQL queries and mutations.
+python-gitlab provides basic support for executing GraphQL queries and mutations,
+providing both a synchronous and asynchronous client.
 
 .. danger::
 
@@ -13,10 +14,11 @@ python-gitlab provides basic support for executing GraphQL queries and mutations
    It is currently unstable and its implementation may change. You can expect a more
    mature client in one of the upcoming versions.
 
-The ``gitlab.GraphQL`` class
-==================================
+The ``gitlab.GraphQL`` and ``gitlab.AsyncGraphQL`` classes
+==========================================================
 
-As with the REST client, you connect to a GitLab instance by creating a ``gitlab.GraphQL`` object:
+As with the REST client, you connect to a GitLab instance by creating a ``gitlab.GraphQL``
+(for synchronous code) or ``gitlab.AsyncGraphQL`` instance (for asynchronous code):
 
 .. code-block:: python
 
@@ -34,6 +36,12 @@ As with the REST client, you connect to a GitLab instance by creating a ``gitlab
    # personal access token or OAuth2 token authentication (self-hosted GitLab instance)
    gq = gitlab.GraphQL('https://gitlab.example.com', token='glpat-JVNSESs8EwWRx5yDxM5q')
 
+   # or the async equivalents
+   async_gq = gitlab.AsyncGraphQL()
+   async_gq = gitlab.AsyncGraphQL('https://gitlab.example.com')
+   async_gq = gitlab.AsyncGraphQL(token='glpat-JVNSESs8EwWRx5yDxM5q')
+   async_gq = gitlab.AsyncGraphQL('https://gitlab.example.com', token='glpat-JVNSESs8EwWRx5yDxM5q')
+  
 Sending queries
 ===============
 
@@ -50,3 +58,17 @@ Get the result of a query:
     """
 
     result = gq.execute(query)
+
+Get the result of a query using the async client:
+
+.. code-block:: python
+
+    query = """{
+        query {
+          currentUser {
+            name
+          }
+        }
+    """
+
+    result = await async_gq.execute(query)
diff --git a/gitlab/__init__.py b/gitlab/__init__.py
@@ -27,7 +27,7 @@
     __title__,
     __version__,
 )
-from gitlab.client import Gitlab, GitlabList, GraphQL  # noqa: F401
+from gitlab.client import AsyncGraphQL, Gitlab, GitlabList, GraphQL  # noqa: F401
 from gitlab.exceptions import *  # noqa: F401,F403
 
 warnings.filterwarnings("default", category=DeprecationWarning, module="^gitlab")
@@ -42,6 +42,7 @@
     "__version__",
     "Gitlab",
     "GitlabList",
+    "AsyncGraphQL",
     "GraphQL",
 ]
 __all__.extend(gitlab.exceptions.__all__)
diff --git a/gitlab/_backends/graphql.py b/gitlab/_backends/graphql.py
@@ -1,7 +1,7 @@
 from typing import Any
 
 import httpx
-from gql.transport.httpx import HTTPXTransport
+from gql.transport.httpx import HTTPXAsyncTransport, HTTPXTransport
 
 
 class GitlabTransport(HTTPXTransport):
@@ -22,3 +22,23 @@ def connect(self) -> None:
 
     def close(self) -> None:
         pass
+
+
+class GitlabAsyncTransport(HTTPXAsyncTransport):
+    """An async gql httpx transport that reuses an existing httpx.AsyncClient.
+    By default, gql's transports do not have a keep-alive session
+    and do not enable providing your own session that's kept open.
+    This transport lets us provide and close our session on our own
+    and provide additional auth.
+    For details, see https://github.com/graphql-python/gql/issues/91.
+    """
+
+    def __init__(self, *args: Any, client: httpx.AsyncClient, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        self.client = client
+
+    async def connect(self) -> None:
+        pass
+
+    async def close(self) -> None:
+        pass
diff --git a/gitlab/client.py b/gitlab/client.py
@@ -32,7 +32,7 @@
     import graphql
     import httpx
 
-    from ._backends.graphql import GitlabTransport
+    from ._backends.graphql import GitlabAsyncTransport, GitlabTransport
 
     _GQL_INSTALLED = True
 except ImportError:  # pragma: no cover
@@ -1278,14 +1278,13 @@ def next(self) -> Dict[str, Any]:
         raise StopIteration
 
 
-class GraphQL:
+class _BaseGraphQL:
     def __init__(
         self,
         url: Optional[str] = None,
         *,
         token: Optional[str] = None,
         ssl_verify: Union[bool, str] = True,
-        client: Optional[httpx.Client] = None,
         timeout: Optional[float] = None,
         user_agent: str = gitlab.const.USER_AGENT,
         fetch_schema_from_transport: bool = False,
@@ -1308,9 +1307,50 @@ def __init__(
         self._max_retries = max_retries
         self._obey_rate_limit = obey_rate_limit
         self._retry_transient_errors = retry_transient_errors
+        self._client_opts = self._get_client_opts()
+        self._fetch_schema_from_transport = fetch_schema_from_transport
+
+    def _get_client_opts(self) -> Dict[str, Any]:
+        headers = {"User-Agent": self._user_agent}
+
+        if self._token:
+            headers["Authorization"] = f"Bearer {self._token}"
+
+        return {
+            "headers": headers,
+            "timeout": self._timeout,
+            "verify": self._ssl_verify,
+        }
+
 
-        opts = self._get_client_opts()
-        self._http_client = client or httpx.Client(**opts)
+class GraphQL(_BaseGraphQL):
+    def __init__(
+        self,
+        url: Optional[str] = None,
+        *,
+        token: Optional[str] = None,
+        ssl_verify: Union[bool, str] = True,
+        client: Optional[httpx.Client] = None,
+        timeout: Optional[float] = None,
+        user_agent: str = gitlab.const.USER_AGENT,
+        fetch_schema_from_transport: bool = False,
+        max_retries: int = 10,
+        obey_rate_limit: bool = True,
+        retry_transient_errors: bool = False,
+    ) -> None:
+        super().__init__(
+            url=url,
+            token=token,
+            ssl_verify=ssl_verify,
+            timeout=timeout,
+            user_agent=user_agent,
+            fetch_schema_from_transport=fetch_schema_from_transport,
+            max_retries=max_retries,
+            obey_rate_limit=obey_rate_limit,
+            retry_transient_errors=retry_transient_errors,
+        )
+
+        self._http_client = client or httpx.Client(**self._client_opts)
         self._transport = GitlabTransport(self._url, client=self._http_client)
         self._client = gql.Client(
             transport=self._transport,
@@ -1324,19 +1364,81 @@ def __enter__(self) -> "GraphQL":
     def __exit__(self, *args: Any) -> None:
         self._http_client.close()
 
-    def _get_client_opts(self) -> Dict[str, Any]:
-        headers = {"User-Agent": self._user_agent}
+    def execute(
+        self, request: Union[str, graphql.Source], *args: Any, **kwargs: Any
+    ) -> Any:
+        parsed_document = self._gql(request)
+        retry = utils.Retry(
+            max_retries=self._max_retries,
+            obey_rate_limit=self._obey_rate_limit,
+            retry_transient_errors=self._retry_transient_errors,
+        )
 
-        if self._token:
-            headers["Authorization"] = f"Bearer {self._token}"
+        while True:
+            try:
+                result = self._client.execute(parsed_document, *args, **kwargs)
+            except gql.transport.exceptions.TransportServerError as e:
+                if retry.handle_retry_on_status(
+                    status_code=e.code, headers=self._transport.response_headers
+                ):
+                    continue
 
-        return {
-            "headers": headers,
-            "timeout": self._timeout,
-            "verify": self._ssl_verify,
-        }
+                if e.code == 401:
+                    raise gitlab.exceptions.GitlabAuthenticationError(
+                        response_code=e.code,
+                        error_message=str(e),
+                    )
 
-    def execute(
+                raise gitlab.exceptions.GitlabHttpError(
+                    response_code=e.code,
+                    error_message=str(e),
+                )
+
+            return result
+
+
+class AsyncGraphQL(_BaseGraphQL):
+    def __init__(
+        self,
+        url: Optional[str] = None,
+        *,
+        token: Optional[str] = None,
+        ssl_verify: Union[bool, str] = True,
+        client: Optional[httpx.AsyncClient] = None,
+        timeout: Optional[float] = None,
+        user_agent: str = gitlab.const.USER_AGENT,
+        fetch_schema_from_transport: bool = False,
+        max_retries: int = 10,
+        obey_rate_limit: bool = True,
+        retry_transient_errors: bool = False,
+    ) -> None:
+        super().__init__(
+            url=url,
+            token=token,
+            ssl_verify=ssl_verify,
+            timeout=timeout,
+            user_agent=user_agent,
+            fetch_schema_from_transport=fetch_schema_from_transport,
+            max_retries=max_retries,
+            obey_rate_limit=obey_rate_limit,
+            retry_transient_errors=retry_transient_errors,
+        )
+
+        self._http_client = client or httpx.AsyncClient(**self._client_opts)
+        self._transport = GitlabAsyncTransport(self._url, client=self._http_client)
+        self._client = gql.Client(
+            transport=self._transport,
+            fetch_schema_from_transport=fetch_schema_from_transport,
+        )
+        self._gql = gql.gql
+
+    async def __aenter__(self) -> "AsyncGraphQL":
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        await self._http_client.aclose()
+
+    async def execute(
         self, request: Union[str, graphql.Source], *args: Any, **kwargs: Any
     ) -> Any:
         parsed_document = self._gql(request)
@@ -1348,7 +1450,9 @@ def execute(
 
         while True:
             try:
-                result = self._client.execute(parsed_document, *args, **kwargs)
+                result = await self._client.execute_async(
+                    parsed_document, *args, **kwargs
+                )
             except gql.transport.exceptions.TransportServerError as e:
                 if retry.handle_retry_on_status(
                     status_code=e.code, headers=self._transport.response_headers
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "python-gitlab"
-description="A python wrapper for the GitLab API"
+description="The python wrapper for the GitLab REST and GraphQL APIs."
 readme = "README.rst"
 authors = [
     {name = "Gauvain Pocentek", email= "gauvain@pocentek.net"}
diff --git a/requirements-test.txt b/requirements-test.txt
@@ -1,4 +1,5 @@
 -r requirements.txt
+anyio==4.6.2.post1
 build==1.2.2.post1
 coverage==7.6.8
 pytest-console-scripts==1.4.1
@@ -8,4 +9,5 @@ pytest==8.3.4
 PyYaml==6.0.2
 responses==0.25.3
 respx==0.21.1
+trio==0.27.0
 wheel==0.45.1
diff --git a/tests/functional/api/test_graphql.py b/tests/functional/api/test_graphql.py
@@ -1,20 +1,28 @@
-import logging
-
 import pytest
 
 import gitlab
 
 
 @pytest.fixture
 def gl_gql(gitlab_url: str, gitlab_token: str) -> gitlab.GraphQL:
-    logging.info("Instantiating gitlab.GraphQL instance")
-    instance = gitlab.GraphQL(gitlab_url, token=gitlab_token)
+    return gitlab.GraphQL(gitlab_url, token=gitlab_token)
+
 
-    return instance
+@pytest.fixture
+def gl_async_gql(gitlab_url: str, gitlab_token: str) -> gitlab.AsyncGraphQL:
+    return gitlab.AsyncGraphQL(gitlab_url, token=gitlab_token)
 
 
 def test_query_returns_valid_response(gl_gql: gitlab.GraphQL):
     query = "query {currentUser {active}}"
 
     response = gl_gql.execute(query)
     assert response["currentUser"]["active"] is True
+
+
+@pytest.mark.anyio
+async def test_async_query_returns_valid_response(gl_async_gql: gitlab.AsyncGraphQL):
+    query = "query {currentUser {active}}"
+
+    response = await gl_async_gql.execute(query)
+    assert response["currentUser"]["active"] is True
diff --git a/tests/unit/test_graphql.py b/tests/unit/test_graphql.py
