GraphQL Obfuscated Queries (#279)

* Format

* Add query obfuscation testing

* Add query as agent attribute for graphql

* Fix inheritance of graphql traces

* Fix graphql2 tests

* Move graphql operation trace to capture sql

Author: TimPansino

newrelic/api/graphql_trace.py

@@ -14,8 +14,9 @@
 
 import functools
 
-from newrelic.common.async_wrapper import async_wrapper
 from newrelic.api.time_trace import TimeTrace, current_trace
+from newrelic.api.transaction import current_transaction
+from newrelic.common.async_wrapper import async_wrapper
 from newrelic.common.object_wrapper import FunctionWrapper, wrap_object
 from newrelic.core.graphql_node import GraphQLOperationNode, GraphQLResolverNode
 
@@ -31,17 +32,39 @@ def __init__(self, **kwargs):
         self.operation_name = None
         self.operation_type = None
         self.deepest_path = None
+        self.graphql = None
+        self.graphql_format = None
+        self.statement = None
 
     def __repr__(self):
         return '<%s %s>' % (self.__class__.__name__, dict())
 
-    def finalize_data(self, *args, **kwargs):
+    @property
+    def formatted(self):
+        if not self.statement:
+            return "<unknown>"
+
+        transaction = current_transaction()
+
+        # Record SQL settings
+        settings = transaction.settings
+        tt = settings.transaction_tracer
+        self.graphql_format = tt.record_sql
+
+        return self.statement.formatted(self.graphql_format)
+
+    def finalize_data(self, transaction, exc=None, value=None, tb=None):
+        # Add attributes
         self._add_agent_attribute("graphql.operation.type", self.operation_type or "<unknown>")
         self._add_agent_attribute("graphql.operation.name", self.operation_name or "<anonymous>")
         self._add_agent_attribute("graphql.operation.deepestPath", self.deepest_path or "<unknown>")
 
-        return super(GraphQLOperationTrace, self).finalize_data(*args, **kwargs)
+        # Attach formatted graphql
+        limit = transaction.settings.agent_limits.sql_query_length_maximum
+        self.graphql = graphql = self.formatted[:limit]
+        self._add_agent_attribute("graphql.operation.query", graphql)
 
+        return super(GraphQLOperationTrace, self).finalize_data(transaction, exc=None, value=None, tb=None)
 
     def create_node(self):
         return GraphQLOperationNode(
@@ -56,6 +79,7 @@ def create_node(self):
             operation_name=self.operation_name,
             operation_type=self.operation_type,
             deepest_path=self.deepest_path,
+            graphql=self.graphql,
         )
 
 
@@ -87,18 +111,19 @@ def wrap_graphql_operation_trace(module, object_path):
     wrap_object(module, object_path, GraphQLOperationTraceWrapper)
 
 
-class GraphQLResolverTrace(GraphQLOperationTrace):
+class GraphQLResolverTrace(TimeTrace):
     def __init__(self, field_name=None, **kwargs):
         super(GraphQLResolverTrace, self).__init__(**kwargs)
         self.field_name = field_name
 
+    def __repr__(self):
+        return '<%s %s>' % (self.__class__.__name__, dict())
 
     def finalize_data(self, *args, **kwargs):
         self._add_agent_attribute("graphql.field.name", self.field_name)
 
         return super(GraphQLResolverTrace, self).finalize_data(*args, **kwargs)
 
-
     def create_node(self):
         return GraphQLResolverNode(
             field_name=self.field_name,

newrelic/config.py

@@ -2377,6 +2377,11 @@ def _process_module_builtin_defaults():
         "gluon.contrib.memcache.memcache", "newrelic.hooks.memcache_memcache"
     )
 
+    _process_module_definition(
+        "graphql.graphql",
+        "newrelic.hooks.framework_graphql",
+        "instrument_graphql",
+    )
     _process_module_definition(
         "graphql.execution.execute",
         "newrelic.hooks.framework_graphql",

newrelic/core/graphql_node.py

@@ -17,14 +17,13 @@
 import newrelic.core.trace_node
 
 from newrelic.core.metric import TimeMetric
-
 from newrelic.core.node_mixin import GenericNodeMixin
 
 
 _GraphQLOperationNode = namedtuple('_GraphQLNode',
-    ['operation_type', 'operation_name', 'deepest_path', 
-    'children', 'start_time', 'end_time', 'duration', 'exclusive', 
-    'guid', 'agent_attributes', 'user_attributes'])
+    ['operation_type', 'operation_name', 'deepest_path', 'graphql', 
+    'children', 'start_time', 'end_time', 'duration', 'exclusive', 'guid',
+    'agent_attributes', 'user_attributes'])
 
 _GraphQLResolverNode = namedtuple('_GraphQLNode',
     ['field_name', 'children', 'start_time', 'end_time', 'duration', 

newrelic/core/graphql_utils.py

@@ -0,0 +1,51 @@
+# Copyright 2010 New Relic, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""GraphQL utilities that wrap SQL utilities for reuse."""
+
+import weakref
+
+from newrelic.core.database_utils import SQLStatement
+
+
+class GraphQLStyle(object):
+    """Helper class to initialize SQLStatement instances."""
+
+    quoting_style = "single+double"
+
+
+class GraphQLStatement(SQLStatement):
+    """Wrap SQLStatements to allow usage with GraphQL."""
+
+    def __init__(self, graphql):
+        super(GraphQLStatement, self).__init__(graphql, GraphQLStyle())
+        # Preset unapplicable fields to empty
+        self._operation = ""
+        self._target = ""
+
+
+_graphql_statements = weakref.WeakValueDictionary()
+
+
+def graphql_statement(graphql):
+    result = _graphql_statements.get(graphql, None)
+
+    if result is not None:
+        return result
+
+    result = GraphQLStatement(graphql)
+
+    _graphql_statements[graphql] = result
+
+    return result

newrelic/hooks/framework_graphql.py

@@ -12,17 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from newrelic.api.time_trace import notice_error, current_trace
 from newrelic.api.error_trace import ErrorTrace, ErrorTraceWrapper
 from newrelic.api.function_trace import FunctionTrace, FunctionTraceWrapper
 from newrelic.api.graphql_trace import GraphQLResolverTrace, GraphQLOperationTrace
+from newrelic.api.time_trace import notice_error, current_trace
 from newrelic.api.transaction import current_transaction
 from newrelic.common.object_names import callable_name, parse_exc_info
 from newrelic.common.object_wrapper import function_wrapper, wrap_function_wrapper
+from newrelic.core.graphql_utils import graphql_statement
 from collections import deque
 from copy import copy
 
 
+def graphql_version():
+    from graphql import __version__ as version
+    return tuple(int(v) for v in version.split("."))
+
+
 def ignore_graphql_duplicate_exception(exc, val, tb):
     from graphql.error import GraphQLError
 
@@ -67,17 +73,6 @@ def ignore_graphql_duplicate_exception(exc, val, tb):
     return None  # Follow original exception matching rules
 
 
-def wrap_execute(wrapped, instance, args, kwargs):
-    transaction = current_transaction()
-    if transaction is None:
-        return wrapped(*args, **kwargs)
-    
-    transaction.set_transaction_name(callable_name(wrapped), priority=1)
-    with GraphQLOperationTrace():
-        with ErrorTrace(ignore=ignore_graphql_duplicate_exception):
-            return wrapped(*args, **kwargs)
-
-
 def wrap_executor_context_init(wrapped, instance, args, kwargs):
     result = wrapped(*args, **kwargs)
 
@@ -250,10 +245,7 @@ def wrap_resolve_field(wrapped, instance, args, kwargs):
     if transaction is None:
         return wrapped(*args, **kwargs)
 
-    from graphql import __version__ as version
-    version = tuple(int(v) for v in version.split("."))
-
-    if version <= (3, 0, 0):
+    if graphql_version() <= (3, 0, 0):
         bind_resolve_field = bind_resolve_field_v2
     else:
         bind_resolve_field = bind_resolve_field_v3
@@ -272,9 +264,46 @@ def wrap_resolve_field(wrapped, instance, args, kwargs):
         return wrapped(*args, **kwargs)
 
 
+def bind_graphql_impl_query(schema, source, *args, **kwargs):
+    return source
+
+
+def bind_execute_graphql_query(
+    schema,
+    request_string="",
+    root=None,
+    context=None,
+    variables=None,
+    operation_name=None,
+    middleware=None,
+    backend=None,
+    **execute_options):
+
+    return request_string
+
+
+def wrap_graphql_impl(wrapped, instance, args, kwargs):
+    transaction = current_transaction()
+
+    if not transaction:
+        return wrapped(*args, **kwargs)
+
+    if graphql_version() <= (3, 0, 0):
+        bind_query = bind_execute_graphql_query
+    else:
+        bind_query = bind_graphql_impl_query
+
+    query = bind_query(*args, **kwargs)
+    if hasattr(query, "body"):
+        query = query.body
+
+    with GraphQLOperationTrace() as trace:
+        trace.statement = graphql_statement(query)
+        with ErrorTrace(ignore=ignore_graphql_duplicate_exception):
+            return wrapped(*args, **kwargs)
+
+
 def instrument_graphql_execute(module):
-    if hasattr(module, "execute"):
-        wrap_function_wrapper(module, "execute", wrap_execute)
     if hasattr(module, "get_field_def"):
         wrap_function_wrapper(module, "get_field_def", wrap_get_field_def)
     if hasattr(module, "ExecutionContext"):
@@ -323,3 +352,9 @@ def instrument_graphql_error_located_error(module):
 
 def instrument_graphql_validate(module):
     wrap_function_wrapper(module, "validate", wrap_validate)
+
+def instrument_graphql(module):
+    if hasattr(module, "graphql_impl"):
+        wrap_function_wrapper(module, "graphql_impl", wrap_graphql_impl)
+    if hasattr(module, "execute_graphql"):
+        wrap_function_wrapper(module, "execute_graphql", wrap_graphql_impl)

tests/framework_graphql/_target_application.py

@@ -68,6 +68,10 @@ def resolve_hello(root, info):
     return "Hello!"
 
 
+def resolve_echo(root, info, echo):
+    return echo
+
+
 def resolve_error(root, info):
     raise RuntimeError("Runtime Error!")
 
@@ -79,6 +83,11 @@ def resolve_error(root, info):
         resolver=resolve_library,
         args={"index": GraphQLArgument(GraphQLNonNull(GraphQLInt))},
     )
+    echo_field = GraphQLField(
+        GraphQLString,
+        resolver=resolve_echo,
+        args={"echo": GraphQLArgument(GraphQLNonNull(GraphQLString))},
+    )
     storage_field = GraphQLField(
         Storage,
         resolver=resolve_storage,
@@ -99,6 +108,11 @@ def resolve_error(root, info):
         resolve=resolve_library,
         args={"index": GraphQLArgument(GraphQLNonNull(GraphQLInt))},
     )
+    echo_field = GraphQLField(
+        GraphQLString,
+        resolve=resolve_echo,
+        args={"echo": GraphQLArgument(GraphQLNonNull(GraphQLString))},
+    )
     storage_field = GraphQLField(
         Storage,
         resolve=resolve_storage,
@@ -117,10 +131,11 @@ def resolve_error(root, info):
     name="Query",
     fields={
         "hello": hello_field,
+        "library": library_field,
+        "echo": echo_field,
         "storage": storage_field,
         "error": error_field,
         "error_non_null": error_non_null_field,
-        "library": library_field,
     },
 )