SWIFT-1322 Support batchSize option for listCollections (#673)

Author: kmahar

Sources/MongoSwift/MongoCursor.swift

@@ -49,8 +49,7 @@ public class MongoCursor<T: Codable>: CursorProtocol {
     public let id: Int64?
 
     /**
-     * Initializes a new `MongoCursor` instance. Not meant to be instantiated directly by a user. When `forceIO` is
-     * true, this initializer will force a connection to the server if one is not already established.
+     * Initializes a new `MongoCursor` instance. Not meant to be instantiated directly by a user.
      *
      * - Throws:
      *   - `MongoError.InvalidArgumentError` if the options passed to the command that generated this cursor formed an

Sources/MongoSwift/MongoError.swift

@@ -522,7 +522,7 @@ internal func toErrorString(_ error: bson_error_t) -> String {
     }
 }
 
-internal let failedToRetrieveCursorMessage = "Couldn't get cursor from the server"
+internal let failedToRetrieveCursorMessage = "Expected libmongoc to return a cursor, unexpectedly got nil"
 
 extension MongoErrorProtocol {
     /// Determines whether this error is an "ns not found" error.

Sources/MongoSwift/Operations/ListCollectionsOperation.swift

@@ -106,23 +106,61 @@ internal struct ListCollectionsOperation: Operation {
     }
 
     internal func execute(using connection: Connection, session: ClientSession?) throws -> ListCollectionsResults {
-        var opts = try encodeOptions(options: self.options, session: session) ?? BSONDocument()
-        opts["nameOnly"] = .bool(self.nameOnly)
+        // Drivers MUST run listCollections on the primary node when in a replica set topology, unless directly
+        // connected to a secondary node in Single topology.
+        let readPref = ReadPreference.primary
+
+        var cmd: BSONDocument = ["listCollections": 1, "nameOnly": .bool(self.nameOnly)]
         if let filterDoc = self.filter {
-            opts["filter"] = .document(filterDoc)
+            cmd["filter"] = .document(filterDoc)
 
             // If `listCollectionNames` is called with a non-name filter key, change server-side nameOnly flag to false.
+            // per spec: drivers MUST NOT set nameOnly if a filter specifies any keys other than name.
             if self.nameOnly && filterDoc.keys.contains(where: { $0 != "name" }) {
-                opts["nameOnly"] = false
+                cmd["nameOnly"] = false
+            }
+        }
+
+        var cursorOpts: BSONDocument = [:]
+
+        if let batchSize = options?.batchSize {
+            guard let i32 = Int32(exactly: batchSize) else {
+                throw MongoError.InvalidArgumentError(
+                    message: "batchSize option must be representable as an Int32. Got: \(batchSize)"
+                )
+            }
+            cursorOpts = ["batchSize": .int32(i32)]
+        }
+
+        let commandOpts = try encodeOptions(options: nil as BSONDocument?, session: session) ?? BSONDocument()
+        cursorOpts = try encodeOptions(options: cursorOpts, session: session) ?? BSONDocument()
+        cmd["cursor"] = .document(cursorOpts)
+
+        // We don't need to clean up this reply ourselves, as `mongoc_cursor_new_from_command_reply_with_opts` will
+        // consume it.
+        var reply = try self.database.withMongocDatabase(from: connection) { dbPtr in
+            try readPref.withMongocReadPreference { rpPtr in
+                try runMongocCommandWithCReply(
+                    command: cmd,
+                    options: commandOpts
+                ) { cmdPtr, optsPtr, replyPtr, error in
+                    mongoc_database_read_command_with_opts(dbPtr, cmdPtr, rpPtr, optsPtr, replyPtr, &error)
+                }
             }
         }
 
-        let collections: OpaquePointer = self.database.withMongocDatabase(from: connection) { dbPtr in
-            opts.withBSONPointer { optsPtr in
-                guard let collections = mongoc_database_find_collections_with_opts(dbPtr, optsPtr) else {
-                    fatalError(failedToRetrieveCursorMessage)
+        let collections = connection.withMongocConnection { connPtr in
+            withUnsafeMutablePointer(to: &reply) { replyPtr -> OpaquePointer in
+                withOptionalBSONPointer(to: cursorOpts) { cursorOptsPtr in
+                    guard let result = mongoc_cursor_new_from_command_reply_with_opts(
+                        connPtr,
+                        replyPtr,
+                        cursorOptsPtr
+                    ) else {
+                        fatalError(failedToRetrieveCursorMessage)
+                    }
+                    return result
                 }
-                return collections
             }
         }
 

Sources/MongoSwift/Operations/MongocCommandHelpers.swift

@@ -33,6 +33,33 @@ internal func runMongocCommandWithReply(
     }
 }
 
+/**
+ * Calls the provided mongoc command method using pointers to the specified command and options. Returns the resulting
+ * filled-out `bson_t` from libmongoc.
+ * The caller of this method is responsible for ensuring the reply `bson_t` is properly cleaned up.
+ * If you need a `BSONDocument` reply, use `runMongocCommandWithReply` instead.
+ * If you don't need the reply, use `runMongocCommand` instead.
+ */
+internal func runMongocCommandWithCReply(
+    command: BSONDocument,
+    options: BSONDocument?,
+    body: MongocCommandFunc
+) throws -> bson_t {
+    var reply = bson_t()
+    do {
+        try withUnsafeMutablePointer(to: &reply) { replyPtr in
+            try _runMongocCommand(command: command, options: options, replyPtr: replyPtr, body: body)
+        }
+    } catch {
+        // on error, we need to clean up the bson_t ourselves here. on success, it is the caller's responsibility.
+        withUnsafeMutablePointer(to: &reply) { ptr in
+            bson_destroy(ptr)
+        }
+        throw error
+    }
+    return reply
+}
+
 /// Calls the provided mongoc command method using pointers to the specified command and options.
 internal func runMongocCommand(command: BSONDocument, options: BSONDocument?, body: MongocCommandFunc) throws {
     try withStackAllocatedMutableBSONPointer { replyPtr in

Tests/LinuxMain.swift

@@ -230,6 +230,7 @@ extension MongoDatabaseTests {
         ("testDropDatabase", testDropDatabase),
         ("testCreateCollection", testCreateCollection),
         ("testListCollections", testListCollections),
+        ("testListCollectionsBatchSize", testListCollectionsBatchSize),
         ("testAggregate", testAggregate),
         ("testAggregateWithOutputType", testAggregateWithOutputType),
         ("testAggregateWithListLocalSessions", testAggregateWithListLocalSessions),

Tests/MongoSwiftSyncTests/MongoDatabaseTests.swift

@@ -2,6 +2,7 @@ import Foundation
 import MongoSwiftSync
 import Nimble
 import TestsCommon
+import XCTest
 
 final class MongoDatabaseTests: MongoSwiftTestCase {
     override func setUp() {
@@ -180,6 +181,35 @@ final class MongoDatabaseTests: MongoSwiftTestCase {
         expect(listNamesEvents[3].command["nameOnly"]).to(equal(false))
     }
 
+    func testListCollectionsBatchSize() throws {
+        try self.withTestNamespace { client, db, _ in
+            // clear out collections
+            try db.drop()
+
+            _ = try db.createCollection("foo")
+            _ = try db.createCollection("bar")
+            _ = try db.createCollection("baz")
+
+            let monitor = client.addCommandMonitor()
+            try monitor.captureEvents {
+                let options = ListCollectionsOptions(batchSize: 2)
+                _ = Array(try db.listCollections(options: options))
+            }
+
+            let events = monitor.commandStartedEvents(withNames: ["listCollections", "getMore"])
+
+            guard events.count == 2 else {
+                XCTFail("Expected to find 2 events, but got \(events.count): \(events)")
+                return
+            }
+
+            expect(events[0].commandName).to(equal("listCollections"))
+            expect(events[0].command["cursor"]?.documentValue?["batchSize"]?.toInt()).to(equal(2))
+            expect(events[1].commandName).to(equal("getMore"))
+            expect(events[1].command["batchSize"]?.toInt()).to(equal(2))
+        }
+    }
+
     func testAggregate() throws {
         let client = try MongoClient.makeTestClient()
         // $currentOp must be run on the admin database