Repository: census-instrumentation/opencensus-python
Pull request #33 (2017-10-09 17:39:25)

Add cloud client libraries integration
For #101. Integrated with the cloud client libraries by tracing the gRPC and HTTP requests made by the client library calls.

googlebot: We found a Contributor License Agreement for you (the sender of this pull request), but were unable to find agreements for the commit author(s).  If you authored these, maybe you used a different email address in the git commits than was used to sign the CLA ([login here](https://cla.developers.google.com/) to double check)?  If these were authored by someone else, then they will need to sign a CLA as well, and confirm that they're okay with these being contributed to Google.
In order to pass this check, please resolve this problem and have the pull request author add another comment and the bot will run again. If the bot doesn't comment, it means it doesn't think anything has changed.

<!-- need_author_cla -->
liyanhui1228: Ready for review.

Patch:
diff --git a/opencensus/trace/ext/google_cloud_clientlibs/__init__.py b/opencensus/trace/ext/google_cloud_clientlibs/__init__.py
new file mode 100644
index 00000000..9e84a0c
--- /dev/null
+++ b/opencensus/trace/ext/google_cloud_clientlibs/__init__.py
@@ -0,0 +1,17 @@
+# 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.
+
+from opencensus.trace.ext.google_cloud_clientlibs import trace
+
+__all__ = ['trace']
diff --git a/opencensus/trace/ext/google_cloud_clientlibs/trace.py b/opencensus/trace/ext/google_cloud_clientlibs/trace.py
new file mode 100644
index 00000000..b76d872
--- /dev/null
+++ b/opencensus/trace/ext/google_cloud_clientlibs/trace.py
@@ -0,0 +1,112 @@
+# 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 grpc
+
+from google.cloud import _helpers
+
+from opencensus.trace import execution_context
+from opencensus.trace.ext.grpc.client_interceptor import (
+    OpenCensusClientInterceptor)
+
+from opencensus.trace.ext.requests.trace import (
+    trace_integration as trace_requests)
+
+log = logging.getLogger(__name__)
+
+MODULE_NAME = 'google_cloud_clientlibs'
+
+MAKE_SECURE_CHANNEL = 'make_secure_channel'
+INSECURE_CHANNEL = 'insecure_channel'
+
+
+def trace_integration():
+    """Trace the Google Cloud Client libraries by integrating with
+    the transport level including HTTP and gRPC.
+    """
+    log.info('Integrated module: {}'.format(MODULE_NAME))
+
+    # Integrate with gRPC
+    trace_grpc()
+
+    # Integrate with HTTP
+    trace_http()
+
+
+def trace_grpc():
+    """Integrate with gRPC."""
+    # Wrap google.cloud._helpers.make_secure_channel
+    make_secure_channel_func = getattr(_helpers, MAKE_SECURE_CHANNEL)
+    make_secure_channel_wrapped = wrap_make_secure_channel(
+        make_secure_channel_func)
+    setattr(
+        _helpers,
+        MAKE_SECURE_CHANNEL,
+        make_secure_channel_wrapped)
+
+    # Wrap the grpc.insecure_channel.
+    insecure_channel_func = getattr(grpc, INSECURE_CHANNEL)
+    insecure_channel_wrapped = wrap_insecure_channel(
+        insecure_channel_func)
+    setattr(
+        grpc,
+        INSECURE_CHANNEL,
+        insecure_channel_wrapped)
+
+
+def trace_http():
+    """Integrate with HTTP (requests library)."""
+    trace_requests()
+
+
+def wrap_make_secure_channel(make_secure_channel_func):
+    """Wrap the google.cloud._helpers.make_secure_channel."""
+    def call(*args, **kwargs):
+        channel = make_secure_channel_func(*args, **kwargs)
+
+        try:
+            host = kwargs.get('host')
+            _tracer = execution_context.get_opencensus_tracer()
+            tracer_interceptor = OpenCensusClientInterceptor(_tracer, host)
+            intercepted_channel = grpc.intercept_channel(
+                channel, tracer_interceptor)
+            return intercepted_channel  # pragma: NO COVER
+        except Exception:
+            log.warning(
+                'Failed to wrap secure channel, '
+                'clientlibs grpc calls not traced.')
+            return channel
+    return call
+
+
+def wrap_insecure_channel(insecure_channel_func):
+    """Wrap the grpc.insecure_channel."""
+    def call(*args, **kwargs):
+        channel = insecure_channel_func(*args, **kwargs)
+
+        try:
+            target = kwargs.get('target')
+            _tracer = execution_context.get_opencensus_tracer()
+            tracer_interceptor = OpenCensusClientInterceptor(_tracer, target)
+            intercepted_channel = grpc.intercept_channel(
+                channel, tracer_interceptor)
+            return intercepted_channel  # pragma: NO COVER
+        except Exception:
+            log.warning(
+                'Failed to wrap insecure channel, '
+                'clientlibs grpc calls not traced.')
+            return channel
+    return call
diff --git a/opencensus/trace/ext/grpc/client_interceptor.py b/opencensus/trace/ext/grpc/client_interceptor.py
index ed2636f..09af2b7 100644
--- a/opencensus/trace/ext/grpc/client_interceptor.py
+++ b/opencensus/trace/ext/grpc/client_interceptor.py
@@ -97,7 +97,15 @@ class OpenCensusClientInterceptor(grpc.UnaryUnaryClientInterceptor,
         grpc_trace_metadata = {
             oc_grpc.GRPC_TRACE_KEY: header,
         }
-        metadata = metadata + tuple(six.iteritems(grpc_trace_metadata))
+
+        metadata_to_append = None
+
+        if isinstance(metadata, list):
+            metadata_to_append = list(six.iteritems(grpc_trace_metadata))
+        else:
+            metadata_to_append = tuple(six.iteritems(grpc_trace_metadata))
+
+        metadata = metadata + metadata_to_append
 
         client_call_details = _ClientCallDetails(
             client_call_details.method,
diff --git a/opencensus/trace/ext/requests/trace.py b/opencensus/trace/ext/requests/trace.py
index 95a99a7..b60cc26 100644
--- a/opencensus/trace/ext/requests/trace.py
+++ b/opencensus/trace/ext/requests/trace.py
@@ -14,6 +14,7 @@
 
 import logging
 import requests
+import wrapt
 
 from opencensus.trace import execution_context
 
@@ -37,7 +38,8 @@ def trace_integration():
         setattr(requests, requests_func.__name__, wrapped)
 
     # Wrap Session class
-    setattr(requests, SESSION_CLASS_NAME, TraceSession)
+    wrapt.wrap_function_wrapper(
+        MODULE_NAME, 'Session.request', wrap_session_request)
 
 
 def wrap_requests(requests_func):
@@ -62,33 +64,22 @@ def wrap_requests(requests_func):
     return call
 
 
-def wrap_session_request(request_func):
+def wrap_session_request(wrapped, instance, args, kwargs):
     """Wrap the session function to trace it."""
-    def call(method, url, *args, **kwargs):
-        _tracer = execution_context.get_opencensus_tracer()
-        _span = _tracer.start_span()
-        _span.name = '[requests]{}'.format(method)
-
-        # Add the requests url to attributes
-        _tracer.add_attribute_to_current_span('requests/url', url)
-
-        result = request_func(method, url, *args, **kwargs)
-
-        # Add the status code to attributes
-        _tracer.add_attribute_to_current_span(
-            'requests/status_code', str(result.status_code))
-
-        _tracer.end_span()
-        return result
-
-    return call
+    method = kwargs.get('method') or args[0]
+    url = kwargs.get('url') or args[1]
+    _tracer = execution_context.get_opencensus_tracer()
+    _span = _tracer.start_span()
+    _span.name = '[requests]{}'.format(method)
 
+    # Add the requests url to attributes
+    _tracer.add_attribute_to_current_span('requests/url', url)
 
-class TraceSession(requests.Session):
+    result = wrapped(*args, **kwargs)
 
-    def __init__(self, *args, **kwargs):
-        request_func = getattr(self, SESSION_WRAP_METHODS)
-        wrapped = wrap_session_request(request_func)
-        setattr(self, request_func.__name__, wrapped)
+    # Add the status code to attributes
+    _tracer.add_attribute_to_current_span(
+        'requests/status_code', str(result.status_code))
 
-        super(TraceSession, self).__init__(*args, **kwargs)
+    _tracer.end_span()
+    return result


Test patch:
diff --git a/requirements-test.txt b/requirements-test.txt
index 4f237df..6213ff2 100644
--- a/requirements-test.txt
+++ b/requirements-test.txt
@@ -12,3 +12,4 @@ retrying==1.3.3
 SQLAlchemy==1.1.14
 webapp2==2.5.2
 WebOb==1.7.3
+wrapt==1.10.11
diff --git a/tests/unit/trace/ext/google_cloud_clientlibs/test_google_cloud_clientlibs_trace.py b/tests/unit/trace/ext/google_cloud_clientlibs/test_google_cloud_clientlibs_trace.py
new file mode 100644
index 00000000..8e26140
--- /dev/null
+++ b/tests/unit/trace/ext/google_cloud_clientlibs/test_google_cloud_clientlibs_trace.py
@@ -0,0 +1,114 @@
+# 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.google_cloud_clientlibs import trace
+
+
+class Test_google_cloud_clientlibs_trace(unittest.TestCase):
+
+    def test_trace_integration(self):
+        mock_trace_grpc = mock.Mock()
+        mock_trace_http = mock.Mock()
+
+        patch_trace_grpc = mock.patch(
+            'opencensus.trace.ext.google_cloud_clientlibs.trace.trace_grpc',
+            mock_trace_grpc)
+        patch_trace_http = mock.patch(
+            'opencensus.trace.ext.google_cloud_clientlibs.trace.trace_http',
+            mock_trace_http)
+
+        with patch_trace_grpc, patch_trace_http:
+            trace.trace_integration()
+
+        self.assertTrue(mock_trace_grpc.called)
+        self.assertTrue(mock_trace_http.called)
+
+    def test_trace_grpc(self):
+        mock_wrap = mock.Mock()
+        mock__helpers = mock.Mock()
+
+        wrap_result = 'wrap result'
+        mock_wrap.return_value = wrap_result
+
+        mock_make_secure_channel_func = mock.Mock()
+        mock_make_secure_channel_func.__name__ = 'make_secure_channel'
+        setattr(
+            mock__helpers,
+            'make_secure_channel',
+            mock_make_secure_channel_func)
+
+        patch_wrap = mock.patch(
+            'opencensus.trace.ext.google_cloud_clientlibs.trace.wrap_make_secure_channel', mock_wrap)
+        patch__helpers = mock.patch(
+            'opencensus.trace.ext.google_cloud_clientlibs.trace._helpers', mock__helpers)
+
+        with patch_wrap, patch__helpers:
+            trace.trace_integration()
+
+        self.assertEqual(getattr(mock__helpers, 'make_secure_channel'), wrap_result)
+
+    def test_trace_http(self):
+        mock_trace_requests = mock.Mock()
+        patch = mock.patch(
+            'opencensus.trace.ext.google_cloud_clientlibs.trace.trace_requests',
+            mock_trace_requests)
+
+        with patch:
+            trace.trace_http()
+
+        self.assertTrue(mock_trace_requests.called)
+
+    def test_wrap_make_secure_channel(self):
+        mock_tracer = mock.Mock()
+        mock_interceptor = mock.Mock()
+        mock_func = mock.Mock()
+
+        patch_tracer = mock.patch(
+            'opencensus.trace.ext.google_cloud_clientlibs.trace.execution_context.'
+            'get_opencensus_tracer',
+            return_value=mock_tracer)
+        patch_interceptor = mock.patch(
+            'opencensus.trace.ext.google_cloud_clientlibs.trace.OpenCensusClientInterceptor',
+            mock_interceptor)
+
+        wrapped = trace.wrap_make_secure_channel(mock_func)
+
+        with patch_tracer, patch_interceptor:
+            wrapped()
+
+        self.assertTrue(mock_interceptor.called)
+
+    def test_wrap_insecure_channel(self):
+        mock_tracer = mock.Mock()
+        mock_interceptor = mock.Mock()
+        mock_func = mock.Mock()
+
+        patch_tracer = mock.patch(
+            'opencensus.trace.ext.google_cloud_clientlibs.trace.execution_context.'
+            'get_opencensus_tracer',
+            return_value=mock_tracer)
+        patch_interceptor = mock.patch(
+            'opencensus.trace.ext.google_cloud_clientlibs.trace.OpenCensusClientInterceptor',
+            mock_interceptor)
+
+        wrapped = trace.wrap_insecure_channel(mock_func)
+
+        with patch_tracer, patch_interceptor:
+            wrapped()
+
+        self.assertTrue(mock_interceptor.called)
diff --git a/tests/unit/trace/ext/grpc/test_client_interceptor.py b/tests/unit/trace/ext/grpc/test_client_interceptor.py
index cf1c7a0..b4281a4 100644
--- a/tests/unit/trace/ext/grpc/test_client_interceptor.py
+++ b/tests/unit/trace/ext/grpc/test_client_interceptor.py
@@ -90,7 +90,29 @@ class TestOpenCensusClientInterceptor(unittest.TestCase):
 
         self.assertEqual(expected_metadata, client_call_details.metadata)
 
-    def test__intercept_call(self):
+    def test__intercept_call_metadata_list(self):
+        tracer = mock.Mock()
+        tracer.span_context = mock.Mock()
+        test_header = 'test header'
+        mock_propagator = mock.Mock()
+        mock_propagator.to_header.return_value = test_header
+
+        interceptor = client_interceptor.OpenCensusClientInterceptor(
+            tracer=tracer, host_port='test')
+        interceptor._propagator = mock_propagator
+        mock_client_call_details = mock.Mock()
+        mock_client_call_details.metadata = [('test_key', 'test_value'),]
+
+        client_call_details, request_iterator, current_span = interceptor._intercept_call(
+            mock_client_call_details,
+            mock.Mock(),
+            'unary_unary')
+
+        expected_metadata = [('test_key', 'test_value'), ('grpc-trace-bin', test_header),]
+
+        self.assertEqual(expected_metadata, client_call_details.metadata)
+
+    def test__intercept_call_metadata_tuple(self):
         tracer = mock.Mock()
         tracer.span_context = mock.Mock()
         test_header = 'test header'
diff --git a/tests/unit/trace/ext/requests/test_requests_trace.py b/tests/unit/trace/ext/requests/test_requests_trace.py
index c1ea356..a1db9a0 100644
--- a/tests/unit/trace/ext/requests/test_requests_trace.py
+++ b/tests/unit/trace/ext/requests/test_requests_trace.py
@@ -74,11 +74,11 @@ class Test_requests_trace(unittest.TestCase):
         self.assertEqual(expected_name, mock_tracer.current_span.name)
 
     def test_wrap_session_request(self):
-        mock_return = mock.Mock()
-        mock_return.status_code = 200
-        return_value = mock_return
-        mock_func = mock.Mock()
-        mock_func.return_value = return_value
+        def wrapped(*args, **kwargs):
+            result = mock.Mock()
+            result.status_code = 200
+            return result
+
         mock_tracer = MockTracer()
 
         patch = mock.patch(
@@ -86,13 +86,12 @@ class Test_requests_trace(unittest.TestCase):
             'get_opencensus_tracer',
             return_value=mock_tracer)
 
-        wrapped = trace.wrap_session_request(mock_func)
-
         url = 'http://localhost:8080'
         request_method = 'POST'
 
         with patch:
-            wrapped(request_method, url)
+            result = trace.wrap_session_request(
+                wrapped, 'Session.request', (request_method, url), {})
 
         expected_attributes = {
             'requests/url': url,
@@ -103,23 +102,6 @@ class Test_requests_trace(unittest.TestCase):
         self.assertEqual(expected_name, mock_tracer.current_span.name)
 
 
-class TestTraceSession(unittest.TestCase):
-
-    def test___init__(self):
-        import requests
-
-        mock_wrapped = mock.Mock()
-        patch = mock.patch(
-            'opencensus.trace.ext.requests.trace.wrap_session_request',
-            return_value=mock_wrapped)
-
-        with patch:
-            session = trace.TraceSession()
-
-        self.assertEqual(session.request, mock_wrapped)
-        assert isinstance(session, requests.Session)
-
-
 class MockTracer(object):
     def __init__(self):
         self.current_span = None