Add httplib integration (#95)

liyanhui1228 · web-flow · commit 77382666f879 · 2017-12-11T11:10:51.000-08:00
diff --git a/opencensus/trace/config_integration.py b/opencensus/trace/config_integration.py
@@ -17,7 +17,8 @@
 
 log = logging.getLogger(__name__)
 
-SUPPORTED_INTEGRATIONS = ['mysql', 'postgresql', 'sqlalchemy', 'requests']
+SUPPORTED_INTEGRATIONS = ['httplib', 'mysql', 'postgresql', 'pymysql',
+                          'requests', 'sqlalchemy']
 
 PATH_PREFIX = 'opencensus.trace.ext'
 
diff --git a/opencensus/trace/ext/httplib/__init__.py b/opencensus/trace/ext/httplib/__init__.py
@@ -0,0 +1,13 @@
+# 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.
diff --git a/opencensus/trace/ext/httplib/trace.py b/opencensus/trace/ext/httplib/trace.py
@@ -0,0 +1,106 @@
+# 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.
+
+import logging
+import sys
+
+from opencensus.trace import attributes_helper
+from opencensus.trace import execution_context
+
+PYTHON2 = sys.version_info.major == 2
+
+if PYTHON2:
+    import httplib
+else:
+    import http.client as httplib
+
+log = logging.getLogger(__name__)
+
+MODULE_NAME = 'httplib'
+HTTPLIB_REQUEST_FUNC = 'request'
+HTTPLIB_RESPONSE_FUNC = 'getresponse'
+
+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']
+
+
+def trace_integration():
+    """Wrap the httplib to trace."""
+    log.info('Integrated module: {}'.format(MODULE_NAME))
+
+    # Wrap the httplib request function
+    request_func = getattr(
+        httplib.HTTPConnection, HTTPLIB_REQUEST_FUNC)
+    wrapped_request = wrap_httplib_request(request_func)
+    setattr(httplib.HTTPConnection, request_func.__name__, wrapped_request)
+
+    # Wrap the httplib response function
+    response_func = getattr(
+        httplib.HTTPConnection, HTTPLIB_RESPONSE_FUNC)
+    wrapped_response = wrap_httplib_response(response_func)
+    setattr(httplib.HTTPConnection, response_func.__name__, wrapped_response)
+
+
+def wrap_httplib_request(request_func):
+    """Wrap the httplib request function to trace. Create a new span and update
+    and close the span in the response later.
+    """
+    def call(self, method, url, *args, **kwargs):
+        _tracer = execution_context.get_opencensus_tracer()
+        _span = _tracer.start_span()
+        _span.name = '[httplib]{}'.format(request_func.__name__)
+
+        # Add the request url to attributes
+        _tracer.add_attribute_to_current_span(HTTP_URL, url)
+
+        # Add the request method to attributes
+        _tracer.add_attribute_to_current_span(HTTP_METHOD, method)
+
+        # Store the current span id to thread local.
+        execution_context.set_opencensus_attr(
+            'httplib/current_span_id', _span.span_id)
+
+        return request_func(self, method, url, *args, **kwargs)
+
+    return call
+
+
+def wrap_httplib_response(response_func):
+    """Wrap the httplib response function to trace.
+
+    If there is a corresponding httplib request span, update and close it.
+    If not, return the response.
+    """
+    def call(self, *args, **kwargs):
+        _tracer = execution_context.get_opencensus_tracer()
+        current_span_id = execution_context.get_opencensus_attr(
+            'httplib/current_span_id')
+
+        span = _tracer.current_span()
+
+        # No corresponding request span is found, request not traced.
+        if span.span_id != current_span_id:
+            return response_func(self, *args, **kwargs)
+
+        result = response_func(self, *args, **kwargs)
+
+        # Add the status code to attributes
+        _tracer.add_attribute_to_current_span(
+            HTTP_STATUS_CODE, str(result.status))
+
+        _tracer.end_span()
+        return result
+
+    return call
diff --git a/tests/unit/trace/ext/httplib/test_httplib_trace.py b/tests/unit/trace/ext/httplib/test_httplib_trace.py
@@ -0,0 +1,163 @@
+# 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.
+
+import unittest
+
+import mock
+
+from opencensus.trace.ext.httplib import trace
+
+
+class Test_httplib_trace(unittest.TestCase):
+
+    def test_trace_integration(self):
+        mock_wrap_request = mock.Mock()
+        mock_wrap_response = mock.Mock()
+        mock_httplib = mock.Mock()
+
+        wrap_request_result = 'wrap request result'
+        wrap_response_result = 'wrap response result'
+        mock_wrap_request.return_value = wrap_request_result
+        mock_wrap_response.return_value = wrap_response_result
+
+        mock_request_func = mock.Mock()
+        mock_response_func = mock.Mock()
+        mock_request_func.__name__ = 'request'
+        mock_response_func.__name__ = 'getresponse'
+        setattr(mock_httplib.HTTPConnection, 'request', mock_request_func)
+        setattr(mock_httplib.HTTPConnection, 'getresponse', mock_response_func)
+
+        patch_wrap_request = mock.patch(
+            'opencensus.trace.ext.httplib.trace.wrap_httplib_request',
+            mock_wrap_request)
+        patch_wrap_response = mock.patch(
+            'opencensus.trace.ext.httplib.trace.wrap_httplib_response',
+            mock_wrap_response)
+        patch_httplib = mock.patch(
+            'opencensus.trace.ext.httplib.trace.httplib', mock_httplib)
+
+        with patch_wrap_request, patch_wrap_response, patch_httplib:
+            trace.trace_integration()
+
+        self.assertEqual(getattr(mock_httplib.HTTPConnection, 'request'),
+                         wrap_request_result)
+        self.assertEqual(getattr(mock_httplib.HTTPConnection, 'getresponse'),
+                         wrap_response_result)
+
+    def test_wrap_httplib_request(self):
+        mock_tracer = MockTracer()
+        mock_request_func = mock.Mock()
+        mock_request_func.__name__ = 'request'
+
+        patch = mock.patch(
+            'opencensus.trace.ext.requests.trace.execution_context.'
+            'get_opencensus_tracer',
+            return_value=mock_tracer)
+
+        wrapped = trace.wrap_httplib_request(mock_request_func)
+
+        url = 'http://localhost:8080'
+        method = 'GET'
+
+        with patch:
+            wrapped(mock.Mock(), method, url)
+
+        expected_attributes = {
+            '/http/url': url,
+            '/http/method': method}
+        expected_name = '[httplib]request'
+
+        self.assertEqual(expected_attributes,
+                         mock_tracer.span.attributes)
+        self.assertEqual(expected_name, mock_tracer.span.name)
+
+    def test_wrap_httplib_response(self):
+        mock_span = mock.Mock()
+        span_id = '1234'
+        mock_span.span_id = span_id
+        mock_span.attributes = {}
+        mock_tracer = MockTracer(mock_span)
+        mock_response_func = mock.Mock()
+        mock_result = mock.Mock()
+        mock_result.status = '200'
+        mock_response_func.return_value = mock_result
+
+        patch_tracer = mock.patch(
+            'opencensus.trace.ext.requests.trace.execution_context.'
+            'get_opencensus_tracer',
+            return_value=mock_tracer)
+        patch_attr = mock.patch('opencensus.trace.ext.httplib.trace.'
+                                'execution_context.get_opencensus_attr',
+                                return_value=span_id)
+
+        wrapped = trace.wrap_httplib_response(mock_response_func)
+
+        with patch_tracer, patch_attr:
+            wrapped(mock.Mock())
+
+        expected_attributes = {
+            '/http/status_code': '200'}
+
+        self.assertEqual(expected_attributes,
+                         mock_tracer.span.attributes)
+
+    def test_wrap_httplib_response_no_open_span(self):
+        mock_span = mock.Mock()
+        span_id = '1234'
+        mock_span.span_id = span_id
+        mock_span.attributes = {}
+        mock_tracer = MockTracer(mock_span)
+        mock_response_func = mock.Mock()
+        mock_result = mock.Mock()
+        mock_result.status = '200'
+        mock_response_func.return_value = mock_result
+
+        patch_tracer = mock.patch(
+            'opencensus.trace.ext.requests.trace.execution_context.'
+            'get_opencensus_tracer',
+            return_value=mock_tracer)
+        patch_attr = mock.patch('opencensus.trace.ext.httplib.trace.'
+                                'execution_context.get_opencensus_attr',
+                                return_value='1111')
+
+        wrapped = trace.wrap_httplib_response(mock_response_func)
+
+        with patch_tracer, patch_attr:
+            wrapped(mock.Mock())
+
+        # Attribute should be empty as there is no matching span
+        expected_attributes = {}
+
+        self.assertEqual(expected_attributes,
+                         mock_tracer.span.attributes)
+
+
+class MockTracer(object):
+    def __init__(self, span=None):
+        self.span = span
+
+    def current_span(self):
+        return self.span
+
+    def start_span(self):
+        span = mock.Mock()
+        span.attributes = {}
+        self.span = span
+        return span
+
+    def end_span(self):
+        pass
+
+    def add_attribute_to_current_span(self, key, value):
+        self.span.attributes[key] = value
