diff --git a/.gitignore b/.gitignore index de435d109e..f7ad6563ff 100644 --- a/.gitignore +++ b/.gitignore @@ -15,3 +15,4 @@ pymongo.egg-info/ .tox mongocryptd.pid .idea/ +.nova/ diff --git a/pymongo/aggregation.py b/pymongo/aggregation.py index 84ecffe5fb..62fe4bd055 100644 --- a/pymongo/aggregation.py +++ b/pymongo/aggregation.py @@ -174,6 +174,7 @@ def get_cursor(self, session, server, sock_info, read_preference): max_await_time_ms=self._max_await_time_ms, session=session, explicit_session=self._explicit_session, + comment=self._options.get("comment"), ) cmd_cursor._maybe_pin_connection(sock_info) return cmd_cursor diff --git a/pymongo/collection.py b/pymongo/collection.py index d0ebd9311a..3de1210522 100644 --- a/pymongo/collection.py +++ b/pymongo/collection.py @@ -2180,7 +2180,12 @@ def _cmd(session, server, sock_info, read_preference): raise cursor = {"id": 0, "firstBatch": []} cmd_cursor = CommandCursor( - coll, cursor, sock_info.address, session=session, explicit_session=explicit_session + coll, + cursor, + sock_info.address, + session=session, + explicit_session=explicit_session, + comment=cmd.get("comment"), ) cmd_cursor._maybe_pin_connection(sock_info) return cmd_cursor diff --git a/pymongo/command_cursor.py b/pymongo/command_cursor.py index 0bd99f0bbb..6f3f244419 100644 --- a/pymongo/command_cursor.py +++ b/pymongo/command_cursor.py @@ -43,6 +43,7 @@ def __init__( max_await_time_ms: Optional[int] = None, session: Optional["ClientSession"] = None, explicit_session: bool = False, + comment: Any = None, ) -> None: """Create a new command cursor.""" self.__sock_mgr: Any = None @@ -56,6 +57,7 @@ def __init__( self.__session = session self.__explicit_session = explicit_session self.__killed = self.__id == 0 + self.__comment = comment if self.__killed: self.__end_session(True) @@ -224,6 +226,7 @@ def _refresh(self): self.__max_await_time_ms, self.__sock_mgr, False, + self.__comment, ) ) else: # Cursor id is zero nothing else to return @@ -314,6 +317,7 @@ def __init__( max_await_time_ms: Optional[int] = None, session: Optional["ClientSession"] = None, explicit_session: bool = False, + comment: Any = None, ) -> None: """Create a new cursor / iterator over raw batches of BSON data. @@ -332,6 +336,7 @@ def __init__( max_await_time_ms, session, explicit_session, + comment, ) def _unpack_response( diff --git a/pymongo/cursor.py b/pymongo/cursor.py index 9f6f0898b4..350cc255bb 100644 --- a/pymongo/cursor.py +++ b/pymongo/cursor.py @@ -1183,6 +1183,7 @@ def _refresh(self): self.__max_await_time_ms, self.__sock_mgr, self.__exhaust, + self.__comment, ) self.__send_message(g) diff --git a/pymongo/database.py b/pymongo/database.py index 17cba06b65..d3d1b274fd 100644 --- a/pymongo/database.py +++ b/pymongo/database.py @@ -780,6 +780,7 @@ def _list_collections(self, sock_info, session, read_preference, **kwargs): sock_info.address, session=tmp_session, explicit_session=session is not None, + comment=cmd.get("comment"), ) cmd_cursor._maybe_pin_connection(sock_info) return cmd_cursor diff --git a/pymongo/message.py b/pymongo/message.py index 6aa8e4e7f9..1fdf0ece35 100644 --- a/pymongo/message.py +++ b/pymongo/message.py @@ -222,13 +222,15 @@ def _gen_find_command( return cmd -def _gen_get_more_command(cursor_id, coll, batch_size, max_await_time_ms): +def _gen_get_more_command(cursor_id, coll, batch_size, max_await_time_ms, comment, sock_info): """Generate a getMore command document.""" cmd = SON([("getMore", cursor_id), ("collection", coll)]) if batch_size: cmd["batchSize"] = batch_size if max_await_time_ms is not None: cmd["maxTimeMS"] = max_await_time_ms + if comment is not None and sock_info.max_wire_version >= 9: + cmd["comment"] = comment return cmd @@ -421,6 +423,7 @@ class _GetMore(object): "sock_mgr", "_as_command", "exhaust", + "comment", ) name = "getMore" @@ -438,6 +441,7 @@ def __init__( max_await_time_ms, sock_mgr, exhaust, + comment, ): self.db = db self.coll = coll @@ -451,6 +455,7 @@ def __init__( self.sock_mgr = sock_mgr self._as_command = None self.exhaust = exhaust + self.comment = comment def namespace(self): return "%s.%s" % (self.db, self.coll) @@ -473,9 +478,13 @@ def as_command(self, sock_info): return self._as_command cmd = _gen_get_more_command( - self.cursor_id, self.coll, self.ntoreturn, self.max_await_time_ms + self.cursor_id, + self.coll, + self.ntoreturn, + self.max_await_time_ms, + self.comment, + sock_info, ) - if self.session: self.session._apply_to(cmd, False, self.read_preference, sock_info) sock_info.add_server_api(cmd) diff --git a/pymongo/mongo_client.py b/pymongo/mongo_client.py index 280818ce00..83295fccc9 100644 --- a/pymongo/mongo_client.py +++ b/pymongo/mongo_client.py @@ -1763,7 +1763,7 @@ def list_databases( "firstBatch": res["databases"], "ns": "admin.$cmd", } - return CommandCursor(admin["$cmd"], cursor, None) + return CommandCursor(admin["$cmd"], cursor, None, comment=comment) def list_database_names( self, diff --git a/test/change_streams/unified/change-streams.json b/test/change_streams/unified/change-streams.json index 4aea9a4aa1..5fd2544ce0 100644 --- a/test/change_streams/unified/change-streams.json +++ b/test/change_streams/unified/change-streams.json @@ -247,6 +247,185 @@ ] } ] + }, + { + "description": "Test that comment is set on getMore", + "runOnRequirements": [ + { + "minServerVersion": "4.4.0", + "topologies": [ + "replicaset" + ] + } + ], + "operations": [ + { + "name": "createChangeStream", + "object": "collection0", + "arguments": { + "pipeline": [], + "comment": { + "key": "value" + } + }, + "saveResultAsEntity": "changeStream0" + }, + { + "name": "insertOne", + "object": "collection0", + "arguments": { + "document": { + "_id": 1, + "a": 1 + } + } + }, + { + "name": "iterateUntilDocumentOrError", + "object": "changeStream0" + } + ], + "expectEvents": [ + { + "client": "client0", + "events": [ + { + "commandStartedEvent": { + "command": { + "aggregate": "collection0", + "pipeline": [ + { + "$changeStream": {} + } + ], + "comment": { + "key": "value" + } + } + } + }, + { + "commandStartedEvent": { + "command": { + "insert": "collection0", + "documents": [ + { + "_id": 1, + "a": 1 + } + ] + } + } + }, + { + "commandStartedEvent": { + "command": { + "getMore": { + "$$type": [ + "int", + "long" + ] + }, + "collection": "collection0", + "comment": { + "key": "value" + } + }, + "commandName": "getMore", + "databaseName": "database0" + } + } + ] + } + ] + }, + { + "description": "Test that comment is not set on getMore - pre 4.4", + "runOnRequirements": [ + { + "minServerVersion": "3.6.0", + "maxServerVersion": "4.3.99", + "topologies": [ + "replicaset" + ] + } + ], + "operations": [ + { + "name": "createChangeStream", + "object": "collection0", + "arguments": { + "pipeline": [], + "comment": "comment" + }, + "saveResultAsEntity": "changeStream0" + }, + { + "name": "insertOne", + "object": "collection0", + "arguments": { + "document": { + "_id": 1, + "a": 1 + } + } + }, + { + "name": "iterateUntilDocumentOrError", + "object": "changeStream0" + } + ], + "expectEvents": [ + { + "client": "client0", + "events": [ + { + "commandStartedEvent": { + "command": { + "aggregate": "collection0", + "pipeline": [ + { + "$changeStream": {} + } + ], + "comment": "comment" + } + } + }, + { + "commandStartedEvent": { + "command": { + "insert": "collection0", + "documents": [ + { + "_id": 1, + "a": 1 + } + ] + } + } + }, + { + "commandStartedEvent": { + "command": { + "getMore": { + "$$type": [ + "int", + "long" + ] + }, + "collection": "collection0", + "comment": { + "$$exists": false + } + }, + "commandName": "getMore", + "databaseName": "database0" + } + } + ] + } + ] } ] } diff --git a/test/crud/unified/aggregate.json b/test/crud/unified/aggregate.json index f6da8ff32f..0cbfb4e6e9 100644 --- a/test/crud/unified/aggregate.json +++ b/test/crud/unified/aggregate.json @@ -327,10 +327,131 @@ ] }, { - "description": "aggregate with comment does not set comment on getMore", + "description": "aggregate with comment sets comment on getMore", "runOnRequirements": [ { - "minServerVersion": "3.6.0" + "minServerVersion": "4.4.0" + } + ], + "operations": [ + { + "name": "aggregate", + "arguments": { + "pipeline": [ + { + "$match": { + "_id": { + "$gt": 1 + } + } + } + ], + "batchSize": 2, + "comment": { + "content": "test" + } + }, + "object": "collection0", + "expectResult": [ + { + "_id": 2, + "x": 22 + }, + { + "_id": 3, + "x": 33 + }, + { + "_id": 4, + "x": 44 + }, + { + "_id": 5, + "x": 55 + }, + { + "_id": 6, + "x": 66 + } + ] + } + ], + "expectEvents": [ + { + "client": "client0", + "events": [ + { + "commandStartedEvent": { + "command": { + "aggregate": "coll0", + "pipeline": [ + { + "$match": { + "_id": { + "$gt": 1 + } + } + } + ], + "cursor": { + "batchSize": 2 + }, + "comment": { + "content": "test" + } + }, + "commandName": "aggregate", + "databaseName": "aggregate-tests" + } + }, + { + "commandStartedEvent": { + "command": { + "getMore": { + "$$type": [ + "int", + "long" + ] + }, + "collection": "coll0", + "batchSize": 2, + "comment": { + "content": "test" + } + }, + "commandName": "getMore", + "databaseName": "aggregate-tests" + } + }, + { + "commandStartedEvent": { + "command": { + "getMore": { + "$$type": [ + "int", + "long" + ] + }, + "collection": "coll0", + "batchSize": 2, + "comment": { + "content": "test" + } + }, + "commandName": "getMore", + "databaseName": "aggregate-tests" + } + } + ] + } + ] + }, + { + "description": "aggregate with comment does not set comment on getMore - pre 4.4", + "runOnRequirements": [ + { + "minServerVersion": "3.6.0", + "maxServerVersion": "4.3.99" } ], "operations": [ diff --git a/test/crud/unified/bulkWrite-updateMany-let.json b/test/crud/unified/bulkWrite-updateMany-let.json index 3cc8da4c53..fbeba1a607 100644 --- a/test/crud/unified/bulkWrite-updateMany-let.json +++ b/test/crud/unified/bulkWrite-updateMany-let.json @@ -142,7 +142,7 @@ "description": "BulkWrite updateMany with let option unsupported (server-side error)", "runOnRequirements": [ { - "minServerVersion": "3.6.0", + "minServerVersion": "4.2.0", "maxServerVersion": "4.9" } ], diff --git a/test/crud/unified/bulkWrite-updateOne-let.json b/test/crud/unified/bulkWrite-updateOne-let.json index 2a3e4f79dc..96783c782f 100644 --- a/test/crud/unified/bulkWrite-updateOne-let.json +++ b/test/crud/unified/bulkWrite-updateOne-let.json @@ -144,7 +144,7 @@ "description": "BulkWrite updateOne with let option unsupported (server-side error)", "runOnRequirements": [ { - "minServerVersion": "3.6.0", + "minServerVersion": "4.2.0", "maxServerVersion": "4.9" } ], diff --git a/test/crud/unified/find-comment.json b/test/crud/unified/find-comment.json index 6000bb0172..600a3723f1 100644 --- a/test/crud/unified/find-comment.json +++ b/test/crud/unified/find-comment.json @@ -195,10 +195,115 @@ ] }, { - "description": "find with comment does not set comment on getMore", + "description": "find with comment sets comment on getMore", "runOnRequirements": [ { - "minServerVersion": "3.6" + "minServerVersion": "4.4.0" + } + ], + "operations": [ + { + "name": "find", + "object": "collection0", + "arguments": { + "filter": { + "_id": { + "$gt": 1 + } + }, + "batchSize": 2, + "comment": { + "key": "value" + } + }, + "expectResult": [ + { + "_id": 2, + "x": 22 + }, + { + "_id": 3, + "x": 33 + }, + { + "_id": 4, + "x": 44 + }, + { + "_id": 5, + "x": 55 + }, + { + "_id": 6, + "x": 66 + } + ] + } + ], + "expectEvents": [ + { + "client": "client0", + "events": [ + { + "commandStartedEvent": { + "command": { + "find": "coll0", + "filter": { + "_id": { + "$gt": 1 + } + }, + "batchSize": 2, + "comment": { + "key": "value" + } + } + } + }, + { + "commandStartedEvent": { + "command": { + "getMore": { + "$$type": [ + "int", + "long" + ] + }, + "collection": "coll0", + "batchSize": 2, + "comment": { + "key": "value" + } + } + } + }, + { + "commandStartedEvent": { + "command": { + "getMore": { + "$$type": [ + "int", + "long" + ] + }, + "collection": "coll0", + "batchSize": 2, + "comment": { + "key": "value" + } + } + } + } + ] + } + ] + }, + { + "description": "find with comment does not set comment on getMore - pre 4.4", + "runOnRequirements": [ + { + "minServerVersion": "3.6.0", + "maxServerVersion": "4.3.99" } ], "operations": [ diff --git a/test/test_client.py b/test/test_client.py index a0d6e22d53..7a66792873 100644 --- a/test/test_client.py +++ b/test/test_client.py @@ -1344,6 +1344,7 @@ def test_stale_getmore(self): None, None, False, + None, ), unpack_res=Cursor(client.pymongo_test.collection)._unpack_response, address=("not-a-member", 27017),