initial commit, for some reason some spec tests not passing

This commit is contained in:
julius 2022-03-17 03:22:50 -07:00
parent 225d131c2d
commit ae63fbcda6
11 changed files with 113 additions and 20 deletions

1
.gitignore vendored
View File

@ -15,3 +15,4 @@ pymongo.egg-info/
.tox
mongocryptd.pid
.idea/
.nova/

View File

@ -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

View File

@ -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

View File

@ -1182,6 +1182,7 @@ class Cursor(Generic[_DocumentType]):
self.__max_await_time_ms,
self.__sock_mgr,
self.__exhaust,
self.__comment,
)
self.__send_message(g)

View File

@ -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

View File

@ -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)

View File

@ -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,

View File

@ -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"
}
}
]
}
]
}
]
}

View File

@ -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"

View File

@ -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"
}
}
}

View File

@ -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),