Skip to content
This repository was archived by the owner on Sep 17, 2025. It is now read-only.

Commit 39299b8

Browse files
liyanhui1228liyanhui1228
authored
Add cloud client libraries integration (#33)
1 parent 75b7a36 commit 39299b8

File tree

8 files changed

+300
-53
lines changed

8 files changed

+300
-53
lines changed

‎opencensus/trace/ext/google_cloud_clientlibs/__init__.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
# Copyright 2017, OpenCensus Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from opencensus.trace.ext.google_cloud_clientlibs import trace
16+
17+
__all__ = ['trace']

‎opencensus/trace/ext/google_cloud_clientlibs/trace.py‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
# Copyright 2017, OpenCensus Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import logging
16+
17+
import grpc
18+
19+
from google.cloud import _helpers
20+
21+
from opencensus.trace import execution_context
22+
from opencensus.trace.ext.grpc.client_interceptor import (
23+
OpenCensusClientInterceptor)
24+
25+
from opencensus.trace.ext.requests.trace import (
26+
trace_integration as trace_requests)
27+
28+
log = logging.getLogger(__name__)
29+
30+
MODULE_NAME = 'google_cloud_clientlibs'
31+
32+
MAKE_SECURE_CHANNEL = 'make_secure_channel'
33+
INSECURE_CHANNEL = 'insecure_channel'
34+
35+
36+
def trace_integration():
37+
"""Trace the Google Cloud Client libraries by integrating with
38+
the transport level including HTTP and gRPC.
39+
"""
40+
log.info('Integrated module: {}'.format(MODULE_NAME))
41+
42+
# Integrate with gRPC
43+
trace_grpc()
44+
45+
# Integrate with HTTP
46+
trace_http()
47+
48+
49+
def trace_grpc():
50+
"""Integrate with gRPC."""
51+
# Wrap google.cloud._helpers.make_secure_channel
52+
make_secure_channel_func = getattr(_helpers, MAKE_SECURE_CHANNEL)
53+
make_secure_channel_wrapped = wrap_make_secure_channel(
54+
make_secure_channel_func)
55+
setattr(
56+
_helpers,
57+
MAKE_SECURE_CHANNEL,
58+
make_secure_channel_wrapped)
59+
60+
# Wrap the grpc.insecure_channel.
61+
insecure_channel_func = getattr(grpc, INSECURE_CHANNEL)
62+
insecure_channel_wrapped = wrap_insecure_channel(
63+
insecure_channel_func)
64+
setattr(
65+
grpc,
66+
INSECURE_CHANNEL,
67+
insecure_channel_wrapped)
68+
69+
70+
def trace_http():
71+
"""Integrate with HTTP (requests library)."""
72+
trace_requests()
73+
74+
75+
def wrap_make_secure_channel(make_secure_channel_func):
76+
"""Wrap the google.cloud._helpers.make_secure_channel."""
77+
def call(*args, **kwargs):
78+
channel = make_secure_channel_func(*args, **kwargs)
79+
80+
try:
81+
host = kwargs.get('host')
82+
_tracer = execution_context.get_opencensus_tracer()
83+
tracer_interceptor = OpenCensusClientInterceptor(_tracer, host)
84+
intercepted_channel = grpc.intercept_channel(
85+
channel, tracer_interceptor)
86+
return intercepted_channel # pragma: NO COVER
87+
except Exception:
88+
log.warning(
89+
'Failed to wrap secure channel, '
90+
'clientlibs grpc calls not traced.')
91+
return channel
92+
return call
93+
94+
95+
def wrap_insecure_channel(insecure_channel_func):
96+
"""Wrap the grpc.insecure_channel."""
97+
def call(*args, **kwargs):
98+
channel = insecure_channel_func(*args, **kwargs)
99+
100+
try:
101+
target = kwargs.get('target')
102+
_tracer = execution_context.get_opencensus_tracer()
103+
tracer_interceptor = OpenCensusClientInterceptor(_tracer, target)
104+
intercepted_channel = grpc.intercept_channel(
105+
channel, tracer_interceptor)
106+
return intercepted_channel # pragma: NO COVER
107+
except Exception:
108+
log.warning(
109+
'Failed to wrap insecure channel, '
110+
'clientlibs grpc calls not traced.')
111+
return channel
112+
return call

‎opencensus/trace/ext/grpc/client_interceptor.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,15 @@ def _intercept_call(
9797
grpc_trace_metadata = {
9898
oc_grpc.GRPC_TRACE_KEY: header,
9999
}
100-
metadata = metadata + tuple(six.iteritems(grpc_trace_metadata))
100+
101+
metadata_to_append = None
102+
103+
if isinstance(metadata, list):
104+
metadata_to_append = list(six.iteritems(grpc_trace_metadata))
105+
else:
106+
metadata_to_append = tuple(six.iteritems(grpc_trace_metadata))
107+
108+
metadata = metadata + metadata_to_append
101109

102110
client_call_details = _ClientCallDetails(
103111
client_call_details.method,

‎opencensus/trace/ext/requests/trace.py‎

Lines changed: 17 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
import logging
1616
import requests
17+
import wrapt
1718

1819
from opencensus.trace import execution_context
1920

@@ -37,7 +38,8 @@ def trace_integration():
3738
setattr(requests, requests_func.__name__, wrapped)
3839

3940
# Wrap Session class
40-
setattr(requests, SESSION_CLASS_NAME, TraceSession)
41+
wrapt.wrap_function_wrapper(
42+
MODULE_NAME, 'Session.request', wrap_session_request)
4143

4244

4345
def wrap_requests(requests_func):
@@ -62,33 +64,22 @@ def call(url, *args, **kwargs):
6264
return call
6365

6466

65-
def wrap_session_request(request_func):
67+
def wrap_session_request(wrapped, instance, args, kwargs):
6668
"""Wrap the session function to trace it."""
67-
def call(method, url, *args, **kwargs):
68-
_tracer = execution_context.get_opencensus_tracer()
69-
_span = _tracer.start_span()
70-
_span.name = '[requests]{}'.format(method)
71-
72-
# Add the requests url to attributes
73-
_tracer.add_attribute_to_current_span('requests/url', url)
74-
75-
result = request_func(method, url, *args, **kwargs)
76-
77-
# Add the status code to attributes
78-
_tracer.add_attribute_to_current_span(
79-
'requests/status_code', str(result.status_code))
80-
81-
_tracer.end_span()
82-
return result
83-
84-
return call
69+
method = kwargs.get('method') or args[0]
70+
url = kwargs.get('url') or args[1]
71+
_tracer = execution_context.get_opencensus_tracer()
72+
_span = _tracer.start_span()
73+
_span.name = '[requests]{}'.format(method)
8574

75+
# Add the requests url to attributes
76+
_tracer.add_attribute_to_current_span('requests/url', url)
8677

87-
class TraceSession(requests.Session):
78+
result = wrapped(*args, **kwargs)
8879

89-
def __init__(self, *args, **kwargs):
90-
request_func = getattr(self, SESSION_WRAP_METHODS)
91-
wrapped = wrap_session_request(request_func)
92-
setattr(self, request_func.__name__, wrapped)
80+
# Add the status code to attributes
81+
_tracer.add_attribute_to_current_span(
82+
'requests/status_code', str(result.status_code))
9383

94-
super(TraceSession, self).__init__(*args, **kwargs)
84+
_tracer.end_span()
85+
return result

‎requirements-test.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,3 +12,4 @@ retrying==1.3.3
1212
SQLAlchemy==1.1.14
1313
webapp2==2.5.2
1414
WebOb==1.7.3
15+
wrapt==1.10.11

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

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
# Copyright 2017, OpenCensus Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import unittest
16+
17+
import mock
18+
19+
from opencensus.trace.ext.google_cloud_clientlibs import trace
20+
21+
22+
class Test_google_cloud_clientlibs_trace(unittest.TestCase):
23+
24+
def test_trace_integration(self):
25+
mock_trace_grpc = mock.Mock()
26+
mock_trace_http = mock.Mock()
27+
28+
patch_trace_grpc = mock.patch(
29+
'opencensus.trace.ext.google_cloud_clientlibs.trace.trace_grpc',
30+
mock_trace_grpc)
31+
patch_trace_http = mock.patch(
32+
'opencensus.trace.ext.google_cloud_clientlibs.trace.trace_http',
33+
mock_trace_http)
34+
35+
with patch_trace_grpc, patch_trace_http:
36+
trace.trace_integration()
37+
38+
self.assertTrue(mock_trace_grpc.called)
39+
self.assertTrue(mock_trace_http.called)
40+
41+
def test_trace_grpc(self):
42+
mock_wrap = mock.Mock()
43+
mock__helpers = mock.Mock()
44+
45+
wrap_result = 'wrap result'
46+
mock_wrap.return_value = wrap_result
47+
48+
mock_make_secure_channel_func = mock.Mock()
49+
mock_make_secure_channel_func.__name__ = 'make_secure_channel'
50+
setattr(
51+
mock__helpers,
52+
'make_secure_channel',
53+
mock_make_secure_channel_func)
54+
55+
patch_wrap = mock.patch(
56+
'opencensus.trace.ext.google_cloud_clientlibs.trace.wrap_make_secure_channel', mock_wrap)
57+
patch__helpers = mock.patch(
58+
'opencensus.trace.ext.google_cloud_clientlibs.trace._helpers', mock__helpers)
59+
60+
with patch_wrap, patch__helpers:
61+
trace.trace_integration()
62+
63+
self.assertEqual(getattr(mock__helpers, 'make_secure_channel'), wrap_result)
64+
65+
def test_trace_http(self):
66+
mock_trace_requests = mock.Mock()
67+
patch = mock.patch(
68+
'opencensus.trace.ext.google_cloud_clientlibs.trace.trace_requests',
69+
mock_trace_requests)
70+
71+
with patch:
72+
trace.trace_http()
73+
74+
self.assertTrue(mock_trace_requests.called)
75+
76+
def test_wrap_make_secure_channel(self):
77+
mock_tracer = mock.Mock()
78+
mock_interceptor = mock.Mock()
79+
mock_func = mock.Mock()
80+
81+
patch_tracer = mock.patch(
82+
'opencensus.trace.ext.google_cloud_clientlibs.trace.execution_context.'
83+
'get_opencensus_tracer',
84+
return_value=mock_tracer)
85+
patch_interceptor = mock.patch(
86+
'opencensus.trace.ext.google_cloud_clientlibs.trace.OpenCensusClientInterceptor',
87+
mock_interceptor)
88+
89+
wrapped = trace.wrap_make_secure_channel(mock_func)
90+
91+
with patch_tracer, patch_interceptor:
92+
wrapped()
93+
94+
self.assertTrue(mock_interceptor.called)
95+
96+
def test_wrap_insecure_channel(self):
97+
mock_tracer = mock.Mock()
98+
mock_interceptor = mock.Mock()
99+
mock_func = mock.Mock()
100+
101+
patch_tracer = mock.patch(
102+
'opencensus.trace.ext.google_cloud_clientlibs.trace.execution_context.'
103+
'get_opencensus_tracer',
104+
return_value=mock_tracer)
105+
patch_interceptor = mock.patch(
106+
'opencensus.trace.ext.google_cloud_clientlibs.trace.OpenCensusClientInterceptor',
107+
mock_interceptor)
108+
109+
wrapped = trace.wrap_insecure_channel(mock_func)
110+
111+
with patch_tracer, patch_interceptor:
112+
wrapped()
113+
114+
self.assertTrue(mock_interceptor.called)

‎tests/unit/trace/ext/grpc/test_client_interceptor.py‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,29 @@ def test__intercept_call_metadata_none(self):
9090

9191
self.assertEqual(expected_metadata, client_call_details.metadata)
9292

93-
def test__intercept_call(self):
93+
def test__intercept_call_metadata_list(self):
94+
tracer = mock.Mock()
95+
tracer.span_context = mock.Mock()
96+
test_header = 'test header'
97+
mock_propagator = mock.Mock()
98+
mock_propagator.to_header.return_value = test_header
99+
100+
interceptor = client_interceptor.OpenCensusClientInterceptor(
101+
tracer=tracer, host_port='test')
102+
interceptor._propagator = mock_propagator
103+
mock_client_call_details = mock.Mock()
104+
mock_client_call_details.metadata = [('test_key', 'test_value'),]
105+
106+
client_call_details, request_iterator, current_span = interceptor._intercept_call(
107+
mock_client_call_details,
108+
mock.Mock(),
109+
'unary_unary')
110+
111+
expected_metadata = [('test_key', 'test_value'), ('grpc-trace-bin', test_header),]
112+
113+
self.assertEqual(expected_metadata, client_call_details.metadata)
114+
115+
def test__intercept_call_metadata_tuple(self):
94116
tracer = mock.Mock()
95117
tracer.span_context = mock.Mock()
96118
test_header = 'test header'

0 commit comments

Comments
 (0)