Add Stackdriver Trace V2 models (#87)

Author: liyanhui1228

opencensus/trace/attributes.py

@@ -0,0 +1,75 @@
+# Copyright 2017, OpenCensus Authors
+#
+# 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.
+
+from opencensus.trace import utils
+
+
+def _format_attribute_value(value):
+    if isinstance(value, bool):
+        value_type = 'bool_value'
+    elif isinstance(value, int):
+        value_type = 'int_value'
+    elif isinstance(value, str):
+        value_type = 'string_value'
+        value = utils._get_truncatable_str(value)
+    else:
+        return None
+
+    return {value_type: value}
+
+
+class Attributes(object):
+    """A set of attributes, each in the format [KEY]:[VALUE].
+
+    :type attributes: dict
+    :param attributes: The set of attributes. Each attribute's key can be up
+                       to 128 bytes long. The value can be a string up to 256
+                       bytes, an integer, or the Boolean values true and false.
+    """
+    def __init__(self, attributes=None):
+        self.attributes = dict(attributes or {})
+
+    def set_attribute(self, key, value):
+        """Set a key value pair."""
+        self.attributes[key] = value
+
+    def delete_attribute(self, key):
+        """Delete an attribute given a key if existed."""
+        self.attributes.pop(key, None)
+
+    def get_attribute(self, key):
+        """Get a attribute value."""
+        return self.attributes.get(key, None)
+
+    def format_attributes_json(self):
+        """Convert the Attributes object to json format."""
+        attributes_json = {
+            utils.check_str_length(key)[0]: _format_attribute_value(value)
+            for key, value in self.attributes.items()
+        }
+
+        attributes_json = {}
+
+        for key, value in self.attributes.items():
+            key = utils.check_str_length(key)[0]
+            value = _format_attribute_value(value)
+
+            if value is not None:
+                attributes_json[key] = value
+
+        result = {
+            'attributeMap': attributes_json
+        }
+
+        return result

opencensus/trace/enums.py

@@ -19,6 +19,7 @@
 
 from google.cloud.trace.client import Client
 
+
 # Environment variable set in App Engine when vm:true is set.
 _APPENGINE_FLEXIBLE_ENV_VM = 'GAE_APPENGINE_HOSTNAME'
 
@@ -90,54 +91,67 @@ class StackdriverExporter(base.Exporter):
     """
     def __init__(self, client=None, project_id=None,
                  transport=sync.SyncTransport):
-        # The client will handler the case when project_id is None
+        # The client will handle the case when project_id is None
         if client is None:
             client = Client(project=project_id)
 
         self.client = client
         self.project_id = client.project
         self.transport = transport(self)
 
-    def emit(self, trace):
+    def emit(self, spans):
         """
-        :type trace: dict
-        :param trace: Trace collected.
+        :type spans: dict
+        :param spans: Spans collected.
         """
-        stackdriver_traces = self.translate_to_stackdriver(trace)
-        self.client.patch_traces(stackdriver_traces)
+        name = 'projects/{}'.format(self.project_id)
+        stackdriver_spans = self.translate_to_stackdriver(spans)
+        self.client.batch_write_spans(name, stackdriver_spans)
 
     def export(self, trace):
         self.transport.export(trace)
 
-    def translate_to_stackdriver(self, trace):
-        """
-        :type trace: dict
-        :param trace: Trace collected.
+    def translate_to_stackdriver(self, spans):
+        """Translate the spans json to Stackdriver format.
+
+        See: https://cloud.google.com/trace/docs/reference/v2/rest/v2/
+             projects.traces/batchWrite
+
+        :type spans: dict
+        :param spans: Spans collected.
 
         :rtype: dict
-        :returns: Traces in Google Cloud StackDriver Trace format.
+        :returns: Spans in Google Cloud StackDriver Trace format.
         """
-        set_attributes(trace)
-        spans = trace.get('spans')
-        trace_id = trace.get('traceId')
-        spans_json = []
+        set_attributes(spans)
+        spans_json = spans.get('spans')
+        trace_id = spans.get('traceId')
+        spans_list = []
+
+        for span in spans_json:
+            span_name = 'projects/{}/traces/{}/spans/{}'.format(
+                self.project_id, trace_id, span.get('spanId'))
 
-        for span in spans:
             span_json = {
-                'name': span.get('name'),
+                'name': span_name,
+                'displayName': span.get('displayName'),
                 'startTime': span.get('startTime'),
                 'endTime': span.get('endTime'),
-                'spanId': span.get('spanId'),
-                'parentSpanId': span.get('parentSpanId'),
-                'labels': span.get('attributes')
+                'spanId': str(span.get('spanId')),
+                'attributes': span.get('attributes'),
+                'links': span.get('links'),
+                'status': span.get('status'),
+                'stackTrace': span.get('stackTrace'),
+                'timeEvents': span.get('timeEvents'),
+                'sameProcessAsParentSpan': span.get('sameProcessAsParentSpan'),
+                'childSpanCount': span.get('childSpanCount')
             }
-            spans_json.append(span_json)
 
-        trace_json = {
-            'projectId': self.project_id,
-            'traceId': trace_id,
-            'spans': spans_json
-        }
+            if span.get('parentSpanId') is not None:
+                parent_span_id = str(span.get('parentSpanId'))
+                span_json['parentSpanId'] = parent_span_id
+
+            spans_list.append(span_json)
 
-        traces = {'traces': [trace_json]}
-        return traces
+        spans = {'spans': spans_list}
+        return spans

opencensus/trace/exporters/stackdriver_exporter.py

@@ -62,7 +62,7 @@ def _before_cursor_execute(conn, cursor, statement, parameters,
 
     _tracer = execution_context.get_opencensus_tracer()
     _span = _tracer.start_span()
-    _span.name = '[{}.query]{}'.format(MODULE_NAME, statement)
+    _span.name = '{}.query'.format(MODULE_NAME)
 
     # Set query statement attribute
     _tracer.add_attribute_to_current_span(

opencensus/trace/ext/sqlalchemy/trace.py

@@ -0,0 +1,71 @@
+# Copyright 2017, OpenCensus Authors
+#
+# 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.
+
+
+class Type(object):
+    """The relationship of the current span relative to the linked span: child,
+    parent, or unspecified.
+
+    Attributes:
+      TYPE_UNSPECIFIED (int): The relationship of the two spans is unknown.
+      CHILD_LINKED_SPAN (int): The linked span is a child of the current span.
+      PARENT_LINKED_SPAN (int): The linked span is a parent of the current
+      span.
+    """
+    TYPE_UNSPECIFIED = 0
+    CHILD_LINKED_SPAN = 1
+    PARENT_LINKED_SPAN = 2
+
+
+class Link(object):
+    """A pointer from the current span to another span in the same trace or in
+    a different trace. For example, this can be used in batching operations,
+    where a single batch handler processes multiple requests from different
+    traces or when the handler receives a request from a different project.
+
+    :type trace_id: str
+    :param trace_id: The [TRACE_ID] for a trace within a project.
+
+    :type span_id: str
+    :param span_id: The [SPAN_ID] for a span within a trace.
+
+    :type type: Enum of :class:`~opencensus.trace.link.Type`
+    :param type: The relationship of the current span relative to the linked
+                 span.
+
+    :type attributes: :class:`~opencensus.trace.attributes.Attributes`
+    :param attributes: A set of attributes on the link. You have have up to 32
+                       attributes per link.
+    """
+    def __init__(self, trace_id, span_id, type=None, attributes=None):
+        self.trace_id = trace_id
+        self.span_id = span_id
+
+        if type is None:
+            type = Type.TYPE_UNSPECIFIED
+
+        self.type = type
+        self.attributes = attributes
+
+    def format_link_json(self):
+        """Convert a Link object to json format."""
+        link_json = {}
+        link_json['trace_id'] = self.trace_id
+        link_json['span_id'] = self.span_id
+        link_json['type'] = self.type
+
+        if self.attributes is not None:
+            link_json['attributes'] = self.attributes
+
+        return link_json