@@ -57,14 +57,12 @@ def test_started(self):
57
57
# pylint: disable=protected-access
58
58
span = command_tracer ._pop_span (mock_event )
59
59
self .assertIs (span .kind , trace_api .SpanKind .CLIENT )
60
- self .assertEqual (span .name , "database_name.command_name " )
60
+ self .assertEqual (span .name , "database_name.find " )
61
61
self .assertEqual (span .attributes [SpanAttributes .DB_SYSTEM ], "mongodb" )
62
62
self .assertEqual (
63
63
span .attributes [SpanAttributes .DB_NAME ], "database_name"
64
64
)
65
- self .assertEqual (
66
- span .attributes [SpanAttributes .DB_STATEMENT ], "command_name"
67
- )
65
+ self .assertEqual (span .attributes [SpanAttributes .DB_STATEMENT ], "find" )
68
66
self .assertEqual (
69
67
span .attributes [SpanAttributes .NET_PEER_NAME ], "test.com"
70
68
)
@@ -181,7 +179,7 @@ def test_int_command(self):
181
179
182
180
self .assertEqual (len (spans_list ), 1 )
183
181
span = spans_list [0 ]
184
- self .assertEqual (span .name , "database_name.command_name " )
182
+ self .assertEqual (span .name , "database_name.123 " )
185
183
186
184
def test_no_op_tracer (self ):
187
185
mock_event = MockEvent ({})
@@ -194,6 +192,90 @@ def test_no_op_tracer(self):
194
192
spans_list = self .memory_exporter .get_finished_spans ()
195
193
self .assertEqual (len (spans_list ), 0 )
196
194
195
+ def test_capture_statement_getmore (self ):
196
+ command_attrs = {
197
+ "command_name" : "getMore" ,
198
+ "collection" : "test_collection" ,
199
+ }
200
+ mock_event = MockEvent (command_attrs )
201
+
202
+ command_tracer = CommandTracer (self .tracer , capture_statement = True )
203
+ command_tracer .started (event = mock_event )
204
+ command_tracer .succeeded (event = mock_event )
205
+
206
+ spans_list = self .memory_exporter .get_finished_spans ()
207
+
208
+ self .assertEqual (len (spans_list ), 1 )
209
+ span = spans_list [0 ]
210
+
211
+ self .assertEqual (
212
+ span .attributes [SpanAttributes .DB_STATEMENT ],
213
+ "getMore test_collection" ,
214
+ )
215
+
216
+ def test_capture_statement_aggregate (self ):
217
+ pipeline = [
218
+ {"$match" : {"status" : "active" }},
219
+ {"$group" : {"_id" : "$category" , "count" : {"$sum" : 1 }}},
220
+ ]
221
+ command_attrs = {
222
+ "command_name" : "aggregate" ,
223
+ "pipeline" : pipeline ,
224
+ }
225
+ command_tracer = CommandTracer (self .tracer , capture_statement = True )
226
+ mock_event = MockEvent (command_attrs )
227
+ command_tracer .started (event = mock_event )
228
+ command_tracer .succeeded (event = mock_event )
229
+
230
+ spans_list = self .memory_exporter .get_finished_spans ()
231
+
232
+ self .assertEqual (len (spans_list ), 1 )
233
+ span = spans_list [0 ]
234
+
235
+ expected_statement = f"aggregate { pipeline } "
236
+ self .assertEqual (
237
+ span .attributes [SpanAttributes .DB_STATEMENT ], expected_statement
238
+ )
239
+
240
+ def test_capture_statement_disabled_getmore (self ):
241
+ command_attrs = {
242
+ "command_name" : "getMore" ,
243
+ "collection" : "test_collection" ,
244
+ }
245
+ command_tracer = CommandTracer (self .tracer , capture_statement = False )
246
+ mock_event = MockEvent (command_attrs )
247
+ command_tracer .started (event = mock_event )
248
+ command_tracer .succeeded (event = mock_event )
249
+
250
+ spans_list = self .memory_exporter .get_finished_spans ()
251
+
252
+ self .assertEqual (len (spans_list ), 1 )
253
+ span = spans_list [0 ]
254
+
255
+ self .assertEqual (
256
+ span .attributes [SpanAttributes .DB_STATEMENT ], "getMore"
257
+ )
258
+
259
+ def test_capture_statement_disabled_aggregate (self ):
260
+ pipeline = [{"$match" : {"status" : "active" }}]
261
+ command_attrs = {
262
+ "command_name" : "aggregate" ,
263
+ "pipeline" : pipeline ,
264
+ }
265
+ command_tracer = CommandTracer (self .tracer , capture_statement = False )
266
+ mock_event = MockEvent (command_attrs )
267
+ command_tracer .started (event = mock_event )
268
+ command_tracer .succeeded (event = mock_event )
269
+
270
+ spans_list = self .memory_exporter .get_finished_spans ()
271
+
272
+ self .assertEqual (len (spans_list ), 1 )
273
+ span = spans_list [0 ]
274
+
275
+ self .assertEqual (
276
+ span .attributes [SpanAttributes .DB_STATEMENT ], "aggregate"
277
+ )
278
+
197
279
198
280
class MockCommand :
199
281
def __init__ (self , command_attrs ):
@@ -206,6 +288,7 @@ def get(self, key, default=""):
206
288
class MockEvent :
207
289
def __init__ (self , command_attrs , connection_id = None , request_id = "" ):
208
290
self .command = MockCommand (command_attrs )
291
+ self .command_name = self .command .get ("command_name" )
209
292
self .connection_id = connection_id
210
293
self .request_id = request_id
211
294
0 commit comments