Add agent label to stackdriver exporter (#111)

census-instrumentation · Feb 1, 2018 · ae72795 · ae72795
1 parent 54d11c9
commit ae72795
Show file tree

Hide file tree

Showing 3 changed files with 115 additions and 31 deletions.
diff --git a/opencensus/trace/attributes_helper.py b/opencensus/trace/attributes_helper.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 COMMON_ATTRIBUTES = {
-    'AGENT': '/agent',
+    'AGENT': 'g.co/agent',
     'COMPONENT': '/component',
     'ERROR_MESSAGE': '/error/message',
     'ERROR_NAME': '/error/name',

diff --git a/opencensus/trace/exporters/stackdriver_exporter.py b/opencensus/trace/exporters/stackdriver_exporter.py
@@ -14,12 +14,20 @@
 
 import os
 
+from opencensus.trace.attributes import Attributes
+from opencensus.trace import attributes_helper
 from opencensus.trace.exporters import base
 from opencensus.trace.exporters.transports import sync
 
 from google.cloud.trace.client import Client
 
 
+# OpenCensus Version
+VERSION = '0.1.1'
+
+# Agent
+AGENT = 'opencensus-python [{}]'.format(VERSION)
+
 # Environment variable set in App Engine when vm:true is set.
 _APPENGINE_FLEXIBLE_ENV_VM = 'GAE_APPENGINE_HOSTNAME'
 
@@ -45,24 +53,49 @@
 }
 
 
+def _update_attr_map(span, attrs):
+    attr_map = span.get('attributes', {}).get('attributeMap', {})
+    attr_map.update(attrs)
+    span['attributes']['attributeMap'] = attr_map
+
+
 def set_attributes(trace):
     """Automatically set attributes for Google Cloud environment."""
-    if is_gae_environment():
-        set_gae_attributes(trace)
+    spans = trace.get('spans')
+    for span in spans:
+        if span.get('attributes') is None:
+            span['attributes'] = {}
 
+        if is_gae_environment():
+            set_gae_attributes(span)
 
-def set_gae_attributes(trace):
-    """Set the GAE environment common attributes."""
-    spans = trace.get('spans')
+        set_common_attributes(span)
 
+
+def set_common_attributes(span):
+    """Set the common attributes."""
+    common = {
+        attributes_helper.COMMON_ATTRIBUTES.get('AGENT'): AGENT,
+    }
+    common_attrs = Attributes(common)\
+        .format_attributes_json()\
+        .get('attributeMap')
+
+    _update_attr_map(span, common_attrs)
+
+
+def set_gae_attributes(span):
+    """Set the GAE environment common attributes."""
     for env_var, attribute_key in GAE_ATTRIBUTES.items():
         attribute_value = os.environ.get(env_var)
 
         if attribute_value is not None:
-            for span in spans:
-                attributes = span.get('attributes')
-                attributes[attribute_key] = attribute_value
-                span['attributes'] = attributes
+            pair = {attribute_key: attribute_value}
+            pair_attrs = Attributes(pair)\
+                .format_attributes_json()\
+                .get('attributeMap')
+
+            _update_attr_map(span, pair_attrs)
 
 
 def is_gae_environment():

diff --git a/tests/unit/trace/exporters/test_stackdriver_exporter.py b/tests/unit/trace/exporters/test_stackdriver_exporter.py
@@ -87,7 +87,16 @@ def test_emit(self):
                     'links': None,
                     'startTime': None,
                     'spanId': '1111',
-                    'attributes': None,
+                    'attributes': {
+                        'attributeMap': {
+                            'g.co/agent': {
+                                'string_value': {
+                                    'truncated_byte_count': 0,
+                                    'value': 'opencensus-python [0.1.1]'
+                                }
+                            }
+                        }
+                    },
                     'stackTrace': None,
                     'displayName':
                         {
@@ -122,8 +131,15 @@ def test_translate_to_stackdriver(self):
         trace_id = '6e0c63257de34c92bf9efcd03927272e'
         span_name = 'test span'
         span_id = 1234
-        attributes = {'attributeMap': {
-            'key': 'value'}
+        attributes = {
+            'attributeMap': {
+                'key': {
+                    'string_value': {
+                        'truncated_byte_count': 0,
+                        'value': 'value'
+                    }
+               }
+            }
         }
         parent_span_id = 1111
         start_time = 'test start time'
@@ -164,7 +180,21 @@ def test_translate_to_stackdriver(self):
                         'value': span_name,
                         'truncated_byte_count': 0
                     },
-                    'attributes': {'attributeMap': {'key': 'value'}},
+                    'attributes': {
+                        'attributeMap': {
+                            'g.co/agent': {
+                                'string_value': {
+                                    'truncated_byte_count': 0,
+                                    'value': 'opencensus-python [0.1.1]'}
+                            },
+                            'key': {
+                                'string_value': {
+                                    'truncated_byte_count': 0,
+                                    'value': 'value'
+                                }
+                            }
+                        }
+                    },
                     'spanId': str(span_id),
                     'startTime': start_time,
                     'endTime': end_time,
@@ -179,9 +209,6 @@ def test_translate_to_stackdriver(self):
             ]
         }
 
-        print(spans)
-        print(expected_traces)
-
         self.assertEqual(spans, expected_traces)
 
 
@@ -190,19 +217,41 @@ class Test_set_attributes_gae(unittest.TestCase):
     def test_set_attributes_gae(self):
         import os
 
-        trace = {
-            'spans': [
-                {
-                    'attributes':{},
-                    'span_id': 123,
-                },
-            ],
-        }
-
-        expected_attributes = {
-            stackdriver_exporter.GAE_ATTRIBUTES['GAE_FLEX_PROJECT']: 'project',
-            stackdriver_exporter.GAE_ATTRIBUTES['GAE_FLEX_SERVICE']: 'service',
-            stackdriver_exporter.GAE_ATTRIBUTES['GAE_FLEX_VERSION']: 'version',
+        trace = {'spans': [
+            {
+                'attributes': {}
+            }
+        ]}
+
+        expected = {
+            'attributes': {
+                'attributeMap': {
+                    'g.co/gae/app/service': {
+                        'string_value': {
+                            'truncated_byte_count': 0,
+                            'value': 'service'
+                        }
+                    },
+                    'g.co/gae/app/version': {
+                        'string_value': {
+                            'truncated_byte_count': 0,
+                            'value': 'version'
+                        }
+                    },
+                    'g.co/gae/app/project': {
+                        'string_value': {
+                            'truncated_byte_count': 0,
+                            'value': 'project'
+                        }
+                    },
+                    'g.co/agent': {
+                        'string_value': {
+                            'truncated_byte_count': 0,
+                            'value': 'opencensus-python [0.1.1]'
+                        }
+                    },
+                }
+            }
         }
 
         with mock.patch.dict(
@@ -212,9 +261,11 @@ def test_set_attributes_gae(self):
                  'GAE_FLEX_PROJECT': 'project',
                  'GAE_FLEX_SERVICE': 'service',
                  'GAE_FLEX_VERSION': 'version'}):
+            self.assertTrue(stackdriver_exporter.is_gae_environment())
             stackdriver_exporter.set_attributes(trace)
 
-        self.assertEqual(trace['spans'][0]['attributes'], expected_attributes)
+        span = trace.get('spans')[0]
+        self.assertEqual(span, expected)
 
 
 class MockTransport(object):