Skip to content

Commit a05acca

Browse files
tammy-baylis-switammy-baylis-swi
committed
Add tests
1 parent 30d7b96 commit a05acca

File tree

1 file changed

+173
-1
lines changed

1 file changed

+173
-1
lines changed

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

Lines changed: 173 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,59 @@ def test_sqlcommenter_enabled(self):
6161
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
6262
)
6363

64-
def test_sqlcommenter_enabled_stmt_disabled_default_matches_db_statement_attribute(
64+
def test_sqlcommenter_default_stmt_enabled_no_comments_anywhere(self):
65+
engine = create_engine("sqlite:///:memory:")
66+
SQLAlchemyInstrumentor().instrument(
67+
engine=engine,
68+
tracer_provider=self.tracer_provider,
69+
# enable_commenter not set
70+
enable_attribute_commenter=True,
71+
)
72+
cnx = engine.connect()
73+
cnx.execute(text("SELECT 1;")).fetchall()
74+
query_log = self.caplog.records[-2].getMessage()
75+
self.assertEqual(
76+
query_log,
77+
"SELECT 1;",
78+
)
79+
spans = self.memory_exporter.get_finished_spans()
80+
self.assertEqual(len(spans), 2)
81+
# first span is connection to db
82+
self.assertEqual(spans[0].name, "connect")
83+
# second span is query itself
84+
query_span = spans[1]
85+
self.assertEqual(
86+
query_span.attributes[SpanAttributes.DB_STATEMENT],
87+
"SELECT 1;",
88+
)
89+
90+
def test_sqlcommenter_disabled_stmt_enabled_no_comments_anywhere(self):
91+
engine = create_engine("sqlite:///:memory:")
92+
SQLAlchemyInstrumentor().instrument(
93+
engine=engine,
94+
tracer_provider=self.tracer_provider,
95+
enable_commenter=False,
96+
enable_attribute_commenter=True,
97+
)
98+
cnx = engine.connect()
99+
cnx.execute(text("SELECT 1;")).fetchall()
100+
query_log = self.caplog.records[-2].getMessage()
101+
self.assertEqual(
102+
query_log,
103+
"SELECT 1;",
104+
)
105+
spans = self.memory_exporter.get_finished_spans()
106+
self.assertEqual(len(spans), 2)
107+
# first span is connection to db
108+
self.assertEqual(spans[0].name, "connect")
109+
# second span is query itself
110+
query_span = spans[1]
111+
self.assertEqual(
112+
query_span.attributes[SpanAttributes.DB_STATEMENT],
113+
"SELECT 1;",
114+
)
115+
116+
def test_sqlcommenter_enabled_stmt_disabled_default(
65117
self,
66118
):
67119
engine = create_engine("sqlite:///:memory:")
@@ -142,6 +194,45 @@ def test_sqlcommenter_enabled_otel_values_false(self):
142194
self.caplog.records[-2].getMessage(),
143195
r"SELECT 1 /\*db_driver='(.*)'\*/;",
144196
)
197+
spans = self.memory_exporter.get_finished_spans()
198+
self.assertEqual(len(spans), 2)
199+
# first span is connection to db
200+
self.assertEqual(spans[0].name, "connect")
201+
# second span is query itself
202+
query_span = spans[1]
203+
self.assertEqual(
204+
query_span.attributes[SpanAttributes.DB_STATEMENT],
205+
r"SELECT 1;",
206+
)
207+
208+
def test_sqlcommenter_enabled_stmt_enabled_otel_values_false(self):
209+
engine = create_engine("sqlite:///:memory:")
210+
SQLAlchemyInstrumentor().instrument(
211+
engine=engine,
212+
tracer_provider=self.tracer_provider,
213+
enable_commenter=True,
214+
commenter_options={
215+
"db_framework": False,
216+
"opentelemetry_values": False,
217+
},
218+
enable_attribute_commenter=True,
219+
)
220+
cnx = engine.connect()
221+
cnx.execute(text("SELECT 1;")).fetchall()
222+
self.assertRegex(
223+
self.caplog.records[-2].getMessage(),
224+
r"SELECT 1 /\*db_driver='(.*)'\*/;",
225+
)
226+
spans = self.memory_exporter.get_finished_spans()
227+
self.assertEqual(len(spans), 2)
228+
# first span is connection to db
229+
self.assertEqual(spans[0].name, "connect")
230+
# second span is query itself
231+
query_span = spans[1]
232+
self.assertRegex(
233+
query_span.attributes[SpanAttributes.DB_STATEMENT],
234+
r"SELECT 1 /\*db_driver='(.*)'*/;",
235+
)
145236

146237
def test_sqlcommenter_flask_integration(self):
147238
engine = create_engine("sqlite:///:memory:")
@@ -164,6 +255,49 @@ def test_sqlcommenter_flask_integration(self):
164255
self.caplog.records[-2].getMessage(),
165256
r"SELECT 1 /\*db_driver='(.*)',flask=1,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
166257
)
258+
spans = self.memory_exporter.get_finished_spans()
259+
self.assertEqual(len(spans), 2)
260+
# first span is connection to db
261+
self.assertEqual(spans[0].name, "connect")
262+
# second span is query itself
263+
query_span = spans[1]
264+
self.assertEqual(
265+
query_span.attributes[SpanAttributes.DB_STATEMENT],
266+
"SELECT 1;",
267+
)
268+
269+
def test_sqlcommenter_stmt_enabled_flask_integration(self):
270+
engine = create_engine("sqlite:///:memory:")
271+
SQLAlchemyInstrumentor().instrument(
272+
engine=engine,
273+
tracer_provider=self.tracer_provider,
274+
enable_commenter=True,
275+
commenter_options={"db_framework": False},
276+
enable_attribute_commenter=True,
277+
)
278+
cnx = engine.connect()
279+
280+
current_context = context.get_current()
281+
sqlcommenter_context = context.set_value(
282+
"SQLCOMMENTER_ORM_TAGS_AND_VALUES", {"flask": 1}, current_context
283+
)
284+
context.attach(sqlcommenter_context)
285+
286+
cnx.execute(text("SELECT 1;")).fetchall()
287+
self.assertRegex(
288+
self.caplog.records[-2].getMessage(),
289+
r"SELECT 1 /\*db_driver='(.*)',flask=1,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
290+
)
291+
spans = self.memory_exporter.get_finished_spans()
292+
self.assertEqual(len(spans), 2)
293+
# first span is connection to db
294+
self.assertEqual(spans[0].name, "connect")
295+
# second span is query itself
296+
query_span = spans[1]
297+
self.assertRegex(
298+
query_span.attributes[SpanAttributes.DB_STATEMENT],
299+
r"SELECT 1 /\*db_driver='(.*)',flask=1,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
300+
)
167301

168302
def test_sqlcommenter_enabled_create_engine_after_instrumentation(self):
169303
SQLAlchemyInstrumentor().instrument(
@@ -179,6 +313,44 @@ def test_sqlcommenter_enabled_create_engine_after_instrumentation(self):
179313
self.caplog.records[-2].getMessage(),
180314
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
181315
)
316+
spans = self.memory_exporter.get_finished_spans()
317+
self.assertEqual(len(spans), 2)
318+
# first span is connection to db
319+
self.assertEqual(spans[0].name, "connect")
320+
# second span is query itself
321+
query_span = spans[1]
322+
self.assertEqual(
323+
query_span.attributes[SpanAttributes.DB_STATEMENT],
324+
"SELECT 1;",
325+
)
326+
327+
def test_sqlcommenter_enabled_stmt_enabled_create_engine_after_instrumentation(
328+
self,
329+
):
330+
SQLAlchemyInstrumentor().instrument(
331+
tracer_provider=self.tracer_provider,
332+
enable_commenter=True,
333+
enable_attribute_commenter=True,
334+
)
335+
from sqlalchemy import create_engine # pylint: disable-all
336+
337+
engine = create_engine("sqlite:///:memory:")
338+
cnx = engine.connect()
339+
cnx.execute(text("SELECT 1;")).fetchall()
340+
self.assertRegex(
341+
self.caplog.records[-2].getMessage(),
342+
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
343+
)
344+
spans = self.memory_exporter.get_finished_spans()
345+
self.assertEqual(len(spans), 2)
346+
# first span is connection to db
347+
self.assertEqual(spans[0].name, "connect")
348+
# second span is query itself
349+
query_span = spans[1]
350+
self.assertRegex(
351+
query_span.attributes[SpanAttributes.DB_STATEMENT],
352+
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
353+
)
182354

183355
def test_sqlcommenter_disabled_create_engine_after_instrumentation(self):
184356
SQLAlchemyInstrumentor().instrument(

0 commit comments

Comments
 (0)