Graphql Tracing Middleware for Apollo Engine.

tracing.py

from graphql.execution import ExecutionResult

from .schema import schema

import datetime
import dateutil.parser
import dateutil.tz

import monotonic


def monotonic_ms():
    return monotonic.monotonic() * 1000000000


class TracingMiddleware(object):
    """
    Middleware adding Apollo Tracing data for Apollo Engine.
    """

    def _update_tracing_extension(self, info, start, monototic_start, end, monototic_end):
        """
        Augment extensions.tracing with the Apollo Tracing metadata.

        This will initialize the global tracing info if it doesn't exist
        then add a execution.resolvers entry for the current info.

        Args:
            info (ResolveInfo): about the traced resolver.
            start (datetime.datetime): when the traced resolver started.
            monototic_start (float): when the traced resolver started.
            end (datetime.datetime): when the traced resolver ended.
            monototic_end (float): when the traced resolver started.

        Returns:
            dict: updated info.extensions that should be added to the response.

        """

        extensions = info.extensions or {}

        # TODO: Figure out how to measure parsing and validation times.

        tracing = extensions.setdefault('tracing', {
            "version": 1,
            "startTime": start.isoformat(),
            "endTime": end.isoformat(),
            "duration": monototic_end - monototic_start,
            "parsing": {
                "startOffset": 0,
                "duration": 0,
            },
            "validation": {
                "startOffset": 0,
                "duration": 0,
            },
            "execution": {
                "resolvers": []
            },
            # The following two fields are for our own use.
            "_systemStartOffset": monototic_start,
            "_systemEndOffset": monototic_end,
        })

        # Since we don't store anywhere else we re-parse what we wrote.
        global_monototic_end = max(monototic_end, tracing['_systemStartOffset'])
        global_monototic_start = min(monototic_start, tracing['_systemEndOffset'])
        global_start = min(start, dateutil.parser.parse(tracing['startTime']))
        global_end = max(end, dateutil.parser.parse(tracing['endTime']))

        tracing.update(dict(
            _systemStartOffset=global_monototic_start,
            _systemEndOffset=global_monototic_end,
            startTime=global_start.isoformat(),
            endTime=global_end.isoformat(),
            duration=global_monototic_end - global_monototic_start,
        ))

        tracing['execution']['resolvers'].append({
            "path": info.path,
            "parentType": str(info.parent_type),
            "fieldName": info.field_name,
            "returnType": str(info.return_type),
            "startOffset": monototic_start - global_monototic_start,
            "duration": monototic_end - monototic_start,
        })

        return extensions

    def resolve(self, next_, root, info, **args):
        """
        Measure the time the resolver takes.

        Args:
            next_ (callable): resolver to call and measure.
            root (object): passed through to the next resolver.
            info (ResolveInfo): info about the field being resolved.

        Returns:
            object: passed through from the next resolver.

        """

        start = datetime.datetime.now(dateutil.tz.tzlocal())
        monototic_start = monotonic_ms()

        try:
            result = next_(root, info, **args)
        finally:
            extensions = self._update_tracing_extension(
                info,
                start=start,
                monototic_start=monototic_start,
                end=datetime.datetime.now(dateutil.tz.tzlocal()),
                monototic_end=monotonic_ms(),
            )

        return ExecutionResult(data=result, extensions=extensions)

On line 118 you will want to check to see if result is Promise.is_thenable and handle the return appropriately if the execution is async. Also, it can no longer be assumed that result from next is the raw data given chained middleware could also return ExecutionResult.

    def handle_resolved_result(result, error=None):
        extensions = self._update_tracing_extension(
                info,
                start=start,
                monototic_start=monototic_start,
                end=datetime.datetime.now(dateutil.tz.tzlocal()),
                monototic_end=monotonic_ms(),
            )

        data = result
        errors = [error] if error else None
        if isinstance(result, ExecutionResult):
            data = result.data
            if result.extensions:
                extensions.update(result.extensions)
            if result.errors:
                errors = result.errors if not errors else result.errors + errors

        return ExecutionResult(data=data, extensions=extensions, errors=errors)
        
    try:
        result = next_(root, info, **args)
        if result and Promise.is_thenable(result):
            return result.then(handle_resolved_result).catch(partial(handle_resolved_result, None))
        return handle_resolved_result(result)
    except Exception as error:
        return handle_resolved_result(None, error)

Something like this but without the inevitable bugs and syntax errors