Skip to content

Commit affe911

Browse files
lzchenlzchen
authored
Use is_recording flag in aiopg, asyncpg, dbapi, psycopg2, pymemcache, pymongo, redis, sqlalchemy instrumentations (open-telemetry#1212)
1 parent 0e852ea commit affe911

File tree

21 files changed

+304
-103
lines changed

21 files changed

+304
-103
lines changed

‎instrumentation/opentelemetry-instrumentation-aiopg/src/opentelemetry/instrumentation/aiopg/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ class AiopgInstrumentor(BaseInstrumentor):
6363

6464
def _instrument(self, **kwargs):
6565
"""Integrate with PostgreSQL aiopg library.
66-
aiopg: https://github.com/aio-libs/aiopg
66+
aiopg: https://github.com/aio-libs/aiopg
6767
"""
6868

6969
tracer_provider = kwargs.get("tracer_provider")

‎instrumentation/opentelemetry-instrumentation-aiopg/src/opentelemetry/instrumentation/aiopg/aiopg_integration.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,7 @@ async def wrapped_connection(
3737
args: typing.Tuple[typing.Any, typing.Any],
3838
kwargs: typing.Dict[typing.Any, typing.Any],
3939
):
40-
"""Add object proxy to connection object.
41-
"""
40+
"""Add object proxy to connection object."""
4241
connection = await connect_method(*args, **kwargs)
4342
# pylint: disable=protected-access
4443
self.get_connection_attributes(connection._conn)
@@ -109,10 +108,14 @@ async def traced_execution(
109108
self._populate_span(span, *args)
110109
try:
111110
result = await query_method(*args, **kwargs)
112-
span.set_status(Status(StatusCanonicalCode.OK))
111+
if span.is_recording():
112+
span.set_status(Status(StatusCanonicalCode.OK))
113113
return result
114114
except Exception as ex: # pylint: disable=broad-except
115-
span.set_status(Status(StatusCanonicalCode.UNKNOWN, str(ex)))
115+
if span.is_recording():
116+
span.set_status(
117+
Status(StatusCanonicalCode.UNKNOWN, str(ex))
118+
)
116119
raise ex
117120

118121

‎instrumentation/opentelemetry-instrumentation-aiopg/src/opentelemetry/instrumentation/aiopg/wrappers.py

Lines changed: 26 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -56,16 +56,16 @@ def trace_integration(
5656
tracer_provider: typing.Optional[TracerProvider] = None,
5757
):
5858
"""Integrate with aiopg library.
59-
based on dbapi integration, where replaced sync wrap methods to async
60-
61-
Args:
62-
database_component: Database driver name or
63-
database name "postgreSQL".
64-
database_type: The Database type. For any SQL database, "sql".
65-
connection_attributes: Attribute names for database, port, host and
66-
user in Connection object.
67-
tracer_provider: The :class:`opentelemetry.trace.TracerProvider` to
68-
use. If ommited the current configured one is used.
59+
based on dbapi integration, where replaced sync wrap methods to async
60+
61+
Args:
62+
database_component: Database driver name or
63+
database name "postgreSQL".
64+
database_type: The Database type. For any SQL database, "sql".
65+
connection_attributes: Attribute names for database, port, host and
66+
user in Connection object.
67+
tracer_provider: The :class:`opentelemetry.trace.TracerProvider` to
68+
use. If ommited the current configured one is used.
6969
"""
7070

7171
wrap_connect(
@@ -87,18 +87,18 @@ def wrap_connect(
8787
tracer_provider: typing.Optional[TracerProvider] = None,
8888
):
8989
"""Integrate with aiopg library.
90-
https://github.com/aio-libs/aiopg
91-
92-
Args:
93-
name: Name of opentelemetry extension for aiopg.
94-
database_component: Database driver name
95-
or database name "postgreSQL".
96-
database_type: The Database type. For any SQL database, "sql".
97-
connection_attributes: Attribute names for database, port, host and
98-
user in Connection object.
99-
version: Version of opentelemetry extension for aiopg.
100-
tracer_provider: The :class:`opentelemetry.trace.TracerProvider` to
101-
use. If ommited the current configured one is used.
90+
https://github.com/aio-libs/aiopg
91+
92+
Args:
93+
name: Name of opentelemetry extension for aiopg.
94+
database_component: Database driver name
95+
or database name "postgreSQL".
96+
database_type: The Database type. For any SQL database, "sql".
97+
connection_attributes: Attribute names for database, port, host and
98+
user in Connection object.
99+
version: Version of opentelemetry extension for aiopg.
100+
tracer_provider: The :class:`opentelemetry.trace.TracerProvider` to
101+
use. If ommited the current configured one is used.
102102
"""
103103

104104
# pylint: disable=unused-argument
@@ -125,8 +125,8 @@ async def wrap_connect_(
125125

126126

127127
def unwrap_connect():
128-
""""Disable integration with aiopg library.
129-
https://github.com/aio-libs/aiopg
128+
"""Disable integration with aiopg library.
129+
https://github.com/aio-libs/aiopg
130130
"""
131131

132132
unwrap(aiopg, "connect")
@@ -217,7 +217,7 @@ async def wrap_create_pool_(
217217

218218

219219
def unwrap_create_pool():
220-
""""Disable integration with aiopg library.
221-
https://github.com/aio-libs/aiopg
220+
"""Disable integration with aiopg library.
221+
https://github.com/aio-libs/aiopg
222222
"""
223223
unwrap(aiopg, "create_pool")

‎instrumentation/opentelemetry-instrumentation-aiopg/tests/test_aiopg_integration.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,40 @@ def test_span_succeeded(self):
230230
trace_api.status.StatusCanonicalCode.OK,
231231
)
232232

233+
def test_span_not_recording(self):
234+
connection_props = {
235+
"database": "testdatabase",
236+
"server_host": "testhost",
237+
"server_port": 123,
238+
"user": "testuser",
239+
}
240+
connection_attributes = {
241+
"database": "database",
242+
"port": "server_port",
243+
"host": "server_host",
244+
"user": "user",
245+
}
246+
mock_tracer = mock.Mock()
247+
mock_span = mock.Mock()
248+
mock_span.is_recording.return_value = False
249+
mock_tracer.start_span.return_value = mock_span
250+
mock_tracer.use_span.return_value.__enter__ = mock_span
251+
mock_tracer.use_span.return_value.__exit__ = True
252+
db_integration = AiopgIntegration(
253+
mock_tracer, "testcomponent", "testtype", connection_attributes
254+
)
255+
mock_connection = async_call(
256+
db_integration.wrapped_connection(
257+
mock_connect, {}, connection_props
258+
)
259+
)
260+
cursor = async_call(mock_connection.cursor())
261+
async_call(cursor.execute("Test query", ("param1Value", False)))
262+
self.assertFalse(mock_span.is_recording())
263+
self.assertTrue(mock_span.is_recording.called)
264+
self.assertFalse(mock_span.set_attribute.called)
265+
self.assertFalse(mock_span.set_status.called)
266+
233267
def test_span_failed(self):
234268
db_integration = AiopgIntegration(self.tracer, "testcomponent")
235269
mock_connection = async_call(

‎instrumentation/opentelemetry-instrumentation-asyncpg/src/opentelemetry/instrumentation/asyncpg/__init__.py

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -112,31 +112,34 @@ def _uninstrument(self, **__):
112112
unwrap(asyncpg.Connection, method)
113113

114114
async def _do_execute(self, func, instance, args, kwargs):
115-
span_attributes = _hydrate_span_from_args(
116-
instance, args[0], args[1:] if self.capture_parameters else None,
117-
)
118115
tracer = getattr(asyncpg, _APPLIED)
119116

120117
exception = None
121118

122119
with tracer.start_as_current_span(
123120
"postgresql", kind=SpanKind.CLIENT
124121
) as span:
125-
126-
for attribute, value in span_attributes.items():
127-
span.set_attribute(attribute, value)
122+
if span.is_recording():
123+
span_attributes = _hydrate_span_from_args(
124+
instance,
125+
args[0],
126+
args[1:] if self.capture_parameters else None,
127+
)
128+
for attribute, value in span_attributes.items():
129+
span.set_attribute(attribute, value)
128130

129131
try:
130132
result = await func(*args, **kwargs)
131133
except Exception as exc: # pylint: disable=W0703
132134
exception = exc
133135
raise
134136
finally:
135-
if exception is not None:
136-
span.set_status(
137-
Status(_exception_to_canonical_code(exception))
138-
)
139-
else:
140-
span.set_status(Status(StatusCanonicalCode.OK))
137+
if span.is_recording():
138+
if exception is not None:
139+
span.set_status(
140+
Status(_exception_to_canonical_code(exception))
141+
)
142+
else:
143+
span.set_status(Status(StatusCanonicalCode.OK))
141144

142145
return result

‎instrumentation/opentelemetry-instrumentation-boto/tests/test_boto_instrumentation.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ def test_not_recording(self):
8888
mock_span.is_recording.return_value = False
8989
mock_tracer.start_span.return_value = mock_span
9090
mock_tracer.use_span.return_value.__enter__ = mock_span
91-
mock_tracer.use_span.return_value.__exit__ = mock_span
91+
mock_tracer.use_span.return_value.__exit__ = True
9292
with patch("opentelemetry.trace.get_tracer") as tracer:
9393
tracer.return_value = mock_tracer
9494
ec2 = boto.ec2.connect_to_region("us-west-2")

‎instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_instrumentation.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ def test_not_recording(self):
8484
mock_span.is_recording.return_value = False
8585
mock_tracer.start_span.return_value = mock_span
8686
mock_tracer.use_span.return_value.__enter__ = mock_span
87-
mock_tracer.use_span.return_value.__exit__ = mock_span
87+
mock_tracer.use_span.return_value.__exit__ = True
8888
with patch("opentelemetry.trace.get_tracer") as tracer:
8989
tracer.return_value = mock_tracer
9090
ec2 = self.session.create_client("ec2", region_name="us-west-2")

‎instrumentation/opentelemetry-instrumentation-dbapi/src/opentelemetry/instrumentation/dbapi/__init__.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -311,6 +311,8 @@ def __init__(self, db_api_integration: DatabaseApiIntegration):
311311
def _populate_span(
312312
self, span: trace_api.Span, *args: typing.Tuple[typing.Any, typing.Any]
313313
):
314+
if not span.is_recording():
315+
return
314316
statement = args[0] if args else ""
315317
span.set_attribute(
316318
"component", self._db_api_integration.database_component
@@ -341,10 +343,14 @@ def traced_execution(
341343
self._populate_span(span, *args)
342344
try:
343345
result = query_method(*args, **kwargs)
344-
span.set_status(Status(StatusCanonicalCode.OK))
346+
if span.is_recording():
347+
span.set_status(Status(StatusCanonicalCode.OK))
345348
return result
346349
except Exception as ex: # pylint: disable=broad-except
347-
span.set_status(Status(StatusCanonicalCode.UNKNOWN, str(ex)))
350+
if span.is_recording():
351+
span.set_status(
352+
Status(StatusCanonicalCode.UNKNOWN, str(ex))
353+
)
348354
raise ex
349355

350356

‎instrumentation/opentelemetry-instrumentation-dbapi/tests/test_dbapi_integration.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,38 @@ def test_span_succeeded(self):
6969
trace_api.status.StatusCanonicalCode.OK,
7070
)
7171

72+
def test_span_not_recording(self):
73+
connection_props = {
74+
"database": "testdatabase",
75+
"server_host": "testhost",
76+
"server_port": 123,
77+
"user": "testuser",
78+
}
79+
connection_attributes = {
80+
"database": "database",
81+
"port": "server_port",
82+
"host": "server_host",
83+
"user": "user",
84+
}
85+
mock_tracer = mock.Mock()
86+
mock_span = mock.Mock()
87+
mock_span.is_recording.return_value = False
88+
mock_tracer.start_span.return_value = mock_span
89+
mock_tracer.use_span.return_value.__enter__ = mock_span
90+
mock_tracer.use_span.return_value.__exit__ = True
91+
db_integration = dbapi.DatabaseApiIntegration(
92+
mock_tracer, "testcomponent", "testtype", connection_attributes
93+
)
94+
mock_connection = db_integration.wrapped_connection(
95+
mock_connect, {}, connection_props
96+
)
97+
cursor = mock_connection.cursor()
98+
cursor.execute("Test query", ("param1Value", False))
99+
self.assertFalse(mock_span.is_recording())
100+
self.assertTrue(mock_span.is_recording.called)
101+
self.assertFalse(mock_span.set_attribute.called)
102+
self.assertFalse(mock_span.set_status.called)
103+
72104
def test_span_failed(self):
73105
db_integration = dbapi.DatabaseApiIntegration(
74106
self.tracer, "testcomponent"

‎instrumentation/opentelemetry-instrumentation-django/tests/test_middleware.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ def test_not_recording(self):
9595
mock_span.is_recording.return_value = False
9696
mock_tracer.start_span.return_value = mock_span
9797
mock_tracer.use_span.return_value.__enter__ = mock_span
98-
mock_tracer.use_span.return_value.__exit__ = mock_span
98+
mock_tracer.use_span.return_value.__exit__ = True
9999
with patch("opentelemetry.trace.get_tracer") as tracer:
100100
tracer.return_value = mock_tracer
101101
Client().get("/traced/")

‎instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,29 @@ def test_instrumentor(self, mock_connect):
6060
spans_list = self.memory_exporter.get_finished_spans()
6161
self.assertEqual(len(spans_list), 1)
6262

63+
@mock.patch("psycopg2.connect")
64+
# pylint: disable=unused-argument
65+
def test_not_recording(self, mock_connect):
66+
mock_tracer = mock.Mock()
67+
mock_span = mock.Mock()
68+
mock_span.is_recording.return_value = False
69+
mock_tracer.start_span.return_value = mock_span
70+
mock_tracer.use_span.return_value.__enter__ = mock_span
71+
mock_tracer.use_span.return_value.__exit__ = True
72+
Psycopg2Instrumentor().instrument()
73+
with mock.patch("opentelemetry.trace.get_tracer") as tracer:
74+
tracer.return_value = mock_tracer
75+
cnx = psycopg2.connect(database="test")
76+
cursor = cnx.cursor()
77+
query = "SELECT * FROM test"
78+
cursor.execute(query)
79+
self.assertFalse(mock_span.is_recording())
80+
self.assertTrue(mock_span.is_recording.called)
81+
self.assertFalse(mock_span.set_attribute.called)
82+
self.assertFalse(mock_span.set_status.called)
83+
84+
Psycopg2Instrumentor().uninstrument()
85+
6386
@mock.patch("psycopg2.connect")
6487
# pylint: disable=unused-argument
6588
def test_custom_tracer_provider(self, mock_connect):

‎instrumentation/opentelemetry-instrumentation-pymemcache/src/opentelemetry/instrumentation/pymemcache/__init__.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,13 +90,14 @@
9090

9191

9292
def _set_connection_attributes(span, instance):
93+
if not span.is_recording():
94+
return
9395
for key, value in _get_address_attributes(instance).items():
9496
span.set_attribute(key, value)
9597

9698

9799
def _with_tracer_wrapper(func):
98-
"""Helper for providing tracer for wrapper functions.
99-
"""
100+
"""Helper for providing tracer for wrapper functions."""
100101

101102
def _with_tracer(tracer, cmd):
102103
def wrapper(wrapped, instance, args, kwargs):

‎instrumentation/opentelemetry-instrumentation-pymemcache/tests/test_pymemcache.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
14+
from unittest import mock
1415

1516
import pymemcache
1617
from pymemcache.exceptions import (
@@ -84,6 +85,23 @@ def test_set_success(self):
8485

8586
self.check_spans(spans, 1, ["set key"])
8687

88+
def test_set_not_recording(self):
89+
mock_tracer = mock.Mock()
90+
mock_span = mock.Mock()
91+
mock_span.is_recording.return_value = False
92+
mock_tracer.start_span.return_value = mock_span
93+
mock_tracer.use_span.return_value.__enter__ = mock_span
94+
mock_tracer.use_span.return_value.__exit__ = True
95+
with mock.patch("opentelemetry.trace.get_tracer") as tracer:
96+
tracer.return_value = mock_tracer
97+
client = self.make_client([b"STORED\r\n"])
98+
result = client.set(b"key", b"value", noreply=False)
99+
self.assertTrue(result)
100+
self.assertFalse(mock_span.is_recording())
101+
self.assertTrue(mock_span.is_recording.called)
102+
self.assertFalse(mock_span.set_attribute.called)
103+
self.assertFalse(mock_span.set_status.called)
104+
87105
def test_get_many_none_found(self):
88106
client = self.make_client([b"END\r\n"])
89107
result = client.get_many([b"key1", b"key2"])

0 commit comments

Comments
 (0)