Merge pull request #227 from mirumee/opentelemetry

Add opentelemetry to base clients

.github/workflows/tests.yml

@@ -25,7 +25,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install wheel
-        pip install -e .[subscriptions,dev]
+        pip install -e .[subscriptions,opentelemetry,dev]
     - name: Pytest
       run: |
         pytest

CHANGELOG.md

@@ -9,6 +9,7 @@
 - Fixed parsing of unions and interfaces to always add `__typename` to generated result models.
 - Added escaping of enum values which are Python keywords by appending `_` to them.
 - Fixed `enums_module_name` option not being passed to generators.
+- Added additional base clients supporting the Open Telemetry tracing. Added `opentelemetry_client` config option.
 
 
 ## 0.9.0 (2023-09-11)

README.md

@@ -76,6 +76,7 @@ Optional settings:
 - `include_comments` (defaults to `"stable"`) - option which sets content of comments included at the top of every generated file. Valid choices are: `"none"` (no comments), `"timestamp"` (comment with generation timestamp), `"stable"` (comment contains a message that this is a generated file)
 - `convert_to_snake_case` (defaults to `true`) - a flag that specifies whether to convert fields and arguments names to snake case
 - `async_client` (defaults to `true`) - default generated client is `async`, change this to option `false` to generate synchronous client instead
+- `opentelemetry_client` (defaults to `false`) - default base clients don't support any performance tracing. Change this option to `true` to use the base client with Open Telemetry support.
 - `files_to_include` (defaults to `[]`) - list of files which will be copied into generated package
 - `plugins` (defaults to `[]`) - list of plugins to use during generation
 
@@ -142,6 +143,20 @@ type = "Upload"
 ```
 
 
+### Open Telemetry
+
+When config option `opentelemetry_client` is set to `true` then default, included base client is replaced with one that implements the opt-in Open Telemetry support. By default this support does nothing but when the `opentelemetry-api` package is installed and the `tracer` argument is provided then the client will create spans with data about performed requests.
+
+Tracing arguments handled by `BaseClientOpenTelemetry`:
+- `tracer`: `Optional[Union[str, Tracer]] = None` - tracer object or name which will be passed to the `get_tracer` method
+- `root_context`: `Optional[Context] = None` - optional context added to root span
+- `root_span_name`: `str = "GraphQL Operation"` - name of root span
+
+`AsyncBaseClientOpenTelemetry` supports all arguments which `BaseClientOpenTelemetry` does, but also exposes additional arguments regarding websockets:
+- `ws_root_context`: `Optional[Context] = None` - optional context added to root span for websocket connection
+- `ws_root_span_name`: `str = "GraphQL Subscription"` - name of root span for websocket connection
+
+
 ## Custom scalars
 
 By default, not built-in scalars are represented as `typing.Any` in generated client.

ariadne_codegen/client_generators/constants.py

@@ -68,10 +68,25 @@
 SKIP_DIRECTIVE_NAME = "skip"
 INCLUDE_DIRECTIVE_NAME = "include"
 
+
 DEFAULT_ASYNC_BASE_CLIENT_PATH = (
     Path(__file__).parent / "dependencies" / "async_base_client.py"
 )
+DEFAULT_ASYNC_BASE_CLIENT_NAME = "AsyncBaseClient"
+
+DEFAULT_ASYNC_BASE_CLIENT_OPEN_TELEMETRY_PATH = (
+    Path(__file__).parent / "dependencies" / "async_base_client_open_telemetry.py"
+)
+DEFAULT_ASYNC_BASE_CLIENT_OPEN_TELEMETRY_NAME = "AsyncBaseClientOpenTelemetry"
+
 DEFAULT_BASE_CLIENT_PATH = Path(__file__).parent / "dependencies" / "base_client.py"
+DEFAULT_BASE_CLIENT_NAME = "BaseClient"
+
+DEFAULT_BASE_CLIENT_OPEN_TELEMETRY_PATH = (
+    Path(__file__).parent / "dependencies" / "base_client_open_telemetry.py"
+)
+DEFAULT_BASE_CLIENT_OPEN_TELEMETRY_NAME = "BaseClientOpenTelemetry"
+
 
 GRAPHQL_CLIENT_EXCEPTIONS_NAMES = [
     "GraphQLClientError",

ariadne_codegen/client_generators/dependencies/async_base_client.py

@@ -29,7 +29,9 @@ async def ws_connect(*args, **kwargs):  # pylint: disable=unused-argument
     WebSocketClientProtocol = Any  # type: ignore
     Data = Any  # type: ignore
     Origin = Any  # type: ignore
-    Subprotocol = Any  # type: ignore
+
+    def Subprotocol(*args, **kwargs):  # type: ignore # pylint: disable=invalid-name
+        raise NotImplementedError("Subscriptions require 'websockets' package.")
 
 
 Self = TypeVar("Self", bound="AsyncBaseClient")
@@ -64,6 +66,7 @@ def __init__(
         self.http_client = (
             http_client if http_client else httpx.AsyncClient(headers=headers)
         )
+
         self.ws_url = ws_url
         self.ws_headers = ws_headers or {}
         self.ws_origin = Origin(ws_origin) if ws_origin else None
@@ -84,16 +87,16 @@ async def execute(
         self, query: str, variables: Optional[Dict[str, Any]] = None
     ) -> httpx.Response:
         processed_variables, files, files_map = self._process_variables(variables)
-        payload: Dict[str, Any] = {"query": query, "variables": processed_variables}
 
         if files and files_map:
             return await self._execute_multipart(
-                payload=payload,
+                query=query,
+                variables=processed_variables,
                 files=files,
                 files_map=files_map,
             )
 
-        return await self._execute_json(payload=payload)
+        return await self._execute_json(query=query, variables=processed_variables)
 
     def get_data(self, response: httpx.Response) -> Dict[str, Any]:
         if not response.is_success:
@@ -213,21 +216,31 @@ def separate_files(path: str, obj: Any) -> Any:
 
     async def _execute_multipart(
         self,
-        payload: Dict[str, Any],
+        query: str,
+        variables: Dict[str, Any],
         files: Dict[str, Tuple[str, IO[bytes], str]],
         files_map: Dict[str, List[str]],
     ) -> httpx.Response:
         data = {
-            "operations": json.dumps(payload, default=to_jsonable_python),
+            "operations": json.dumps(
+                {"query": query, "variables": variables}, default=to_jsonable_python
+            ),
             "map": json.dumps(files_map, default=to_jsonable_python),
         }
 
         return await self.http_client.post(url=self.url, data=data, files=files)
 
-    async def _execute_json(self, payload: Dict[str, Any]) -> httpx.Response:
-        content = json.dumps(payload, default=to_jsonable_python)
+    async def _execute_json(
+        self,
+        query: str,
+        variables: Dict[str, Any],
+    ) -> httpx.Response:
         return await self.http_client.post(
-            url=self.url, content=content, headers={"Content-Type": "application/json"}
+            url=self.url,
+            content=json.dumps(
+                {"query": query, "variables": variables}, default=to_jsonable_python
+            ),
+            headers={"Content-Type": "application/json"},
         )
 
     async def _send_connection_init(self, websocket: WebSocketClientProtocol) -> None: