Add set_status to span, update pymongo integration (#738)

victoraugustolls · c24t · commit 5831ba9774f2 · 2019-08-02T12:58:50.000-07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,8 @@
 # Changelog
 
 ## Unreleased
+- Added `set_status` to `span`
+  ([#738](https://github.com/census-instrumentation/opencensus-python/pull/738))
 
 ## 0.7.0
 Released 2019-07-31
diff --git a/contrib/opencensus-ext-flask/tests/test_flask_middleware.py b/contrib/opencensus-ext-flask/tests/test_flask_middleware.py
@@ -316,8 +316,10 @@ def test_teardown_include_exception(self):
         exported_spandata = mock_exporter.export.call_args[0][0][0]
         self.assertIsInstance(exported_spandata, span_data.SpanData)
         self.assertIsInstance(exported_spandata.status, status.Status)
-        self.assertEqual(exported_spandata.status.code, code_pb2.UNKNOWN)
-        self.assertEqual(exported_spandata.status.message, 'error')
+        self.assertEqual(
+            exported_spandata.status.canonical_code, code_pb2.UNKNOWN
+        )
+        self.assertEqual(exported_spandata.status.description, 'error')
 
     def test_teardown_include_exception_and_traceback(self):
         mock_exporter = mock.MagicMock()
@@ -331,8 +333,10 @@ def test_teardown_include_exception_and_traceback(self):
         exported_spandata = mock_exporter.export.call_args[0][0][0]
         self.assertIsInstance(exported_spandata, span_data.SpanData)
         self.assertIsInstance(exported_spandata.status, status.Status)
-        self.assertEqual(exported_spandata.status.code, code_pb2.UNKNOWN)
-        self.assertEqual(exported_spandata.status.message, 'error')
+        self.assertEqual(
+            exported_spandata.status.canonical_code, code_pb2.UNKNOWN
+        )
+        self.assertEqual(exported_spandata.status.description, 'error')
         self.assertIsInstance(
             exported_spandata.stack_trace, stack_trace.StackTrace
         )
diff --git a/contrib/opencensus-ext-grpc/tests/test_server_interceptor.py b/contrib/opencensus-ext-grpc/tests/test_server_interceptor.py
@@ -147,8 +147,10 @@ def test_intercept_handler_exception(self):
 
             # check that the status obj is attached to the current span
             self.assertIsNotNone(current_span.status)
-            self.assertEqual(current_span.status.code, code_pb2.UNKNOWN)
-            self.assertEqual(current_span.status.message, 'Test')
+            self.assertEqual(
+                current_span.status.canonical_code, code_pb2.UNKNOWN
+            )
+            self.assertEqual(current_span.status.description, 'Test')
 
     @mock.patch(
         'opencensus.trace.execution_context.get_opencensus_tracer')
diff --git a/contrib/opencensus-ext-jaeger/opencensus/ext/jaeger/trace_exporter/__init__.py b/contrib/opencensus-ext-jaeger/opencensus/ext/jaeger/trace_exporter/__init__.py
@@ -184,12 +184,12 @@ def translate_to_jaeger(self, span_datas):
                 tags.append(jaeger.Tag(
                     key='status.code',
                     vType=jaeger.TagType.LONG,
-                    vLong=status.code))
+                    vLong=status.canonical_code))
 
                 tags.append(jaeger.Tag(
                     key='status.message',
                     vType=jaeger.TagType.STRING,
-                    vStr=status.message))
+                    vStr=status.description))
 
             refs = _extract_refs_from_span(span)
             logs = _extract_logs_from_span(span)
diff --git a/contrib/opencensus-ext-ocagent/opencensus/ext/ocagent/trace_exporter/utils.py b/contrib/opencensus-ext-ocagent/opencensus/ext/ocagent/trace_exporter/utils.py
@@ -43,8 +43,9 @@ def translate_to_trace_proto(span_data):
             span_data.start_time),
         end_time=ocagent_utils.proto_ts_from_datetime_str(span_data.end_time),
         status=trace_pb2.Status(
-            code=span_data.status.code,
-            message=span_data.status.message)
+            code=span_data.status.canonical_code,
+            message=span_data.status.description,
+        )
         if span_data.status is not None else None,
         same_process_as_parent_span=BoolValue(
             value=span_data.same_process_as_parent_span)
diff --git a/contrib/opencensus-ext-pymongo/CHANGELOG.md b/contrib/opencensus-ext-pymongo/CHANGELOG.md
@@ -1,6 +1,8 @@
 # Changelog
 
 ## Unreleased
+- Changed attributes names to make it compatible with [OpenTelemetry](https://github.com/open-telemetry/opentelemetry-specification/blob/master/specification/data-semantic-conventions.md), maintaining OpenCensus specs fidelity
+([#738](https://github.com/census-instrumentation/opencensus-python/pull/738))
 
 ## 0.1.3
 Released 2019-05-31
diff --git a/contrib/opencensus-ext-pymongo/opencensus/ext/pymongo/trace.py b/contrib/opencensus-ext-pymongo/opencensus/ext/pymongo/trace.py
@@ -16,8 +16,11 @@
 
 from pymongo import monitoring
 
+from google.rpc import code_pb2
+
 from opencensus.trace import execution_context
 from opencensus.trace import span as span_module
+from opencensus.trace import status as status_module
 
 
 log = logging.getLogger(__name__)
@@ -34,7 +37,6 @@ def trace_integration(tracer=None):
 
 
 class MongoCommandListener(monitoring.CommandListener):
-
     def __init__(self, tracer=None):
         self._tracer = tracer
 
@@ -44,30 +46,47 @@ def tracer(self):
 
     def started(self, event):
         span = self.tracer.start_span(
-            name='{}.{}.{}.{}'.format(MODULE_NAME,
-                                      event.database_name,
-                                      event.command.get(event.command_name),
-                                      event.command_name))
+            name='{}.{}.{}.{}'.format(
+                MODULE_NAME,
+                event.database_name,
+                event.command.get(event.command_name),
+                event.command_name,
+            )
+        )
         span.span_kind = span_module.SpanKind.CLIENT
 
+        self.tracer.add_attribute_to_current_span('component', 'mongodb')
+        self.tracer.add_attribute_to_current_span('db.type', 'mongodb')
+        self.tracer.add_attribute_to_current_span(
+            'db.instance', event.database_name
+        )
+        self.tracer.add_attribute_to_current_span(
+            'db.statement', event.command.get(event.command_name)
+        )
+
         for attr in COMMAND_ATTRIBUTES:
             _attr = event.command.get(attr)
             if _attr is not None:
                 self.tracer.add_attribute_to_current_span(attr, str(_attr))
 
         self.tracer.add_attribute_to_current_span(
-            'request_id', event.request_id)
+            'request_id', event.request_id
+        )
 
         self.tracer.add_attribute_to_current_span(
-            'connection_id', str(event.connection_id))
+            'connection_id', str(event.connection_id)
+        )
 
     def succeeded(self, event):
-        self._stop('succeeded')
+        self._stop(code_pb2.OK)
 
     def failed(self, event):
-        self._stop('failed')
-
-    def _stop(self, status):
-        self.tracer.add_attribute_to_current_span('status', status)
-
+        self._stop(code_pb2.UNKNOWN, 'MongoDB error', event.failure)
+
+    def _stop(self, code, message='', details=None):
+        span = self.tracer.current_span()
+        status = status_module.Status(
+            code=code, message=message, details=details
+        )
+        span.set_status(status)
         self.tracer.end_span()
diff --git a/contrib/opencensus-ext-pymongo/tests/test_pymongo_trace.py b/contrib/opencensus-ext-pymongo/tests/test_pymongo_trace.py
@@ -49,6 +49,10 @@ def test_started(self):
         }
 
         expected_attrs = {
+            'component': 'mongodb',
+            'db.type': 'mongodb',
+            'db.instance': 'database_name',
+            'db.statement': 'find',
             'filter': 'filter',
             'sort': 'sort',
             'limit': 'limit',
@@ -63,8 +67,8 @@ def test_started(self):
             trace.MongoCommandListener().started(
                 event=MockEvent(command_attrs))
 
-        self.assertEqual(mock_tracer.current_span.attributes, expected_attrs)
-        self.assertEqual(mock_tracer.current_span.name, expected_name)
+        self.assertEqual(mock_tracer.span.attributes, expected_attrs)
+        self.assertEqual(mock_tracer.span.name, expected_name)
 
     def test_succeed(self):
         mock_tracer = MockTracer()
@@ -74,12 +78,16 @@ def test_succeed(self):
             'opencensus.trace.execution_context.get_opencensus_tracer',
             return_value=mock_tracer)
 
-        expected_attrs = {'status': 'succeeded'}
+        expected_status = {
+            'code': 0,
+            'message': '',
+            'details': None
+        }
 
         with patch:
             trace.MongoCommandListener().succeeded(event=MockEvent(None))
 
-        self.assertEqual(mock_tracer.current_span.attributes, expected_attrs)
+        self.assertEqual(mock_tracer.span.status, expected_status)
         mock_tracer.end_span.assert_called_with()
 
     def test_failed(self):
@@ -90,12 +98,16 @@ def test_failed(self):
             'opencensus.trace.execution_context.get_opencensus_tracer',
             return_value=mock_tracer)
 
-        expected_attrs = {'status': 'failed'}
+        expected_status = {
+            'code': 2,
+            'message': 'MongoDB error',
+            'details': 'failure'
+        }
 
         with patch:
             trace.MongoCommandListener().failed(event=MockEvent(None))
 
-        self.assertEqual(mock_tracer.current_span.attributes, expected_attrs)
+        self.assertEqual(mock_tracer.span.status, expected_status)
         mock_tracer.end_span.assert_called_with()
 
 
@@ -115,17 +127,31 @@ def __getattr__(self, item):
         return item
 
 
+class MockSpan(object):
+    def __init__(self):
+        self.status = None
+
+    def set_status(self, status):
+        self.status = {
+            'code': status.canonical_code,
+            'message': status.description,
+            'details': status.details,
+        }
+
+
 class MockTracer(object):
     def __init__(self):
-        self.current_span = None
+        self.span = MockSpan()
         self.end_span = mock.Mock()
 
     def start_span(self, name=None):
-        span = mock.Mock()
-        span.name = name
-        span.attributes = {}
-        self.current_span = span
-        return span
+        self.span.name = name
+        self.span.attributes = {}
+        self.span.status = {}
+        return self.span
 
     def add_attribute_to_current_span(self, key, value):
-        self.current_span.attributes[key] = value
+        self.span.attributes[key] = value
+
+    def current_span(self):
+        return self.span
diff --git a/opencensus/trace/base_span.py b/opencensus/trace/base_span.py
@@ -80,6 +80,14 @@ def add_link(self, link):
         """
         raise NotImplementedError
 
+    def set_status(self, status):
+        """Sets span status.
+
+        :type code: :class: `~opencensus.trace.status.Status`
+        :param code: A Status object.
+        """
+        raise NotImplementedError
+
     def start(self):
         """Set the start time for a span."""
         raise NotImplementedError
diff --git a/opencensus/trace/blank_span.py b/opencensus/trace/blank_span.py
@@ -136,6 +136,14 @@ def add_link(self, link):
         """
         pass
 
+    def set_status(self, status):
+        """No-op implementation of this method.
+
+        :type code: :class: `~opencensus.trace.status.Status`
+        :param code: A Status object.
+        """
+        pass
+
     def start(self):
         """No-op implementation of this method."""
         pass
diff --git a/opencensus/trace/span.py b/opencensus/trace/span.py
@@ -264,9 +264,13 @@ def __init__(
         else:
             self.links = BoundedList.from_seq(MAX_NUM_LINKS, links)
 
+        if status is None:
+            self.status = status_module.Status.as_ok()
+        else:
+            self.status = status
+
         self.span_id = span_id
         self.stack_trace = stack_trace
-        self.status = status
         self.same_process_as_parent_span = same_process_as_parent_span
         self._child_spans = []
         self.context_tracer = context_tracer
@@ -346,6 +350,18 @@ def add_link(self, link):
             raise TypeError("Type Error: received {}, but requires Link.".
                             format(type(link).__name__))
 
+    def set_status(self, status):
+        """Sets span status.
+
+        :type code: :class: `~opencensus.trace.status.Status`
+        :param code: A Status object.
+        """
+        if isinstance(status, status_module.Status):
+            self.status = status
+        else:
+            raise TypeError("Type Error: received {}, but requires Status.".
+                            format(type(status).__name__))
+
     def start(self):
         """Set the start time for a span."""
         self.start_time = utils.to_iso_str()
diff --git a/opencensus/trace/status.py b/opencensus/trace/status.py
@@ -39,17 +39,31 @@ class Status(object):
                     See: https://cloud.google.com/trace/docs/reference/v2/
                          rest/v2/Status#FIELDS.details
     """
-    def __init__(self, code, message, details=None):
+    def __init__(self, code, message=None, details=None):
         self.code = code
         self.message = message
         self.details = details
 
+    @property
+    def canonical_code(self):
+        return self.code
+
+    @property
+    def description(self):
+        return self.message
+
+    @property
+    def is_ok(self):
+        return self.canonical_code == code_pb2.OK
+
     def format_status_json(self):
         """Convert a Status object to json format."""
         status_json = {}
 
-        status_json['code'] = self.code
-        status_json['message'] = self.message
+        status_json['code'] = self.canonical_code
+
+        if self.description is not None:
+            status_json['message'] = self.description
 
         if self.details is not None:
             status_json['details'] = self.details
@@ -62,3 +76,9 @@ def from_exception(cls, exc):
             code=code_pb2.UNKNOWN,
             message=str(exc)
         )
+
+    @classmethod
+    def as_ok(cls):
+        return cls(
+            code=code_pb2.OK,
+        )
diff --git a/tests/unit/trace/test_base_span.py b/tests/unit/trace/test_base_span.py
@@ -73,6 +73,12 @@ def test_add_link_abstract(self):
         with self.assertRaises(NotImplementedError):
             span.add_link(None)
 
+    def test_set_status_abstract(self):
+        span = BaseSpan()
+
+        with self.assertRaises(NotImplementedError):
+            span.set_status(None)
+
     def test_iter_abstract(self):
         span = BaseSpan()
 
diff --git a/tests/unit/trace/test_blank_span.py b/tests/unit/trace/test_blank_span.py
diff --git a/tests/unit/trace/test_span.py b/tests/unit/trace/test_span.py
diff --git a/tests/unit/trace/test_status.py b/tests/unit/trace/test_status.py
