Improve extensions performance & add OpenTelemetry

.github/workflows/benchmark.yml

@@ -29,7 +29,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install wheel
-        pip install -e .[asgi-file-uploads,tracing,test,dev]
+        pip install -e .[asgi-file-uploads,tracing,telemetry,test,dev]
     - name: Benchmark
       run: |
         pytest benchmark --benchmark-storage=file://benchmark/results --benchmark-autosave

.github/workflows/tests.yml

@@ -27,7 +27,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install wheel
-        pip install -e .[asgi-file-uploads,tracing,test,dev]
+        pip install -e .[asgi-file-uploads,tracing,telemetry,test,dev]
     - name: Pytest
       run: |
         pytest --cov=ariadne --cov=tests
@@ -61,7 +61,7 @@ jobs:
         python -m pip install --upgrade pip
         pip install wheel
         pip install -r tests_integrations/${{ matrix.library }}/requirements.txt
-        pip install -e .[asgi-file-uploads,tracing,test,dev]
+        pip install -e .[asgi-file-uploads,tracing,telemetry,test,dev]
     - name: Pytest
       run: |
         pytest tests_integrations/${{ matrix.library }}

CHANGELOG.md

@@ -5,6 +5,10 @@
 - Added `query_validator` option to ASGI and WSGI `GraphQL` applications that enables customization of query validation step.
 - Fixed `ERROR` message in GraphQL-WS protocol having invalid payload type.
 - Fixed query cost validator incorrect handling of inline fragments.
+- Removed `ExtensionSync`. `Extension` now supports both async and sync contexts.
+- Removed `OpenTracingSync` and `opentracing_extension_sync`. `OpenTracing` and `opentracing_extension` now support both async and sync contexts.
+- Removed `ApolloTracingSync`. `ApolloTracing` now supports both async and sync contexts.
+- Added `OpenTelemetry` and `opentelemetry_extension` extension, importable form `ariadne.tracing.opentelemetry`.
 
 
 ## 0.19.1 (2023-03-28)

ariadne/__init__.py

@@ -29,7 +29,7 @@
 from .schema_names import SchemaNameConverter, convert_schema_names
 from .schema_visitor import SchemaDirectiveVisitor
 from .subscriptions import SubscriptionType
-from .types import Extension, ExtensionSync, SchemaBindable
+from .types import Extension, SchemaBindable
 from .unions import UnionType
 from .utils import (
     convert_camel_case_to_snake,
@@ -42,7 +42,6 @@
     "EnumType",
     "Extension",
     "ExtensionManager",
-    "ExtensionSync",
     "FallbackResolversSetter",
     "InputType",
     "InterfaceType",

ariadne/contrib/tracing/apollotracing.py

@@ -1,8 +1,9 @@
 from datetime import datetime
-from inspect import isawaitable
+from inspect import iscoroutinefunction
 from typing import Any, List, Optional, cast
 
 from graphql import GraphQLResolveInfo
+from graphql.pyutils import is_awaitable
 
 from ...types import ContextValue, Extension, Resolver
 from .utils import format_path, should_trace
@@ -35,15 +36,39 @@ def request_started(self, context: ContextValue):
         self.start_date = datetime.utcnow()
         self.start_timestamp = perf_counter_ns()
 
-    async def resolve(
+    def resolve(self, next_: Resolver, obj: Any, info: GraphQLResolveInfo, **kwargs):
+        if not should_trace(info, self.trace_default_resolver):
+            return next_(obj, info, **kwargs)
+
+        if iscoroutinefunction(next_):
+            return self.resolve_async(next_, obj, info, **kwargs)
+
+        return self.resolve_sync(next_, obj, info, **kwargs)
+
+    async def resolve_async(
         self, next_: Resolver, obj: Any, info: GraphQLResolveInfo, **kwargs
     ):
-        if not should_trace(info, self.trace_default_resolver):
+        start_timestamp = perf_counter_ns()
+        record = {
+            "path": format_path(info.path),
+            "parentType": str(info.parent_type),
+            "fieldName": info.field_name,
+            "returnType": str(info.return_type),
+            "startOffset": start_timestamp - cast(int, self.start_timestamp),
+        }
+        self.resolvers.append(record)
+        try:
             result = next_(obj, info, **kwargs)
-            if isawaitable(result):
+            if is_awaitable(result):
                 result = await result
             return result
+        finally:
+            end_timestamp = perf_counter_ns()
+            record["duration"] = end_timestamp - start_timestamp
 
+    def resolve_sync(
+        self, next_: Resolver, obj: Any, info: GraphQLResolveInfo, **kwargs
+    ):
         start_timestamp = perf_counter_ns()
         record = {
             "path": format_path(info.path),
@@ -55,8 +80,6 @@ async def resolve(
         self.resolvers.append(record)
         try:
             result = next_(obj, info, **kwargs)
-            if isawaitable(result):
-                result = await result
             return result
         finally:
             end_timestamp = perf_counter_ns()
@@ -87,28 +110,3 @@ def format(self, context: ContextValue):
                 "execution": {"resolvers": totals["resolvers"]},
             }
         }
-
-
-class ApolloTracingExtensionSync(ApolloTracingExtension):
-    def resolve(
-        self, next_: Resolver, obj: Any, info: GraphQLResolveInfo, **kwargs
-    ):  # pylint: disable=invalid-overridden-method
-        if not should_trace(info):
-            result = next_(obj, info, **kwargs)
-            return result
-
-        start_timestamp = perf_counter_ns()
-        record = {
-            "path": format_path(info.path),
-            "parentType": str(info.parent_type),
-            "fieldName": info.field_name,
-            "returnType": str(info.return_type),
-            "startOffset": start_timestamp - cast(int, self.start_timestamp),
-        }
-        self.resolvers.append(record)
-        try:
-            result = next_(obj, info, **kwargs)
-            return result
-        finally:
-            end_timestamp = perf_counter_ns()
-            record["duration"] = end_timestamp - start_timestamp

ariadne/contrib/tracing/copy_args.py

@@ -0,0 +1,47 @@
+import os
+from typing import Any, Union
+
+from starlette.datastructures import UploadFile
+
+try:
+    from multipart.multipart import File
+except ImportError:
+
+    class File:  # type: ignore
+        """Mock upload file used when python-multipart is not installed."""
+
+
+def copy_args_for_tracing(value: Any) -> Any:
+    if isinstance(value, dict):
+        return {k: copy_args_for_tracing(v) for k, v in value.items()}
+    if isinstance(value, list):
+        return [copy_args_for_tracing(v) for v in value]
+    if isinstance(value, (UploadFile, File)):
+        return repr_upload_file(value)
+    return value
+
+
+def repr_upload_file(upload_file: Union[UploadFile, File]) -> str:
+    if isinstance(upload_file, File):
+        filename = upload_file.file_name
+    else:
+        filename = upload_file.filename
+
+    mime_type: Union[str, None]
+
+    if isinstance(upload_file, File):
+        mime_type = "not/available"
+    else:
+        mime_type = upload_file.content_type
+
+    if isinstance(upload_file, File):
+        size = upload_file.size
+    else:
+        file_ = upload_file.file
+        file_.seek(0, os.SEEK_END)
+        size = file_.tell()
+        file_.seek(0)
+
+    return (
+        f"{type(upload_file)}(mime_type={mime_type}, size={size}, filename={filename})"
+    )

ariadne/contrib/tracing/opentelemetry.py

@@ -0,0 +1,154 @@
+from functools import partial
+from inspect import iscoroutinefunction
+from typing import Any, Callable, Dict, Optional, Union
+
+from graphql import GraphQLResolveInfo
+from graphql.pyutils import is_awaitable
+from opentelemetry.trace import Context, Span, Tracer, get_tracer, set_span_in_context
+
+from ...types import ContextValue, Extension, Resolver
+from .utils import copy_args_for_tracing, format_path, should_trace
+
+
+ArgFilter = Callable[[Dict[str, Any], GraphQLResolveInfo], Dict[str, Any]]
+RootSpanName = Union[str, Callable[[ContextValue], str]]
+
+DEFAULT_OPERATION_NAME = "GraphQL Operation"
+
+
+class OpenTelemetryExtension(Extension):
+    _arg_filter: Optional[ArgFilter]
+    _root_context: Optional[Context]
+    _root_span: Span
+    _root_span_name: Optional[RootSpanName]
+    _tracer: Tracer
+
+    def __init__(
+        self,
+        *,
+        tracer: Optional[Tracer] = None,
+        arg_filter: Optional[ArgFilter] = None,
+        root_context: Optional[Context] = None,
+        root_span_name: Optional[RootSpanName] = None,
+    ) -> None:
+        if tracer:
+            self._tracer = tracer
+        else:
+            self._tracer = get_tracer("ariadne")
+
+        self._arg_filter = arg_filter
+        self._root_context = root_context
+        self._root_span_name = root_span_name
+
+    def request_started(self, context: ContextValue):
+        if self._root_span_name:
+            if callable(self._root_span_name):
+                root_span_name = self._root_span_name(context)
+            else:
+                root_span_name = self._root_span_name
+        else:
+            root_span_name = DEFAULT_OPERATION_NAME
+
+        span_context: Optional[Context] = None
+        if self._root_context:
+            span_context = self._root_context
+
+        root_span = self._tracer.start_span(root_span_name, context=span_context)
+        root_span.set_attribute("component", "GraphQL")
+        self._root_span = root_span
+
+    def request_finished(self, context: ContextValue):
+        self._root_span.end()
+
+    def resolve(
+        self, next_: Resolver, obj: Any, info: GraphQLResolveInfo, **kwargs
+    ) -> Any:
+        if not should_trace(info):
+            return next_(obj, info, **kwargs)
+
+        graphql_path = ".".join(
+            map(str, format_path(info.path))  # pylint: disable=bad-builtin
+        )
+
+        with self._tracer.start_as_current_span(
+            info.field_name, context=set_span_in_context(self._root_span)
+        ) as span:
+            span.set_attribute("component", "GraphQL")
+
+            if info.operation.name:
+                span.set_attribute("graphql.operation.name", info.operation.name.value)
+            else:
+                span.set_attribute("graphql.operation.name", DEFAULT_OPERATION_NAME)
+
+            span.set_attribute("graphql.parentType", info.parent_type.name)
+            span.set_attribute("graphql.path", graphql_path)
+
+            if kwargs:
+                filtered_kwargs = self.filter_resolver_args(kwargs, info)
+                for key, value in filtered_kwargs.items():
+                    span.set_attribute(f"graphql.arg[{key}]", value)
+
+            if iscoroutinefunction(next_):
+                return self.resolve_async(span, next_, obj, info, **kwargs)
+
+            return self.resolve_sync(span, next_, obj, info, **kwargs)
+
+    async def resolve_async(
+        self, span: Span, next_: Resolver, obj: Any, info: GraphQLResolveInfo, **kwargs
+    ) -> Any:
+        with self._tracer.start_as_current_span(
+            "resolve async", context=set_span_in_context(span)
+        ) as child_span:
+            result = next_(obj, info, **kwargs)
+            if is_awaitable(result):
+                with self._tracer.start_as_current_span(
+                    "await result", context=set_span_in_context(child_span)
+                ):
+                    return await result
+            return result
+
+    def resolve_sync(
+        self, span: Span, next_: Resolver, obj: Any, info: GraphQLResolveInfo, **kwargs
+    ) -> Any:
+        with self._tracer.start_as_current_span(
+            "resolve sync", context=set_span_in_context(span)
+        ) as child_span:
+            result = next_(obj, info, **kwargs)
+
+            if is_awaitable(result):
+
+                async def await_sync_result():
+                    with self._tracer.start_as_current_span(
+                        "await result", context=set_span_in_context(child_span)
+                    ):
+                        return await result
+
+                return await_sync_result()
+
+            return result
+
+    def filter_resolver_args(
+        self, args: Dict[str, Any], info: GraphQLResolveInfo
+    ) -> Dict[str, Any]:
+        args_to_trace = copy_args_for_tracing(args)
+
+        if not self._arg_filter:
+            return args_to_trace
+
+        return self._arg_filter(args_to_trace, info)
+
+
+def opentelemetry_extension(
+    *,
+    tracer: Optional[Tracer] = None,
+    arg_filter: Optional[ArgFilter] = None,
+    root_context: Optional[Context] = None,
+    root_span_name: Optional[RootSpanName] = None,
+):
+    return partial(
+        OpenTelemetryExtension,
+        tracer=tracer,
+        arg_filter=arg_filter,
+        root_context=root_context,
+        root_span_name=root_span_name,
+    )