Set SpanKind in extensions (#271)

Liudmila Molkova · songy23 · commit 3b42664967b4 · 2018-08-22T15:58:56.000-07:00
diff --git a/opencensus/trace/ext/dbapi/trace.py b/opencensus/trace/ext/dbapi/trace.py
@@ -15,6 +15,7 @@
 import logging
 
 from opencensus.trace import execution_context
+from opencensus.trace import span as span_module
 
 CURSOR_WRAP_METHOD = 'cursor'
 QUERY_WRAP_METHODS = ['execute', 'executemany']
@@ -55,6 +56,7 @@ def call(query, *args, **kwargs):
         _tracer = execution_context.get_opencensus_tracer()
         _span = _tracer.start_span()
         _span.name = 'mysql.query'
+        _span.span_kind = span_module.SpanKind.CLIENT
         _tracer.add_attribute_to_current_span('mysql/query', query)
         _tracer.add_attribute_to_current_span(
             'mysql/cursor/method/name',
diff --git a/opencensus/trace/ext/django/middleware.py b/opencensus/trace/ext/django/middleware.py
@@ -18,8 +18,9 @@
 from opencensus.trace.ext import utils
 from opencensus.trace.ext.django.config import (settings, convert_to_import)
 from opencensus.trace import attributes_helper
-from opencensus.trace import tracer as tracer_module
 from opencensus.trace import execution_context
+from opencensus.trace import span as span_module
+from opencensus.trace import tracer as tracer_module
 from opencensus.trace.samplers import probability
 
 try:
@@ -164,7 +165,8 @@ def process_request(self, request):
                 propagator=self.propagator)
 
             # Span name is being set at process_view
-            tracer.start_span()
+            span = tracer.start_span()
+            span.span_kind = span_module.SpanKind.SERVER
             tracer.add_attribute_to_current_span(
                 attribute_key=HTTP_METHOD,
                 attribute_value=request.method)
diff --git a/opencensus/trace/ext/flask/flask_middleware.py b/opencensus/trace/ext/flask/flask_middleware.py
@@ -21,6 +21,7 @@
 
 from opencensus.trace import attributes_helper
 from opencensus.trace import execution_context
+from opencensus.trace import span as span_module
 from opencensus.trace import stack_trace
 from opencensus.trace import status
 from opencensus.trace import tracer as tracer_module
@@ -175,7 +176,7 @@ def _before_request(self):
                 propagator=self.propagator)
 
             span = tracer.start_span()
-
+            span.span_kind = span_module.SpanKind.SERVER
             # Set the span name as the name of the current module name
             span.name = '[{}]{}'.format(
                 flask.request.method,
diff --git a/opencensus/trace/ext/grpc/client_interceptor.py b/opencensus/trace/ext/grpc/client_interceptor.py
@@ -20,6 +20,7 @@
 
 from opencensus.trace import attributes_helper
 from opencensus.trace import execution_context
+from opencensus.trace import span as span_module
 from opencensus.trace import time_event
 from opencensus.trace.ext import grpc as oc_grpc
 from opencensus.trace.ext.grpc import utils as grpc_utils
@@ -67,6 +68,7 @@ def _start_client_span(self, client_call_details):
             name=_get_span_name(client_call_details)
         )
 
+        span.span_kind = span_module.SpanKind.CLIENT
         # Add the component grpc to span attribute
         self.tracer.add_attribute_to_current_span(
             attribute_key=attributes_helper.COMMON_ATTRIBUTES.get(
diff --git a/opencensus/trace/ext/grpc/server_interceptor.py b/opencensus/trace/ext/grpc/server_interceptor.py
@@ -19,8 +19,10 @@
 
 from opencensus.trace import attributes_helper
 from opencensus.trace import execution_context
+from opencensus.trace import span as span_module
 from opencensus.trace import stack_trace as stack_trace
 from opencensus.trace import status
+
 from opencensus.trace import time_event
 from opencensus.trace import tracer as tracer_module
 from opencensus.trace.ext import grpc as oc_grpc
@@ -108,6 +110,8 @@ def _start_server_span(self, servicer_context):
         span = tracer.start_span(
             name=_get_span_name(servicer_context)
         )
+
+        span.span_kind = span_module.SpanKind.SERVER
         tracer.add_attribute_to_current_span(
             attribute_key=attributes_helper.COMMON_ATTRIBUTES.get(
                 ATTRIBUTE_COMPONENT),
diff --git a/opencensus/trace/ext/httplib/trace.py b/opencensus/trace/ext/httplib/trace.py
@@ -17,6 +17,7 @@
 
 from opencensus.trace import attributes_helper
 from opencensus.trace import execution_context
+from opencensus.trace import span as span_module
 
 PYTHON2 = sys.version_info.major == 2
 
@@ -60,6 +61,7 @@ def wrap_httplib_request(request_func):
     def call(self, method, url, body, headers, *args, **kwargs):
         _tracer = execution_context.get_opencensus_tracer()
         _span = _tracer.start_span()
+        _span.span_kind = span_module.SpanKind.CLIENT
         _span.name = '[httplib]{}'.format(request_func.__name__)
 
         # Add the request url to attributes
diff --git a/opencensus/trace/ext/postgresql/trace.py b/opencensus/trace/ext/postgresql/trace.py
@@ -16,6 +16,7 @@
 import logging
 
 from opencensus.trace import execution_context
+from opencensus.trace import span as span_module
 
 import psycopg2
 from psycopg2 import connect as pg_connect
@@ -55,6 +56,7 @@ def call(query, *args, **kwargs):
             # here
             _span = _tracer.start_span()
             _span.name = '{}.query'.format(MODULE_NAME)
+            _span.span_kind = span_module.SpanKind.CLIENT
             _tracer.add_attribute_to_current_span(
                 '{}/query'.format(MODULE_NAME), query)
             _tracer.add_attribute_to_current_span(
diff --git a/opencensus/trace/ext/pyramid/pyramid_middleware.py b/opencensus/trace/ext/pyramid/pyramid_middleware.py
@@ -19,9 +19,9 @@
 
 from opencensus.trace import attributes_helper
 from opencensus.trace import execution_context
+from opencensus.trace import span as span_module
 from opencensus.trace import tracer as tracer_module
 
-
 HTTP_METHOD = attributes_helper.COMMON_ATTRIBUTES['HTTP_METHOD']
 HTTP_URL = attributes_helper.COMMON_ATTRIBUTES['HTTP_URL']
 HTTP_STATUS_CODE = attributes_helper.COMMON_ATTRIBUTES['HTTP_STATUS_CODE']
@@ -91,6 +91,7 @@ def _before_request(self, request):
                 request.method,
                 request.path)
 
+            span.span_kind = span_module.SpanKind.SERVER
             tracer.add_attribute_to_current_span(
                 attribute_key=HTTP_METHOD,
                 attribute_value=request.method)
diff --git a/opencensus/trace/ext/requests/trace.py b/opencensus/trace/ext/requests/trace.py
@@ -17,6 +17,7 @@
 import wrapt
 
 from opencensus.trace import execution_context
+from opencensus.trace import span as span_module
 
 log = logging.getLogger(__name__)
 
@@ -50,6 +51,7 @@ def call(url, *args, **kwargs):
         _tracer = execution_context.get_opencensus_tracer()
         _span = _tracer.start_span()
         _span.name = '[requests]{}'.format(requests_func.__name__)
+        _span.span_kind = span_module.SpanKind.CLIENT
 
         # Add the requests url to attributes
         _tracer.add_attribute_to_current_span('requests/url', url)
@@ -72,7 +74,9 @@ def wrap_session_request(wrapped, instance, args, kwargs):
     url = kwargs.get('url') or args[1]
     _tracer = execution_context.get_opencensus_tracer()
     _span = _tracer.start_span()
+
     _span.name = '[requests]{}'.format(method)
+    _span.span_kind = span_module.SpanKind.CLIENT
 
     # Add the requests url to attributes
     _tracer.add_attribute_to_current_span('requests/url', url)
diff --git a/opencensus/trace/ext/sqlalchemy/trace.py b/opencensus/trace/ext/sqlalchemy/trace.py
@@ -19,6 +19,7 @@
 
 
 from opencensus.trace import execution_context
+from opencensus.trace import span as span_module
 
 log = logging.getLogger(__name__)
 
@@ -63,6 +64,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)
+    _span.span_kind = span_module.SpanKind.CLIENT
 
     # Set query statement attribute
     _tracer.add_attribute_to_current_span(
diff --git a/tests/unit/trace/ext/django/test_middleware.py b/tests/unit/trace/ext/django/test_middleware.py
@@ -19,6 +19,7 @@
 from django.test.utils import teardown_test_environment
 
 from opencensus.trace import execution_context
+from opencensus.trace import span as span_module
 from opencensus.trace.exporters import print_exporter
 from opencensus.trace.exporters import zipkin_exporter
 from opencensus.trace.exporters.transports import sync
@@ -192,6 +193,7 @@ def test_process_request(self):
             '/http/url': u'/',
             '/http/method': 'GET',
         }
+        self.assertEqual(span.span_kind, span_module.SpanKind.SERVER)
         self.assertEqual(span.attributes, expected_attributes)
         self.assertEqual(span.parent_span.span_id, span_id)
 
diff --git a/tests/unit/trace/ext/flask/test_flask_middleware.py b/tests/unit/trace/ext/flask/test_flask_middleware.py
@@ -23,6 +23,7 @@
 
 from opencensus.trace import execution_context
 from opencensus.trace import span_data
+from opencensus.trace import span as span_module
 from opencensus.trace import stack_trace
 from opencensus.trace import status
 from opencensus.trace.exporters import print_exporter, stackdriver_exporter, \
@@ -180,6 +181,7 @@ def test__before_request(self):
                 '/http/method': 'GET',
             }
 
+            self.assertEqual(span.span_kind, span_module.SpanKind.SERVER)
             self.assertEqual(span.attributes, expected_attributes)
             self.assertEqual(span.parent_span.span_id, span_id)
 
diff --git a/tests/unit/trace/ext/grpc/test_client_interceptor.py b/tests/unit/trace/ext/grpc/test_client_interceptor.py
@@ -16,6 +16,7 @@
 
 import mock
 from opencensus.trace import execution_context
+from opencensus.trace import span as span_module
 from opencensus.trace.ext.grpc import client_interceptor
 from opencensus.trace.tracers.noop_tracer import NoopTracer
 
diff --git a/tests/unit/trace/ext/grpc/test_server_interceptor.py b/tests/unit/trace/ext/grpc/test_server_interceptor.py
@@ -137,6 +137,8 @@ def test_intercept_handler_exception(self):
                 execution_context.get_opencensus_tracer().current_span().attributes,
                 expected_attributes)
 
+            self.assertEqual(current_span.span_kind, span_module.SpanKind.SERVER)
+
             # check that the stack trace is attached to the current span
             self.assertIsNotNone(current_span.stack_trace)
             self.assertIsNotNone(current_span.stack_trace.stack_trace_hash_id)
diff --git a/tests/unit/trace/ext/httplib/test_httplib_trace.py b/tests/unit/trace/ext/httplib/test_httplib_trace.py
@@ -16,6 +16,7 @@
 
 import mock
 
+from opencensus.trace import span as span_module
 from opencensus.trace.ext.httplib import trace
 from opencensus.trace.propagation import trace_context_http_header_format
 
@@ -98,6 +99,7 @@ def test_wrap_httplib_request(self):
         self.assertEqual(expected_attributes,
                          mock_tracer.span.attributes)
         self.assertEqual(expected_name, mock_tracer.span.name)
+        self.assertEqual(span_module.SpanKind.CLIENT, mock_tracer.span.span_kind)        
 
     def test_wrap_httplib_response(self):
         mock_span = mock.Mock()
diff --git a/tests/unit/trace/ext/pyramid/test_pyramid_middleware.py b/tests/unit/trace/ext/pyramid/test_pyramid_middleware.py
@@ -23,6 +23,7 @@
 from pyramid.testing import DummyRequest
 
 from opencensus.trace import execution_context
+from opencensus.trace import span as span_module
 from opencensus.trace.exporters import print_exporter
 from opencensus.trace.exporters import zipkin_exporter
 from opencensus.trace.exporters.transports import sync
@@ -151,6 +152,7 @@ def dummy_handler(request):
             '/http/method': 'GET',
         }
 
+        self.assertEqual(span.span_kind, span_module.SpanKind.SERVER)
         self.assertEqual(span.attributes, expected_attributes)
         self.assertEqual(span.parent_span.span_id, span_id)
 
diff --git a/tests/unit/trace/ext/requests/test_requests_trace.py b/tests/unit/trace/ext/requests/test_requests_trace.py
@@ -16,6 +16,7 @@
 
 import mock
 
+from opencensus.trace import span as span_module
 from opencensus.trace.ext.requests import trace
 
 
@@ -70,6 +71,7 @@ def test_wrap_requests(self):
             'requests/status_code': '200'}
         expected_name = '[requests]get'
 
+        self.assertEqual(span_module.SpanKind.CLIENT, mock_tracer.current_span.span_kind)
         self.assertEqual(expected_attributes, mock_tracer.current_span.attributes)
         self.assertEqual(expected_name, mock_tracer.current_span.name)
 
@@ -98,6 +100,7 @@ def wrapped(*args, **kwargs):
             'requests/status_code': '200'}
         expected_name = '[requests]POST'
 
+        self.assertEqual(span_module.SpanKind.CLIENT, mock_tracer.current_span.span_kind)
         self.assertEqual(expected_attributes, mock_tracer.current_span.attributes)
         self.assertEqual(expected_name, mock_tracer.current_span.name)
 
diff --git a/tests/unit/trace/ext/sqlalchemy/test_sqlalchemy_trace.py b/tests/unit/trace/ext/sqlalchemy/test_sqlalchemy_trace.py
@@ -16,6 +16,7 @@
 
 import mock
 
+from opencensus.trace import span as span_module
 from opencensus.trace.ext.sqlalchemy import trace
 
 
@@ -77,6 +78,7 @@ def test__before_cursor_execute(self):
 
         expected_name = 'sqlalchemy.query'
 
+        self.assertEqual(mock_tracer.current_span.span_kind, span_module.SpanKind.CLIENT)
         self.assertEqual(mock_tracer.current_span.attributes, expected_attributes)
         self.assertEqual(mock_tracer.current_span.name, expected_name)
 
