Remove vendor specific labels from core library (#62)

Author: liyanhui1228

trace/opencensus/trace/exporters/stackdriver_exporter.py

@@ -12,10 +12,63 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+
 from opencensus.trace.exporters import base
 
 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'
+
+# Environment variable set in App Engine when env:flex is set.
+_APPENGINE_FLEXIBLE_ENV_FLEX = 'GAE_INSTANCE'
+
+# GAE common labels
+# See: https://cloud.google.com/appengine/docs/flexible/python/runtime#
+#      environment_variables
+GAE_LABELS = {
+    'GAE_FLEX_VERSION': 'g.co/gae/app/version',
+    'GAE_FLEX_SERVICE': 'g.co/gae/app/service',
+    'GAE_FLEX_PROJECT': 'g.co/gae/app/project',
+    'GAE_FLEX_INSTANCE': 'g.co/gae/app/instance',
+    'GAE_FLEX_MEMORY_MB': 'g.co/gae/app/memory_mb',
+    'GAE_FLEX_PORT': 'g.co/gae/app/port',
+}
+
+# GCE common labels
+GCE_LABELS = {
+    'GCE_INSTANCE_ID': 'g.co/gce/instanceid',
+    'GCE_HOSTNAME': 'g.co/gce/hostname',
+}
+
+
+def set_labels(trace):
+    """Automatically set labels for Google Cloud environment."""
+    if is_gae_environment():
+        set_gae_labels(trace)
+
+
+def set_gae_labels(trace):
+    """Set the GAE environment common labels."""
+    spans = trace.get('spans')
+
+    for env_var, label_key in GAE_LABELS.items():
+        label_value = os.environ.get(env_var)
+
+        if label_value is not None:
+            for span in spans:
+                labels = span.get('labels')
+                labels[label_key] = label_value
+                span['labels'] = labels
+
+
+def is_gae_environment():
+    """Return True if the GAE related env vars is detected."""
+    if (_APPENGINE_FLEXIBLE_ENV_VM in os.environ or
+            _APPENGINE_FLEXIBLE_ENV_FLEX in os.environ):
+        return True
+
 
 class StackdriverExporter(base.Exporter):
     """A exporter that send traces and trace spans to Google Cloud Stackdriver
@@ -45,6 +98,7 @@ def translate_to_stackdriver(self, trace):
         :rtype: dict
         :returns: Traces in Google Cloud StackDriver Trace format.
         """
+        set_labels(trace)
         trace['projectId'] = self.project_id
         traces = {'traces': [trace]}
         return traces

trace/opencensus/trace/ext/django/middleware.py

@@ -22,9 +22,9 @@
 from opencensus.trace import execution_context
 from opencensus.trace.samplers import probability
 
-HTTP_METHOD = labels_helper.STACKDRIVER_LABELS['HTTP_METHOD']
-HTTP_URL = labels_helper.STACKDRIVER_LABELS['HTTP_URL']
-HTTP_STATUS_CODE = labels_helper.STACKDRIVER_LABELS['HTTP_STATUS_CODE']
+HTTP_METHOD = labels_helper.COMMON_LABELS['HTTP_METHOD']
+HTTP_URL = labels_helper.COMMON_LABELS['HTTP_URL']
+HTTP_STATUS_CODE = labels_helper.COMMON_LABELS['HTTP_STATUS_CODE']
 
 REQUEST_THREAD_LOCAL_KEY = 'django_request'
 

trace/opencensus/trace/ext/flask/flask_middleware.py

@@ -24,9 +24,9 @@
 
 _FLASK_TRACE_HEADER = 'X_CLOUD_TRACE_CONTEXT'
 
-HTTP_METHOD = labels_helper.STACKDRIVER_LABELS['HTTP_METHOD']
-HTTP_URL = labels_helper.STACKDRIVER_LABELS['HTTP_URL']
-HTTP_STATUS_CODE = labels_helper.STACKDRIVER_LABELS['HTTP_STATUS_CODE']
+HTTP_METHOD = labels_helper.COMMON_LABELS['HTTP_METHOD']
+HTTP_URL = labels_helper.COMMON_LABELS['HTTP_URL']
+HTTP_STATUS_CODE = labels_helper.COMMON_LABELS['HTTP_STATUS_CODE']
 
 log = logging.getLogger(__name__)
 

trace/opencensus/trace/labels_helper.py

@@ -12,18 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-
-# Environment variable set in App Engine when vm:true is set.
-_APPENGINE_FLEXIBLE_ENV_VM = 'GAE_APPENGINE_HOSTNAME'
-
-# Environment variable set in App Engine when env:flex is set.
-_APPENGINE_FLEXIBLE_ENV_FLEX = 'GAE_INSTANCE'
-
-# Stackdriver Trace common labels.
-# See: https://cloud.google.com/trace/docs/reference/v1/rpc/
-#      google.devtools.cloudtrace.v1#google.devtools.cloudtrace.v1.TraceSpan
-STACKDRIVER_LABELS = {
+COMMON_LABELS = {
     'AGENT': '/agent',
     'COMPONENT': '/component',
     'ERROR_MESSAGE': '/error/message',
@@ -44,47 +33,3 @@
     'STACKTRACE': '/stacktrace',
     'TID': '/tid',
 }
-
-# GAE common labels
-# See: https://cloud.google.com/appengine/docs/flexible/python/runtime#
-#      environment_variables
-GAE_LABELS = {
-    'GAE_FLEX_VERSION': 'g.co/gae/app/version',
-    'GAE_FLEX_SERVICE': 'g.co/gae/app/service',
-    'GAE_FLEX_PROJECT': 'g.co/gae/app/project',
-    'GAE_FLEX_INSTANCE': 'g.co/gae/app/instance',
-    'GAE_FLEX_MEMORY_MB': 'g.co/gae/app/memory_mb',
-    'GAE_FLEX_PORT': 'g.co/gae/app/port',
-}
-
-# GCE common labels
-GCE_LABELS = {
-    'GCE_INSTANCE_ID': 'g.co/gce/instanceid',
-    'GCE_HOSTNAME': 'g.co/gce/hostname',
-}
-
-
-class LabelsHelper(object):
-    """Helper class to automatically set labels."""
-
-    def __init__(self, tracer):
-        self.tracer = tracer
-
-    def set_labels(self):
-        """Automatically set labels for each environment."""
-        if self.is_gae_environment():
-            self.set_gae_labels()
-
-    def set_gae_labels(self):
-        """Set the GAE environment common labels."""
-        for env_var, label_key in GAE_LABELS.items():
-            label_value = os.environ.get(env_var)
-
-            if label_value is not None:
-                self.tracer.add_label_to_spans(label_key, label_value)
-
-    def is_gae_environment(self):
-        """Return True if the GAE related env vars is detected."""
-        if (_APPENGINE_FLEXIBLE_ENV_VM in os.environ or
-                _APPENGINE_FLEXIBLE_ENV_FLEX in os.environ):
-            return True

trace/opencensus/trace/tracer/context_tracer.py

@@ -16,7 +16,6 @@
 
 from opencensus.trace import execution_context
 from opencensus.trace.span_context import SpanContext
-from opencensus.trace import labels_helper
 from opencensus.trace import span as trace_span
 from opencensus.trace.tracer import base
 
@@ -45,10 +44,6 @@ def finish(self):
         :rtype: dict
         :returns: JSON format trace.
         """
-        # Insert the common labels to spans
-        helper = labels_helper.LabelsHelper(self)
-        helper.set_labels()
-
         trace = self._get_trace_json()
 
         return trace

trace/tests/unit/exporters/test_stackdriver_exporter.py

@@ -88,3 +88,35 @@ def test_translate_to_stackdriver(self):
         expected_traces = {'traces': [trace]}
 
         self.assertEqual(traces, expected_traces)
+
+
+class Test_set_labels_gae(unittest.TestCase):
+
+    def test_set_labels_gae(self):
+        import os
+
+        trace = {
+            'spans': [
+                {
+                    'labels':{},
+                    'span_id': 123,
+                },
+            ],
+        }
+
+        expected_labels = {
+            stackdriver_exporter.GAE_LABELS['GAE_FLEX_PROJECT']: 'project',
+            stackdriver_exporter.GAE_LABELS['GAE_FLEX_SERVICE']: 'service',
+            stackdriver_exporter.GAE_LABELS['GAE_FLEX_VERSION']: 'version',
+        }
+
+        with mock.patch.dict(
+                os.environ,
+                {stackdriver_exporter._APPENGINE_FLEXIBLE_ENV_VM: 'vm',
+                 stackdriver_exporter._APPENGINE_FLEXIBLE_ENV_FLEX: 'flex',
+                 'GAE_FLEX_PROJECT': 'project',
+                 'GAE_FLEX_SERVICE': 'service',
+                 'GAE_FLEX_VERSION': 'version'}):
+            stackdriver_exporter.set_labels(trace)
+
+        self.assertEqual(trace['spans'][0]['labels'], expected_labels)