Skip to content

Commit 9562a81

Browse files
juliusgeojuliusgeo
authored
PYTHON-3119 getMore helper should explicitly send inherited comment (#904)
1 parent 861d795 commit 9562a81

14 files changed

+440
-11
lines changed

‎.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,3 +15,4 @@ pymongo.egg-info/
1515
.tox
1616
mongocryptd.pid
1717
.idea/
18+
.nova/

‎pymongo/aggregation.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,7 @@ def get_cursor(self, session, server, sock_info, read_preference):
174174
max_await_time_ms=self._max_await_time_ms,
175175
session=session,
176176
explicit_session=self._explicit_session,
177+
comment=self._options.get("comment"),
177178
)
178179
cmd_cursor._maybe_pin_connection(sock_info)
179180
return cmd_cursor

‎pymongo/collection.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2180,7 +2180,12 @@ def _cmd(session, server, sock_info, read_preference):
21802180
raise
21812181
cursor = {"id": 0, "firstBatch": []}
21822182
cmd_cursor = CommandCursor(
2183-
coll, cursor, sock_info.address, session=session, explicit_session=explicit_session
2183+
coll,
2184+
cursor,
2185+
sock_info.address,
2186+
session=session,
2187+
explicit_session=explicit_session,
2188+
comment=cmd.get("comment"),
21842189
)
21852190
cmd_cursor._maybe_pin_connection(sock_info)
21862191
return cmd_cursor

‎pymongo/command_cursor.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ def __init__(
4343
max_await_time_ms: Optional[int] = None,
4444
session: Optional["ClientSession"] = None,
4545
explicit_session: bool = False,
46+
comment: Any = None,
4647
) -> None:
4748
"""Create a new command cursor."""
4849
self.__sock_mgr: Any = None
@@ -56,6 +57,7 @@ def __init__(
5657
self.__session = session
5758
self.__explicit_session = explicit_session
5859
self.__killed = self.__id == 0
60+
self.__comment = comment
5961
if self.__killed:
6062
self.__end_session(True)
6163

@@ -224,6 +226,7 @@ def _refresh(self):
224226
self.__max_await_time_ms,
225227
self.__sock_mgr,
226228
False,
229+
self.__comment,
227230
)
228231
)
229232
else: # Cursor id is zero nothing else to return
@@ -314,6 +317,7 @@ def __init__(
314317
max_await_time_ms: Optional[int] = None,
315318
session: Optional["ClientSession"] = None,
316319
explicit_session: bool = False,
320+
comment: Any = None,
317321
) -> None:
318322
"""Create a new cursor / iterator over raw batches of BSON data.
319323
@@ -332,6 +336,7 @@ def __init__(
332336
max_await_time_ms,
333337
session,
334338
explicit_session,
339+
comment,
335340
)
336341

337342
def _unpack_response(

‎pymongo/cursor.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1183,6 +1183,7 @@ def _refresh(self):
11831183
self.__max_await_time_ms,
11841184
self.__sock_mgr,
11851185
self.__exhaust,
1186+
self.__comment,
11861187
)
11871188
self.__send_message(g)
11881189

‎pymongo/database.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -780,6 +780,7 @@ def _list_collections(self, sock_info, session, read_preference, **kwargs):
780780
sock_info.address,
781781
session=tmp_session,
782782
explicit_session=session is not None,
783+
comment=cmd.get("comment"),
783784
)
784785
cmd_cursor._maybe_pin_connection(sock_info)
785786
return cmd_cursor

‎pymongo/message.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -222,13 +222,15 @@ def _gen_find_command(
222222
return cmd
223223

224224

225-
def _gen_get_more_command(cursor_id, coll, batch_size, max_await_time_ms):
225+
def _gen_get_more_command(cursor_id, coll, batch_size, max_await_time_ms, comment, sock_info):
226226
"""Generate a getMore command document."""
227227
cmd = SON([("getMore", cursor_id), ("collection", coll)])
228228
if batch_size:
229229
cmd["batchSize"] = batch_size
230230
if max_await_time_ms is not None:
231231
cmd["maxTimeMS"] = max_await_time_ms
232+
if comment is not None and sock_info.max_wire_version >= 9:
233+
cmd["comment"] = comment
232234
return cmd
233235

234236

@@ -421,6 +423,7 @@ class _GetMore(object):
421423
"sock_mgr",
422424
"_as_command",
423425
"exhaust",
426+
"comment",
424427
)
425428

426429
name = "getMore"
@@ -438,6 +441,7 @@ def __init__(
438441
max_await_time_ms,
439442
sock_mgr,
440443
exhaust,
444+
comment,
441445
):
442446
self.db = db
443447
self.coll = coll
@@ -451,6 +455,7 @@ def __init__(
451455
self.sock_mgr = sock_mgr
452456
self._as_command = None
453457
self.exhaust = exhaust
458+
self.comment = comment
454459

455460
def namespace(self):
456461
return "%s.%s" % (self.db, self.coll)
@@ -473,9 +478,13 @@ def as_command(self, sock_info):
473478
return self._as_command
474479

475480
cmd = _gen_get_more_command(
476-
self.cursor_id, self.coll, self.ntoreturn, self.max_await_time_ms
481+
self.cursor_id,
482+
self.coll,
483+
self.ntoreturn,
484+
self.max_await_time_ms,
485+
self.comment,
486+
sock_info,
477487
)
478-
479488
if self.session:
480489
self.session._apply_to(cmd, False, self.read_preference, sock_info)
481490
sock_info.add_server_api(cmd)

‎pymongo/mongo_client.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1763,7 +1763,7 @@ def list_databases(
17631763
"firstBatch": res["databases"],
17641764
"ns": "admin.$cmd",
17651765
}
1766-
return CommandCursor(admin["$cmd"], cursor, None)
1766+
return CommandCursor(admin["$cmd"], cursor, None, comment=comment)
17671767

17681768
def list_database_names(
17691769
self,

‎test/change_streams/unified/change-streams.json

Lines changed: 179 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,185 @@
247247
]
248248
}
249249
]
250+
},
251+
{
252+
"description": "Test that comment is set on getMore",
253+
"runOnRequirements": [
254+
{
255+
"minServerVersion": "4.4.0",
256+
"topologies": [
257+
"replicaset"
258+
]
259+
}
260+
],
261+
"operations": [
262+
{
263+
"name": "createChangeStream",
264+
"object": "collection0",
265+
"arguments": {
266+
"pipeline": [],
267+
"comment": {
268+
"key": "value"
269+
}
270+
},
271+
"saveResultAsEntity": "changeStream0"
272+
},
273+
{
274+
"name": "insertOne",
275+
"object": "collection0",
276+
"arguments": {
277+
"document": {
278+
"_id": 1,
279+
"a": 1
280+
}
281+
}
282+
},
283+
{
284+
"name": "iterateUntilDocumentOrError",
285+
"object": "changeStream0"
286+
}
287+
],
288+
"expectEvents": [
289+
{
290+
"client": "client0",
291+
"events": [
292+
{
293+
"commandStartedEvent": {
294+
"command": {
295+
"aggregate": "collection0",
296+
"pipeline": [
297+
{
298+
"$changeStream": {}
299+
}
300+
],
301+
"comment": {
302+
"key": "value"
303+
}
304+
}
305+
}
306+
},
307+
{
308+
"commandStartedEvent": {
309+
"command": {
310+
"insert": "collection0",
311+
"documents": [
312+
{
313+
"_id": 1,
314+
"a": 1
315+
}
316+
]
317+
}
318+
}
319+
},
320+
{
321+
"commandStartedEvent": {
322+
"command": {
323+
"getMore": {
324+
"$$type": [
325+
"int",
326+
"long"
327+
]
328+
},
329+
"collection": "collection0",
330+
"comment": {
331+
"key": "value"
332+
}
333+
},
334+
"commandName": "getMore",
335+
"databaseName": "database0"
336+
}
337+
}
338+
]
339+
}
340+
]
341+
},
342+
{
343+
"description": "Test that comment is not set on getMore - pre 4.4",
344+
"runOnRequirements": [
345+
{
346+
"minServerVersion": "3.6.0",
347+
"maxServerVersion": "4.3.99",
348+
"topologies": [
349+
"replicaset"
350+
]
351+
}
352+
],
353+
"operations": [
354+
{
355+
"name": "createChangeStream",
356+
"object": "collection0",
357+
"arguments": {
358+
"pipeline": [],
359+
"comment": "comment"
360+
},
361+
"saveResultAsEntity": "changeStream0"
362+
},
363+
{
364+
"name": "insertOne",
365+
"object": "collection0",
366+
"arguments": {
367+
"document": {
368+
"_id": 1,
369+
"a": 1
370+
}
371+
}
372+
},
373+
{
374+
"name": "iterateUntilDocumentOrError",
375+
"object": "changeStream0"
376+
}
377+
],
378+
"expectEvents": [
379+
{
380+
"client": "client0",
381+
"events": [
382+
{
383+
"commandStartedEvent": {
384+
"command": {
385+
"aggregate": "collection0",
386+
"pipeline": [
387+
{
388+
"$changeStream": {}
389+
}
390+
],
391+
"comment": "comment"
392+
}
393+
}
394+
},
395+
{
396+
"commandStartedEvent": {
397+
"command": {
398+
"insert": "collection0",
399+
"documents": [
400+
{
401+
"_id": 1,
402+
"a": 1
403+
}
404+
]
405+
}
406+
}
407+
},
408+
{
409+
"commandStartedEvent": {
410+
"command": {
411+
"getMore": {
412+
"$$type": [
413+
"int",
414+
"long"
415+
]
416+
},
417+
"collection": "collection0",
418+
"comment": {
419+
"$$exists": false
420+
}
421+
},
422+
"commandName": "getMore",
423+
"databaseName": "database0"
424+
}
425+
}
426+
]
427+
}
428+
]
250429
}
251430
]
252431
}

0 commit comments

Comments
 (0)