PYTHON-4921 Eliminate unnecessary killCursors command when batchSize == limit

pymongo/message.py

@@ -252,6 +252,10 @@ def _gen_find_command(
         if limit < 0:
             cmd["singleBatch"] = True
     if batch_size:
+        # When limit and batchSize are equal we increase batchSize by 1 to
+        # avoid an unnecessary killCursors.
+        if limit == batch_size:
+            batch_size += 1
         cmd["batchSize"] = batch_size
     if read_concern.level and not (session and session.in_transaction):
         cmd["readConcern"] = read_concern.document

test/crud/unified/client-bulkWrite-replaceOne-sort.json

@@ -3,7 +3,8 @@
   "schemaVersion": "1.4",
   "runOnRequirements": [
     {
-      "minServerVersion": "8.0"
+      "minServerVersion": "8.0",
+      "serverless": "forbid"
     }
   ],
   "createEntities": [

test/crud/unified/client-bulkWrite-updateOne-sort.json

@@ -3,7 +3,8 @@
   "schemaVersion": "1.4",
   "runOnRequirements": [
     {
-      "minServerVersion": "8.0"
+      "minServerVersion": "8.0",
+      "serverless": "forbid"
     }
   ],
   "createEntities": [

test/crud/unified/distinct-hint.json

@@ -0,0 +1,139 @@
+{
+  "description": "distinct-hint",
+  "schemaVersion": "1.0",
+  "runOnRequirements": [
+    {
+      "minServerVersion": "7.1.0"
+    }
+  ],
+  "createEntities": [
+    {
+      "client": {
+        "id": "client0",
+        "observeEvents": [
+          "commandStartedEvent"
+        ]
+      }
+    },
+    {
+      "database": {
+        "id": "database0",
+        "client": "client0",
+        "databaseName": "distinct-hint-tests"
+      }
+    },
+    {
+      "collection": {
+        "id": "collection0",
+        "database": "database0",
+        "collectionName": "coll0"
+      }
+    }
+  ],
+  "initialData": [
+    {
+      "collectionName": "coll0",
+      "databaseName": "distinct-hint-tests",
+      "documents": [
+        {
+          "_id": 1,
+          "x": 11
+        },
+        {
+          "_id": 2,
+          "x": 22
+        },
+        {
+          "_id": 3,
+          "x": 33
+        }
+      ]
+    }
+  ],
+  "tests": [
+    {
+      "description": "distinct with hint string",
+      "operations": [
+        {
+          "name": "distinct",
+          "object": "collection0",
+          "arguments": {
+            "fieldName": "x",
+            "filter": {
+              "_id": 1
+            },
+            "hint": "_id_"
+          },
+          "expectResult": [
+            11
+          ]
+        }
+      ],
+      "expectEvents": [
+        {
+          "client": "client0",
+          "events": [
+            {
+              "commandStartedEvent": {
+                "command": {
+                  "distinct": "coll0",
+                  "key": "x",
+                  "query": {
+                    "_id": 1
+                  },
+                  "hint": "_id_"
+                },
+                "commandName": "distinct",
+                "databaseName": "distinct-hint-tests"
+              }
+            }
+          ]
+        }
+      ]
+    },
+    {
+      "description": "distinct with hint document",
+      "operations": [
+        {
+          "name": "distinct",
+          "object": "collection0",
+          "arguments": {
+            "fieldName": "x",
+            "filter": {
+              "_id": 1
+            },
+            "hint": {
+              "_id": 1
+            }
+          },
+          "expectResult": [
+            11
+          ]
+        }
+      ],
+      "expectEvents": [
+        {
+          "client": "client0",
+          "events": [
+            {
+              "commandStartedEvent": {
+                "command": {
+                  "distinct": "coll0",
+                  "key": "x",
+                  "query": {
+                    "_id": 1
+                  },
+                  "hint": {
+                    "_id": 1
+                  }
+                },
+                "commandName": "distinct",
+                "databaseName": "distinct-hint-tests"
+              }
+            }
+          ]
+        }
+      ]
+    }
+  ]
+}

test/crud/unified/estimatedDocumentCount.json

@@ -249,7 +249,7 @@
           "name": "estimatedDocumentCount",
           "object": "collection0",
           "expectError": {
-            "isError": true
+            "isClientError": true
           }
         }
       ],

test/crud/unified/find.json

@@ -237,6 +237,68 @@
           ]
         }
       ]
+    },
+    {
+      "description": "Find with batchSize equal to limit",
+      "operations": [
+        {
+          "object": "collection0",
+          "name": "find",
+          "arguments": {
+            "filter": {
+              "_id": {
+                "$gt": 1
+              }
+            },
+            "sort": {
+              "_id": 1
+            },
+            "limit": 4,
+            "batchSize": 4
+          },
+          "expectResult": [
+            {
+              "_id": 2,
+              "x": 22
+            },
+            {
+              "_id": 3,
+              "x": 33
+            },
+            {
+              "_id": 4,
+              "x": 44
+            },
+            {
+              "_id": 5,
+              "x": 55
+            }
+          ]
+        }
+      ],
+      "expectEvents": [
+        {
+          "client": "client0",
+          "events": [
+            {
+              "commandStartedEvent": {
+                "command": {
+                  "find": "coll0",
+                  "filter": {
+                    "_id": {
+                      "$gt": 1
+                    }
+                  },
+                  "limit": 4,
+                  "batchSize": 5
+                },
+                "commandName": "find",
+                "databaseName": "find-tests"
+              }
+            }
+          ]
+        }
+      ]
     }
   ]
 }