From ae63fbcda6778b4cf36b6a2f853b7a5d81b5c326 Mon Sep 17 00:00:00 2001 From: julius Date: Thu, 17 Mar 2022 03:22:50 -0700 Subject: [PATCH] initial commit, for some reason some spec tests not passing --- .gitignore | 1 + pymongo/collection.py | 7 +- pymongo/command_cursor.py | 5 ++ pymongo/cursor.py | 1 + pymongo/database.py | 1 + pymongo/message.py | 11 ++- pymongo/mongo_client.py | 2 +- .../unified/change-streams.json | 82 +++++++++++++++++++ test/crud/unified/aggregate.json | 10 +-- test/crud/unified/find-comment.json | 10 +-- test/test_client.py | 3 +- 11 files changed, 113 insertions(+), 20 deletions(-) diff --git a/.gitignore b/.gitignore index de435d109..f7ad6563f 100644 --- a/.gitignore +++ b/.gitignore @@ -15,3 +15,4 @@ pymongo.egg-info/ .tox mongocryptd.pid .idea/ +.nova/ diff --git a/pymongo/collection.py b/pymongo/collection.py index 46916f98f..517242308 100644 --- a/pymongo/collection.py +++ b/pymongo/collection.py @@ -2179,7 +2179,12 @@ class Collection(common.BaseObject, Generic[_DocumentType]): 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 d10e23f95..0809d34e6 100644 --- a/pymongo/command_cursor.py +++ b/pymongo/command_cursor.py @@ -43,6 +43,7 @@ class CommandCursor(Generic[_DocumentType]): 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 @@ class CommandCursor(Generic[_DocumentType]): self.__session = session self.__explicit_session = explicit_session self.__killed = self.__id == 0 + self.__comment = comment if self.__killed: self.__end_session(True) @@ -156,6 +158,8 @@ class CommandCursor(Generic[_DocumentType]): def __send_message(self, operation): """Send a getmore message and handle the response.""" client = self.__collection.database.client + if self.__comment is not None: + operation["comment"] = self.__comment try: response = client._run_operation( operation, self._unpack_response, address=self.__address @@ -224,6 +228,7 @@ class CommandCursor(Generic[_DocumentType]): self.__max_await_time_ms, self.__sock_mgr, False, + self.__comment, ) ) else: # Cursor id is zero nothing else to return diff --git a/pymongo/cursor.py b/pymongo/cursor.py index a2ccdf586..0ff7f615f 100644 --- a/pymongo/cursor.py +++ b/pymongo/cursor.py @@ -1182,6 +1182,7 @@ class Cursor(Generic[_DocumentType]): 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 c2f2eb4bc..eaebc25c5 100644 --- a/pymongo/database.py +++ b/pymongo/database.py @@ -774,6 +774,7 @@ class Database(common.BaseObject, Generic[_DocumentType]): 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 92d59c3eb..0a016020f 100644 --- a/pymongo/message.py +++ b/pymongo/message.py @@ -220,13 +220,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): """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: + cmd["comment"] = comment return cmd @@ -419,6 +421,7 @@ class _GetMore(object): "sock_mgr", "_as_command", "exhaust", + "comment", ) name = "getMore" @@ -436,6 +439,7 @@ class _GetMore(object): max_await_time_ms, sock_mgr, exhaust, + comment, ): self.db = db self.coll = coll @@ -449,6 +453,7 @@ class _GetMore(object): 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) @@ -467,13 +472,13 @@ class _GetMore(object): def as_command(self, sock_info): """Return a getMore command document for this query.""" # See _Query.as_command for an explanation of this caching. + if self._as_command is not None: 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 ) - 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 cd9067f46..8bdc8146a 100644 --- a/pymongo/mongo_client.py +++ b/pymongo/mongo_client.py @@ -1761,7 +1761,7 @@ class MongoClient(common.BaseObject, Generic[_DocumentType]): "firstBatch": res["databases"], "ns": "admin.$cmd", } - return CommandCursor(admin["$cmd"], cursor, None) + return CommandCursor(admin["$cmd"], cursor, None, 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 4aea9a4aa..9187d4eda 100644 --- a/test/change_streams/unified/change-streams.json +++ b/test/change_streams/unified/change-streams.json @@ -247,6 +247,88 @@ ] } ] + }, + { + "description": "Test that comment is set on getMore", + "runOnRequirements": [ + { + "minServerVersion": "3.6.0" + } + ], + "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": "comment" + }, + "commandName": "getMore", + "databaseName": "database0" + } + } + ] + } + ] } ] } diff --git a/test/crud/unified/aggregate.json b/test/crud/unified/aggregate.json index f6da8ff32..73d37917d 100644 --- a/test/crud/unified/aggregate.json +++ b/test/crud/unified/aggregate.json @@ -327,7 +327,7 @@ ] }, { - "description": "aggregate with comment does not set comment on getMore", + "description": "aggregate with comment sets comment on getMore", "runOnRequirements": [ { "minServerVersion": "3.6.0" @@ -411,9 +411,7 @@ }, "collection": "coll0", "batchSize": 2, - "comment": { - "$$exists": false - } + "comment": "comment" }, "commandName": "getMore", "databaseName": "aggregate-tests" @@ -430,9 +428,7 @@ }, "collection": "coll0", "batchSize": 2, - "comment": { - "$$exists": false - } + "comment": "comment" }, "commandName": "getMore", "databaseName": "aggregate-tests" diff --git a/test/crud/unified/find-comment.json b/test/crud/unified/find-comment.json index 6000bb017..d7cf85968 100644 --- a/test/crud/unified/find-comment.json +++ b/test/crud/unified/find-comment.json @@ -195,7 +195,7 @@ ] }, { - "description": "find with comment does not set comment on getMore", + "description": "find with comment sets comment on getMore", "runOnRequirements": [ { "minServerVersion": "3.6" @@ -267,9 +267,7 @@ }, "collection": "coll0", "batchSize": 2, - "comment": { - "$$exists": false - } + "comment": "comment" } } }, @@ -284,9 +282,7 @@ }, "collection": "coll0", "batchSize": 2, - "comment": { - "$$exists": false - } + "comment": "comment" } } } diff --git a/test/test_client.py b/test/test_client.py index 29a5b0f1d..251da4725 100644 --- a/test/test_client.py +++ b/test/test_client.py @@ -1321,7 +1321,7 @@ class TestClient(IntegrationTest): with self.assertRaises(AutoReconnect): client = rs_client(connect=False, serverSelectionTimeoutMS=100) client._run_operation( - operation=message._GetMore( + operation=message.kge_GetMore( "pymongo_test", "collection", 101, @@ -1333,6 +1333,7 @@ class TestClient(IntegrationTest): None, None, False, + None, ), unpack_res=Cursor(client.pymongo_test.collection)._unpack_response, address=("not-a-member", 27017),