Add integration for grpc channel from api_core (#131)

Author: liyanhui1228

opencensus/trace/config_integration.py

@@ -18,7 +18,7 @@
 log = logging.getLogger(__name__)
 
 SUPPORTED_INTEGRATIONS = ['httplib', 'mysql', 'postgresql', 'pymysql',
-                          'requests', 'sqlalchemy']
+                          'requests', 'sqlalchemy', 'google_cloud_clientlibs']
 
 PATH_PREFIX = 'opencensus.trace.ext'
 

opencensus/trace/ext/google_cloud_clientlibs/trace.py

@@ -17,6 +17,7 @@
 import grpc
 
 from google.cloud import _helpers
+from google.api_core import grpc_helpers
 
 from opencensus.trace import execution_context
 from opencensus.trace.ext.grpc.client_interceptor import (
@@ -31,6 +32,7 @@
 
 MAKE_SECURE_CHANNEL = 'make_secure_channel'
 INSECURE_CHANNEL = 'insecure_channel'
+CREATE_CHANNEL = 'create_channel'
 
 
 def trace_integration():
@@ -66,6 +68,14 @@ def trace_grpc():
         INSECURE_CHANNEL,
         insecure_channel_wrapped)
 
+    # Wrap google.api_core.grpc_helpers.create_channel
+    create_channel_func = getattr(grpc_helpers, CREATE_CHANNEL)
+    create_channel_wrapped = wrap_create_channel(create_channel_func)
+    setattr(
+        grpc_helpers,
+        CREATE_CHANNEL,
+        create_channel_wrapped)
+
 
 def trace_http():
     """Integrate with HTTP (requests library)."""
@@ -110,3 +120,23 @@ def call(*args, **kwargs):
                 'clientlibs grpc calls not traced.')
             return channel
     return call
+
+
+def wrap_create_channel(create_channel_func):
+    """Wrap the google.api_core.grpc_helpers.create_channel."""
+    def call(*args, **kwargs):
+        channel = create_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 create_channel, '
+                'clientlibs grpc calls not traced.')
+            return channel
+    return call

tests/unit/trace/ext/google_cloud_clientlibs/test_google_cloud_clientlibs_trace.py

@@ -112,3 +112,23 @@ def test_wrap_insecure_channel(self):
             wrapped()
 
         self.assertTrue(mock_interceptor.called)
+
+    def test_wrap_create_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_create_channel(mock_func)
+
+        with patch_tracer, patch_interceptor:
+            wrapped()
+
+        self.assertTrue(mock_interceptor.called)