Skip to content

Instantly share code, notes, and snippets.

@wapiflapi
Created September 12, 2018 08:51
Show Gist options
  • Save wapiflapi/6131ed5e8e5dfd5aa0b0a638f86edf16 to your computer and use it in GitHub Desktop.
Save wapiflapi/6131ed5e8e5dfd5aa0b0a638f86edf16 to your computer and use it in GitHub Desktop.
Graphql Tracing Middleware for Apollo Engine.
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)
@nikordaris
Copy link

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment