Skip to content

Commit 2ce69a6

Browse files
shaharglshahargl
authored
Add span for connection phase (#1134)
1 parent 9e2dbec commit 2ce69a6

File tree

9 files changed

+100
-36
lines changed

9 files changed

+100
-36
lines changed

‎CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
1212
### Added
1313
- `opentelemetry-instrumentation-redis` add support to instrument RedisCluster clients
1414
([#1177](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1177))
15+
- `opentelemetry-instrumentation-sqlalchemy` Added span for the connection phase ([#1133](https://github.com/open-telemetry/opentelemetry-python-contrib/issues/1133))
1516

1617
## [1.12.0rc2-0.32b0](https://github.com/open-telemetry/opentelemetry-python/releases/tag/v1.12.0rc2-0.32b0) - 2022-07-01
1718

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,12 +56,14 @@
5656

5757
import sqlalchemy
5858
from packaging.version import parse as parse_version
59+
from sqlalchemy.engine.base import Engine
5960
from wrapt import wrap_function_wrapper as _w
6061

6162
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
6263
from opentelemetry.instrumentation.sqlalchemy.engine import (
6364
EngineTracer,
6465
_get_tracer,
66+
_wrap_connect,
6567
_wrap_create_async_engine,
6668
_wrap_create_engine,
6769
)
@@ -97,13 +99,17 @@ def _instrument(self, **kwargs):
9799
"create_engine",
98100
_wrap_create_engine(tracer_provider),
99101
)
102+
_w(
103+
"sqlalchemy.engine.base",
104+
"Engine.connect",
105+
_wrap_connect(tracer_provider),
106+
)
100107
if parse_version(sqlalchemy.__version__).release >= (1, 4):
101108
_w(
102109
"sqlalchemy.ext.asyncio",
103110
"create_async_engine",
104111
_wrap_create_async_engine(tracer_provider),
105112
)
106-
107113
if kwargs.get("engine") is not None:
108114
return EngineTracer(
109115
_get_tracer(tracer_provider),
@@ -127,5 +133,6 @@ def _instrument(self, **kwargs):
127133
def _uninstrument(self, **kwargs):
128134
unwrap(sqlalchemy, "create_engine")
129135
unwrap(sqlalchemy.engine, "create_engine")
136+
unwrap(Engine, "connect")
130137
if parse_version(sqlalchemy.__version__).release >= (1, 4):
131138
unwrap(sqlalchemy.ext.asyncio, "create_async_engine")

‎instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,23 @@ def _wrap_create_engine_internal(func, module, args, kwargs):
7777
return _wrap_create_engine_internal
7878

7979

80+
def _wrap_connect(tracer_provider=None):
81+
tracer = trace.get_tracer(
82+
_instrumenting_module_name,
83+
__version__,
84+
tracer_provider=tracer_provider,
85+
)
86+
87+
# pylint: disable=unused-argument
88+
def _wrap_connect_internal(func, module, args, kwargs):
89+
with tracer.start_as_current_span(
90+
"connect", kind=trace.SpanKind.CLIENT
91+
):
92+
return func(*args, **kwargs)
93+
94+
return _wrap_connect_internal
95+
96+
8097
class EngineTracer:
8198
def __init__(self, tracer, engine, enable_commenter=False):
8299
self.tracer = tracer

‎instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py

Lines changed: 46 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -46,9 +46,13 @@ def test_trace_integration(self):
4646
cnx.execute("SELECT 1 + 1;").fetchall()
4747
spans = self.memory_exporter.get_finished_spans()
4848

49-
self.assertEqual(len(spans), 1)
50-
self.assertEqual(spans[0].name, "SELECT :memory:")
49+
self.assertEqual(len(spans), 2)
50+
# first span - the connection to the db
51+
self.assertEqual(spans[0].name, "connect")
5152
self.assertEqual(spans[0].kind, trace.SpanKind.CLIENT)
53+
# second span - the query itself
54+
self.assertEqual(spans[1].name, "SELECT :memory:")
55+
self.assertEqual(spans[1].kind, trace.SpanKind.CLIENT)
5256

5357
def test_instrument_two_engines(self):
5458
engine_1 = create_engine("sqlite:///:memory:")
@@ -65,8 +69,20 @@ def test_instrument_two_engines(self):
6569
cnx_2.execute("SELECT 1 + 1;").fetchall()
6670

6771
spans = self.memory_exporter.get_finished_spans()
72+
# 2 queries + 2 engine connect
73+
self.assertEqual(len(spans), 4)
6874

69-
self.assertEqual(len(spans), 2)
75+
def test_instrument_engine_connect(self):
76+
engine = create_engine("sqlite:///:memory:")
77+
78+
SQLAlchemyInstrumentor().instrument(
79+
engine=engine,
80+
tracer_provider=self.tracer_provider,
81+
)
82+
83+
engine.connect()
84+
spans = self.memory_exporter.get_finished_spans()
85+
self.assertEqual(len(spans), 1)
7086

7187
@pytest.mark.skipif(
7288
not sqlalchemy.__version__.startswith("1.4"),
@@ -85,11 +101,15 @@ async def run():
85101
async with engine.connect() as cnx:
86102
await cnx.execute(sqlalchemy.text("SELECT 1 + 1;"))
87103
spans = self.memory_exporter.get_finished_spans()
88-
self.assertEqual(len(spans), 1)
89-
self.assertEqual(spans[0].name, "SELECT :memory:")
104+
self.assertEqual(len(spans), 2)
105+
# first span - the connection to the db
106+
self.assertEqual(spans[0].name, "connect")
90107
self.assertEqual(spans[0].kind, trace.SpanKind.CLIENT)
108+
# second span - the query
109+
self.assertEqual(spans[1].name, "SELECT :memory:")
110+
self.assertEqual(spans[1].kind, trace.SpanKind.CLIENT)
91111
self.assertEqual(
92-
spans[0].instrumentation_scope.name,
112+
spans[1].instrumentation_scope.name,
93113
"opentelemetry.instrumentation.sqlalchemy",
94114
)
95115

@@ -99,7 +119,10 @@ def test_not_recording(self):
99119
mock_tracer = mock.Mock()
100120
mock_span = mock.Mock()
101121
mock_span.is_recording.return_value = False
122+
mock_span.__enter__ = mock.Mock(return_value=(mock.Mock(), None))
123+
mock_span.__exit__ = mock.Mock(return_value=None)
102124
mock_tracer.start_span.return_value = mock_span
125+
mock_tracer.start_as_current_span.return_value = mock_span
103126
with mock.patch("opentelemetry.trace.get_tracer") as tracer:
104127
tracer.return_value = mock_tracer
105128
engine = create_engine("sqlite:///:memory:")
@@ -123,11 +146,15 @@ def test_create_engine_wrapper(self):
123146
cnx.execute("SELECT 1 + 1;").fetchall()
124147
spans = self.memory_exporter.get_finished_spans()
125148

126-
self.assertEqual(len(spans), 1)
127-
self.assertEqual(spans[0].name, "SELECT :memory:")
149+
self.assertEqual(len(spans), 2)
150+
# first span - the connection to the db
151+
self.assertEqual(spans[0].name, "connect")
128152
self.assertEqual(spans[0].kind, trace.SpanKind.CLIENT)
153+
# second span - the query
154+
self.assertEqual(spans[1].name, "SELECT :memory:")
155+
self.assertEqual(spans[1].kind, trace.SpanKind.CLIENT)
129156
self.assertEqual(
130-
spans[0].instrumentation_scope.name,
157+
spans[1].instrumentation_scope.name,
131158
"opentelemetry.instrumentation.sqlalchemy",
132159
)
133160

@@ -153,7 +180,7 @@ def test_custom_tracer_provider(self):
153180
cnx.execute("SELECT 1 + 1;").fetchall()
154181
spans = self.memory_exporter.get_finished_spans()
155182

156-
self.assertEqual(len(spans), 1)
183+
self.assertEqual(len(spans), 2)
157184
self.assertEqual(spans[0].resource.attributes["service.name"], "test")
158185
self.assertEqual(
159186
spans[0].resource.attributes["deployment.environment"], "env"
@@ -177,11 +204,15 @@ async def run():
177204
async with engine.connect() as cnx:
178205
await cnx.execute(sqlalchemy.text("SELECT 1 + 1;"))
179206
spans = self.memory_exporter.get_finished_spans()
180-
self.assertEqual(len(spans), 1)
181-
self.assertEqual(spans[0].name, "SELECT :memory:")
207+
self.assertEqual(len(spans), 2)
208+
# first span - the connection to the db
209+
self.assertEqual(spans[0].name, "connect")
182210
self.assertEqual(spans[0].kind, trace.SpanKind.CLIENT)
211+
# second span - the query
212+
self.assertEqual(spans[1].name, "SELECT :memory:")
213+
self.assertEqual(spans[1].kind, trace.SpanKind.CLIENT)
183214
self.assertEqual(
184-
spans[0].instrumentation_scope.name,
215+
spans[1].instrumentation_scope.name,
185216
"opentelemetry.instrumentation.sqlalchemy",
186217
)
187218

@@ -199,8 +230,8 @@ def test_generate_commenter(self):
199230
cnx = engine.connect()
200231
cnx.execute("SELECT 1 + 1;").fetchall()
201232
spans = self.memory_exporter.get_finished_spans()
202-
self.assertEqual(len(spans), 1)
203-
span = spans[0]
233+
self.assertEqual(len(spans), 2)
234+
span = spans[1]
204235
self.assertIn(
205236
EngineTracer._generate_comment(span),
206237
self.caplog.records[-2].getMessage(),

‎tests/opentelemetry-docker-tests/tests/sqlalchemy_tests/mixins.py

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -128,8 +128,9 @@ def test_orm_insert(self):
128128
self.session.commit()
129129

130130
spans = self.memory_exporter.get_finished_spans()
131-
self.assertEqual(len(spans), 1)
132-
span = spans[0]
131+
# one span for the connection and one for the query
132+
self.assertEqual(len(spans), 2)
133+
span = spans[1]
133134
stmt = "INSERT INTO players (id, name) VALUES "
134135
if span.attributes.get(SpanAttributes.DB_SYSTEM) == "sqlite":
135136
stmt += "(?, ?)"
@@ -148,8 +149,9 @@ def test_session_query(self):
148149
self.assertEqual(len(out), 0)
149150

150151
spans = self.memory_exporter.get_finished_spans()
151-
self.assertEqual(len(spans), 1)
152-
span = spans[0]
152+
# one span for the connection and one for the query
153+
self.assertEqual(len(spans), 2)
154+
span = spans[1]
153155
stmt = "SELECT players.id AS players_id, players.name AS players_name \nFROM players \nWHERE players.name = "
154156
if span.attributes.get(SpanAttributes.DB_SYSTEM) == "sqlite":
155157
stmt += "?"
@@ -170,8 +172,9 @@ def test_engine_connect_execute(self):
170172
self.assertEqual(len(rows), 0)
171173

172174
spans = self.memory_exporter.get_finished_spans()
173-
self.assertEqual(len(spans), 1)
174-
span = spans[0]
175+
# one span for the connection and one for the query
176+
self.assertEqual(len(spans), 2)
177+
span = spans[1]
175178
self._check_span(span, "SELECT")
176179
self.assertEqual(
177180
span.attributes.get(SpanAttributes.DB_STATEMENT),
@@ -190,8 +193,9 @@ def test_parent(self):
190193
self.assertEqual(len(rows), 0)
191194

192195
spans = self.memory_exporter.get_finished_spans()
193-
self.assertEqual(len(spans), 2)
194-
child_span, parent_span = spans
196+
# one span for the connection and two for the queries
197+
self.assertEqual(len(spans), 3)
198+
_, child_span, parent_span = spans
195199

196200
# confirm the parenting
197201
self.assertIsNone(parent_span.parent)
@@ -247,5 +251,5 @@ def insert_players(session):
247251
# batch inserts together which means `insert_players` only generates one span.
248252
# See https://docs.sqlalchemy.org/en/14/changelog/migration_14.html#orm-batch-inserts-with-psycopg2-now-batch-statements-with-returning-in-most-cases
249253
self.assertEqual(
250-
len(spans), 5 if self.VENDOR not in ["postgresql"] else 3
254+
len(spans), 8 if self.VENDOR not in ["postgresql"] else 6
251255
)

‎tests/opentelemetry-docker-tests/tests/sqlalchemy_tests/test_mssql.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -69,8 +69,9 @@ def test_engine_execute_errors(self):
6969
conn.execute("SELECT * FROM a_wrong_table").fetchall()
7070

7171
spans = self.memory_exporter.get_finished_spans()
72-
self.assertEqual(len(spans), 1)
73-
span = spans[0]
72+
# one span for the connection and one for the query
73+
self.assertEqual(len(spans), 2)
74+
span = spans[1]
7475
# span fields
7576
self.assertEqual(span.name, "SELECT opentelemetry-tests")
7677
self.assertEqual(
@@ -96,9 +97,9 @@ def test_orm_insert(self):
9697
self.session.commit()
9798

9899
spans = self.memory_exporter.get_finished_spans()
99-
# identity insert on before the insert, insert, and identity insert off after the insert
100-
self.assertEqual(len(spans), 3)
101-
span = spans[1]
100+
# connect, identity insert on before the insert, insert, and identity insert off after the insert
101+
self.assertEqual(len(spans), 4)
102+
span = spans[2]
102103
self._check_span(span, "INSERT")
103104
self.assertIn(
104105
"INSERT INTO players",

‎tests/opentelemetry-docker-tests/tests/sqlalchemy_tests/test_mysql.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,9 @@ def test_engine_execute_errors(self):
6868
conn.execute("SELECT * FROM a_wrong_table").fetchall()
6969

7070
spans = self.memory_exporter.get_finished_spans()
71-
self.assertEqual(len(spans), 1)
72-
span = spans[0]
71+
# one span for the connection and one for the query
72+
self.assertEqual(len(spans), 2)
73+
span = spans[1]
7374
# span fields
7475
self.assertEqual(span.name, "SELECT opentelemetry-tests")
7576
self.assertEqual(

‎tests/opentelemetry-docker-tests/tests/sqlalchemy_tests/test_postgres.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,8 +62,9 @@ def test_engine_execute_errors(self):
6262
conn.execute("SELECT * FROM a_wrong_table").fetchall()
6363

6464
spans = self.memory_exporter.get_finished_spans()
65-
self.assertEqual(len(spans), 1)
66-
span = spans[0]
65+
# one span for the connection and one for the query
66+
self.assertEqual(len(spans), 2)
67+
span = spans[1]
6768
# span fields
6869
self.assertEqual(span.name, "SELECT opentelemetry-tests")
6970
self.assertEqual(

‎tests/opentelemetry-docker-tests/tests/sqlalchemy_tests/test_sqlite.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,9 @@ def test_engine_execute_errors(self):
3838
conn.execute(stmt).fetchall()
3939

4040
spans = self.memory_exporter.get_finished_spans()
41-
self.assertEqual(len(spans), 1)
42-
span = spans[0]
41+
# one span for the connection and one span for the query
42+
self.assertEqual(len(spans), 2)
43+
span = spans[1]
4344
# span fields
4445
self.assertEqual(span.name, "SELECT :memory:")
4546
self.assertEqual(

0 commit comments

Comments
 (0)