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

Commit 124cdda

Browse files
liyanhui1228liyanhui1228
authored
Add pymysql integration (#50)
1 parent 70fcb6c commit 124cdda

10 files changed

Lines changed: 292 additions & 140 deletions

File tree

‎trace/opencensus/trace/ext/dbapi/__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.dbapi import trace
16+
17+
__all__ = ['trace']

‎trace/opencensus/trace/ext/dbapi/trace.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
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+
from opencensus.trace import execution_context
18+
19+
CURSOR_WRAP_METHOD = 'cursor'
20+
QUERY_WRAP_METHODS = ['execute', 'executemany']
21+
22+
23+
def wrap_conn(conn_func):
24+
"""Wrap the mysql conn object with TraceConnection."""
25+
def call(*args, **kwargs):
26+
try:
27+
conn = conn_func(*args, **kwargs)
28+
cursor_func = getattr(conn, CURSOR_WRAP_METHOD)
29+
wrapped = wrap_cursor(cursor_func)
30+
setattr(conn, cursor_func.__name__, wrapped)
31+
return conn
32+
except Exception: # pragma: NO COVER
33+
logging.warning('Fail to wrap conn, mysql not traced.')
34+
return conn_func(*args, **kwargs)
35+
return call
36+
37+
38+
def wrap_cursor(cursor_func):
39+
def call(*args, **kwargs):
40+
try:
41+
cursor = cursor_func(*args, **kwargs)
42+
for func in QUERY_WRAP_METHODS:
43+
query_func = getattr(cursor, func)
44+
wrapped = trace_cursor_query(query_func)
45+
setattr(cursor, query_func.__name__, wrapped)
46+
return cursor
47+
except Exception: # pragma: NO COVER
48+
logging.warning('Fail to wrap cursor, mysql not traced.')
49+
return cursor_func(*args, **kwargs)
50+
return call
51+
52+
53+
def trace_cursor_query(query_func):
54+
def call(query, *args, **kwargs):
55+
_tracer = execution_context.get_opencensus_tracer()
56+
_span = _tracer.start_span()
57+
_span.name = '[mysql.query]{}'.format(query)
58+
_tracer.add_label_to_current_span('mysql/query', query)
59+
_tracer.add_label_to_current_span(
60+
'mysql/cursor/method/name',
61+
query_func.__name__)
62+
63+
result = query_func(query, *args, **kwargs)
64+
65+
_tracer.end_span()
66+
return result
67+
return call

‎trace/opencensus/trace/ext/mysql/trace.py‎

Lines changed: 4 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -16,68 +16,17 @@
1616
import logging
1717
import mysql.connector
1818

19-
from opencensus.trace import execution_context
20-
21-
log = logging.getLogger(__name__)
19+
from opencensus.trace.ext.dbapi import trace
2220

2321
MODULE_NAME = 'mysql'
2422

2523
CONN_WRAP_METHOD = 'connect'
26-
CURSOR_WRAP_METHOD = 'cursor'
27-
QUERY_WRAP_METHODS = ['execute', 'executemany']
2824

2925

3026
def trace_integration():
3127
"""Wrap the mysql connector to trace it."""
32-
log.info('Integrated module: {}'.format(MODULE_NAME))
28+
logging.info('Integrated module: {}'.format(MODULE_NAME))
3329
conn_func = getattr(mysql.connector, CONN_WRAP_METHOD)
3430
conn_module = inspect.getmodule(conn_func)
35-
wrapped = wrap_conn(conn_func)
36-
setattr(conn_module, conn_func.__name__, wrapped)
37-
38-
39-
def wrap_conn(conn_func):
40-
"""Wrap the mysql conn object with TraceConnection."""
41-
def call(*args, **kwargs):
42-
try:
43-
conn = conn_func(*args, **kwargs)
44-
cursor_func = getattr(conn, CURSOR_WRAP_METHOD)
45-
wrapped = wrap_cursor(cursor_func)
46-
setattr(conn, cursor_func.__name__, wrapped)
47-
return conn
48-
except Exception: # pragma: NO COVER
49-
log.warning('Fail to wrap conn, mysql not traced.')
50-
return conn_func(*args, **kwargs)
51-
return call
52-
53-
54-
def wrap_cursor(cursor_func):
55-
def call(*args, **kwargs):
56-
try:
57-
cursor = cursor_func(*args, **kwargs)
58-
for func in QUERY_WRAP_METHODS:
59-
query_func = getattr(cursor, func)
60-
wrapped = trace_cursor_query(query_func)
61-
setattr(cursor, query_func.__name__, wrapped)
62-
return cursor
63-
except Exception: # pragma: NO COVER
64-
log.warning('Fail to wrap cursor, mysql not traced.')
65-
return cursor_func(*args, **kwargs)
66-
return call
67-
68-
69-
def trace_cursor_query(query_func):
70-
def call(query, *args, **kwargs):
71-
_tracer = execution_context.get_opencensus_tracer()
72-
_span = _tracer.start_span()
73-
_span.name = '[mysql.query]{}'.format(query)
74-
_tracer.add_label_to_current_span('mysql/query', query)
75-
_tracer.add_label_to_current_span(
76-
'mysql/cursor/method/name',
77-
query_func.__name__)
78-
79-
result = query_func(query, *args, **kwargs)
80-
81-
_tracer.end_span()
82-
return result
83-
return call
31+
wrapped = trace.wrap_conn(conn_func)
32+
setattr(conn_module, CONN_WRAP_METHOD, wrapped)

‎trace/opencensus/trace/ext/pymysql/__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.pymysql import trace
16+
17+
__all__ = ['trace']

‎trace/opencensus/trace/ext/pymysql/trace.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
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 inspect
16+
import logging
17+
import pymysql
18+
19+
from opencensus.trace.ext.dbapi import trace
20+
21+
MODULE_NAME = 'pymysql'
22+
23+
CONN_WRAP_METHOD = 'connect'
24+
25+
26+
def trace_integration():
27+
"""Wrap the mysql connector to trace it."""
28+
logging.info('Integrated module: {}'.format(MODULE_NAME))
29+
conn_func = getattr(pymysql, CONN_WRAP_METHOD)
30+
conn_module = inspect.getmodule(conn_func)
31+
wrapped = trace.wrap_conn(conn_func)
32+
setattr(conn_module, CONN_WRAP_METHOD, wrapped)

‎trace/requirements-test.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ google-cloud-trace==0.15.5
44
mock==2.0.0
55
mysql-connector==2.1.6
66
psycopg2==2.7.3.1
7+
pymysql==0.7.11
78
pytest==3.2.2
89
pytest-cov==2.5.1
910
SQLAlchemy==1.1.14

‎trace/tests/unit/ext/dbapi/test_dbapi_trace.py‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
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.dbapi import trace
20+
21+
22+
class TestDBAPITrace(unittest.TestCase):
23+
24+
def test_wrap_conn(self):
25+
wrap_func_name = 'wrap func'
26+
cursor_func_name = 'cursor_func'
27+
28+
def mock_wrap(func):
29+
return wrap_func_name
30+
31+
mock_func = mock.Mock()
32+
mock_return = mock.Mock()
33+
mock_cursor_func = mock.Mock()
34+
mock_cursor_func.__name__ = cursor_func_name
35+
mock_return.cursor = mock_cursor_func
36+
mock_func.return_value = mock_return
37+
wrapped = trace.wrap_conn(mock_func)
38+
39+
patch_wrap = mock.patch(
40+
'opencensus.trace.ext.dbapi.trace.wrap_cursor',
41+
side_effect=mock_wrap)
42+
43+
with patch_wrap:
44+
wrapped()
45+
46+
self.assertEqual(wrapped.__class__.__name__, 'function')
47+
self.assertEqual(
48+
getattr(mock_return, cursor_func_name, None),
49+
wrap_func_name)
50+
51+
52+
def test_wrap_cursor(self):
53+
wrap_func_name = 'wrap func'
54+
55+
def mock_trace_cursor_query(func):
56+
return wrap_func_name + func.__name__
57+
58+
mock_func = mock.Mock()
59+
mock_return = mock.Mock()
60+
61+
for func in trace.QUERY_WRAP_METHODS:
62+
query_func = mock.Mock()
63+
query_func.__name__ = func
64+
setattr(mock_return, func, query_func)
65+
66+
mock_func.return_value = mock_return
67+
68+
wrapped = trace.wrap_cursor(mock_func)
69+
70+
patch_wrap = mock.patch(
71+
'opencensus.trace.ext.dbapi.trace.trace_cursor_query',
72+
side_effect=mock_trace_cursor_query)
73+
74+
with patch_wrap:
75+
wrapped()
76+
77+
for func in trace.QUERY_WRAP_METHODS:
78+
self.assertEqual(
79+
getattr(mock_return, func, None),
80+
wrap_func_name + func)
81+
82+
83+
def test_trace_cursor_query(self):
84+
return_value = 'trace test'
85+
query = 'SELECT 1'
86+
mock_func = mock.Mock()
87+
mock_func.__name__ = 'execute'
88+
mock_func.return_value = return_value
89+
mock_tracer = mock.Mock()
90+
mock_cursor = mock.Mock()
91+
92+
patch = mock.patch(
93+
'opencensus.trace.ext.dbapi.trace.execution_context.'
94+
'get_opencensus_tracer',
95+
return_value=mock_tracer)
96+
97+
wrapped = trace.trace_cursor_query(mock_func)
98+
99+
with patch:
100+
result = wrapped(mock_cursor, query)
101+
102+
self.assertEqual(result, return_value)
103+
self.assertTrue(mock_tracer.start_span.called)
104+
self.assertTrue(mock_tracer.add_label_to_current_span.called)
105+
self.assertTrue(mock_tracer.end_span.called)

0 commit comments

Comments
 (0)