fix(plugin-mongodb): stop large collections hanging on open and make Stop cancel the query

CHANGELOG.md

@@ -7,6 +7,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- **Count exactly** next to an estimated row total, to replace the estimate with a real count when you want it.
+
+### Fixed
+
+- Opening a large MongoDB collection no longer hangs. Row counts that back the pagination display now give up after 5 seconds and leave the estimate in place instead of holding the tab.
+- Stop now cancels a running MongoDB query on the server, rather than leaving it running until it finishes on its own.
+- A MongoDB query that runs out of time now says so, and points at the missing index that usually causes it.
+- Running a MongoDB query with no limit no longer pulls the whole collection into memory before applying the row limit.
+- Exporting a MongoDB query that has a skip or limit now exports that range, not the whole collection.
+- A row count estimate of zero from a table that was never analyzed no longer displays as an empty table.
+
 ## [0.61.0] - 2026-07-30
 
 ### Added

Plugins/MongoDBDriverPlugin/MongoDBConnection+SyncHelpers.swift

@@ -111,6 +111,14 @@ extension MongoDBConnection {
         }
         defer { bson_destroy(optsBson) }
 
+        let session = attachCancellableSession(client: client, opts: optsBson)
+        defer {
+            if let session {
+                releaseSessionLsid()
+                mongoc_client_session_destroy(session)
+            }
+        }
+
         let col = try getCollection(client, database: database, collection: collection)
         defer { mongoc_collection_destroy(col) }
 
@@ -124,6 +132,20 @@ extension MongoDBConnection {
         return try iterateCursor(cursor)
     }
 
+    /// Binds the operation to a server-side session so `cancelCurrentQuery` can abort it with
+    /// `killSessions`. Returns nil when the deployment does not support sessions, in which case the
+    /// operation still runs, just without server-side cancellation.
+    func attachCancellableSession(client: OpaquePointer, opts: OpaquePointer) -> OpaquePointer? {
+        var error = bson_error_t()
+        guard let session = mongoc_client_start_session(client, nil, &error) else { return nil }
+        guard mongoc_client_session_append(session, opts, &error) else {
+            mongoc_client_session_destroy(session)
+            return nil
+        }
+        adoptSessionLsid(session)
+        return session
+    }
+
     func aggregateSync(
         client: OpaquePointer, database: String, collection: String, pipeline: String
     ) throws -> (docs: [[String: Any]], isTruncated: Bool) {
@@ -157,7 +179,8 @@ extension MongoDBConnection {
     }
 
     func countDocumentsSync(
-        client: OpaquePointer, database: String, collection: String, filter: String
+        client: OpaquePointer, database: String, collection: String,
+        filter: String, background: Bool
     ) throws -> Int64 {
         try checkCancelled()
 
@@ -169,19 +192,61 @@ extension MongoDBConnection {
         let col = try getCollection(client, database: database, collection: collection)
         defer { mongoc_collection_destroy(col) }
 
-        let timeoutMS = queryTimeoutMS
+        let maxTimeMS = effectiveMaxTimeMS(background: background)
+        let isUnfiltered = filter.trimmingCharacters(in: .whitespacesAndNewlines) == "{}"
+
+        let hinted = try runCountDocuments(
+            collection: col, filter: filterBson, maxTimeMS: maxTimeMS, hintIdIndex: isUnfiltered
+        )
+        try checkCancelled()
+        switch hinted {
+        case .count(let value):
+            return value
+        case .failed(let error):
+            guard isUnfiltered, MongoDBTimeoutPolicy.shouldRetryWithoutHint(errorCode: error.code) else {
+                throw error
+            }
+        }
+
+        let unhinted = try runCountDocuments(
+            collection: col, filter: filterBson, maxTimeMS: maxTimeMS, hintIdIndex: false
+        )
+        try checkCancelled()
+        switch unhinted {
+        case .count(let value):
+            return value
+        case .failed(let error):
+            throw error
+        }
+    }
+
+    private enum CountOutcome {
+        case count(Int64)
+        case failed(MongoDBError)
+    }
+
+    private func runCountDocuments(
+        collection col: OpaquePointer, filter: OpaquePointer,
+        maxTimeMS: Int32?, hintIdIndex: Bool
+    ) throws -> CountOutcome {
+        var optsFields: [String] = []
+        if let maxTimeMS {
+            optsFields.append("\"maxTimeMS\": \(maxTimeMS)")
+        }
+        if hintIdIndex {
+            optsFields.append("\"hint\": \"_id_\"")
+        }
+
         var optsBson: OpaquePointer?
-        if timeoutMS > 0 {
-            optsBson = jsonToBson("{\"maxTimeMS\": \(timeoutMS)}")
+        if !optsFields.isEmpty {
+            optsBson = jsonToBson("{\(optsFields.joined(separator: ", "))}")
         }
-        defer { if let opts = optsBson { bson_destroy(opts) } }
+        defer { if let optsBson { bson_destroy(optsBson) } }
 
         var error = bson_error_t()
-        let count = mongoc_collection_count_documents(col, filterBson, optsBson, nil, nil, &error)
-
-        try checkCancelled()
-        guard count >= 0 else { throw makeError(error) }
-        return count
+        let count = mongoc_collection_count_documents(col, filter, optsBson, nil, nil, &error)
+        guard count >= 0 else { return .failed(makeError(error)) }
+        return .count(count)
     }
 
     func insertOneSync(
@@ -367,6 +432,7 @@ extension MongoDBConnection {
         var docPtr: OpaquePointer?
         var sample: [[String: Any]] = []
         var projection: MongoStreamProjection?
+        var emitted = 0
 
         while mongoc_cursor_next(cursor, &docPtr) {
             if Task.isCancelled {
@@ -380,12 +446,18 @@ extension MongoDBConnection {
 
             if let projection {
                 continuation.yield(.rows([projection.row(for: dict, convert: streamCellValue)]))
+                emitted += 1
+                if emitted >= PluginRowLimits.emergencyMax {
+                    logger.warning("Streamed result truncated at \(PluginRowLimits.emergencyMax) documents")
+                    break
+                }
                 continue
             }
 
             sample.append(dict)
             if sample.count >= MongoStreamProjection.sampleSize {
                 projection = openStream(sample: sample, continuation: continuation)
+                emitted += sample.count
                 sample = []
             }
         }

Plugins/MongoDBDriverPlugin/MongoDBConnection.swift

@@ -68,12 +68,17 @@ final class MongoDBConnection: @unchecked Sendable {
     private let replicaSet: String?
     private let extraUriParams: [String: String]
 
+    private let controlQueue = DispatchQueue(label: "com.TablePro.mongodb.control", qos: .userInitiated)
+
     private let stateLock = NSLock()
     private var _isConnected: Bool = false
     private var _isShuttingDown: Bool = false
     private var _cachedServerVersion: String?
     private var _isCancelled: Bool = false
     private var _queryTimeoutMS: Int32 = 0
+    #if canImport(CLibMongoc)
+    private var _activeSessionLsid: OpaquePointer?
+    #endif
 
     var isConnected: Bool {
         stateLock.lock()
@@ -106,6 +111,10 @@ final class MongoDBConnection: @unchecked Sendable {
         stateLock.unlock()
     }
 
+    func effectiveMaxTimeMS(background: Bool) -> Int32? {
+        MongoDBTimeoutPolicy.resolveMaxTimeMS(ambientMS: queryTimeoutMS, background: background)
+    }
+
     // MARK: - Initialization
 
     init(
@@ -339,6 +348,8 @@ final class MongoDBConnection: @unchecked Sendable {
         #if canImport(CLibMongoc)
         let handle = client
         client = nil
+        let staleLsid = _activeSessionLsid
+        _activeSessionLsid = nil
         #endif
         _isConnected = false
         _cachedServerVersion = nil
@@ -347,6 +358,7 @@ final class MongoDBConnection: @unchecked Sendable {
         stateLock.unlock()
 
         #if canImport(CLibMongoc)
+        if let staleLsid { bson_destroy(staleLsid) }
         if let handle = handle {
             queue.async { mongoc_client_destroy(handle) }
         }
@@ -358,8 +370,63 @@ final class MongoDBConnection: @unchecked Sendable {
     func cancelCurrentQuery() {
         stateLock.lock()
         _isCancelled = true
+        #if canImport(CLibMongoc)
+        let lsid = _activeSessionLsid.map { bson_copy($0) }
+        #endif
         stateLock.unlock()
+
+        #if canImport(CLibMongoc)
+        guard let lsid else { return }
+        killSession(lsid: lsid)
+        #endif
+    }
+
+    #if canImport(CLibMongoc)
+    func adoptSessionLsid(_ session: OpaquePointer) {
+        guard let lsid = mongoc_client_session_get_lsid(session) else { return }
+        let copy = bson_copy(lsid)
+        stateLock.lock()
+        let previous = _activeSessionLsid
+        _activeSessionLsid = copy
+        stateLock.unlock()
+        if let previous { bson_destroy(previous) }
+    }
+
+    func releaseSessionLsid() {
+        stateLock.lock()
+        let previous = _activeSessionLsid
+        _activeSessionLsid = nil
+        stateLock.unlock()
+        if let previous { bson_destroy(previous) }
+    }
+
+    /// Aborts the in-flight operation server-side. The connection's own queue is blocked inside a
+    /// libmongoc call at this point, so the `killSessions` command needs its own client.
+    private func killSession(lsid: OpaquePointer) {
+        let uriString = buildUri()
+        controlQueue.async {
+            defer { bson_destroy(lsid) }
+            guard let controlClient = mongoc_client_new(uriString) else { return }
+            defer { mongoc_client_destroy(controlClient) }
+
+            let command = bson_new()
+            defer { bson_destroy(command) }
+            let sessions = bson_new()
+            defer { bson_destroy(sessions) }
+
+            bson_append_document(sessions, "0", -1, lsid)
+            bson_append_array(command, "killSessions", -1, sessions)
+
+            let reply = bson_new()
+            defer { bson_destroy(reply) }
+            var error = bson_error_t()
+            let killed = mongoc_client_command_simple(controlClient, "admin", command, nil, reply, &error)
+            if !killed {
+                logger.warning("killSessions failed with code \(error.code, privacy: .public)")
+            }
+        }
     }
+    #endif
 
     /// Throws if cancellation was requested, resetting the flag atomically.
     /// Safe to call from any thread.
@@ -493,7 +560,12 @@ final class MongoDBConnection: @unchecked Sendable {
         #endif
     }
 
-    func countDocuments(database: String, collection: String, filter: String) async throws -> Int64 {
+    func countDocuments(
+        database: String,
+        collection: String,
+        filter: String,
+        background: Bool
+    ) async throws -> Int64 {
         #if canImport(CLibMongoc)
         resetCancellation()
         return try await pluginDispatchAsync(on: queue) { [self] in
@@ -502,7 +574,8 @@ final class MongoDBConnection: @unchecked Sendable {
             }
             try checkCancelled()
             let count = try countDocumentsSync(
-                client: client, database: database, collection: collection, filter: filter
+                client: client, database: database, collection: collection,
+                filter: filter, background: background
             )
             try checkCancelled()
             return count
@@ -512,7 +585,11 @@ final class MongoDBConnection: @unchecked Sendable {
         #endif
     }
 
-    func estimatedDocumentCount(database: String, collection: String) async throws -> Int64 {
+    func estimatedDocumentCount(
+        database: String,
+        collection: String,
+        background: Bool
+    ) async throws -> Int64 {
         #if canImport(CLibMongoc)
         resetCancellation()
         return try await pluginDispatchAsync(on: queue) { [self] in
@@ -523,8 +600,14 @@ final class MongoDBConnection: @unchecked Sendable {
             let col = try getCollection(client, database: database, collection: collection)
             defer { mongoc_collection_destroy(col) }
 
+            var opts: OpaquePointer?
+            if let maxTimeMS = effectiveMaxTimeMS(background: background) {
+                opts = jsonToBson("{\"maxTimeMS\": \(maxTimeMS)}")
+            }
+            defer { if let opts { bson_destroy(opts) } }
+
             var error = bson_error_t()
-            let count = mongoc_collection_estimated_document_count(col, nil, nil, nil, &error)
+            let count = mongoc_collection_estimated_document_count(col, opts, nil, nil, &error)
             if count < 0 {
                 throw makeError(error)
             }
@@ -639,7 +722,9 @@ final class MongoDBConnection: @unchecked Sendable {
         collection: String,
         filter: String,
         sort: String?,
-        projection: String?
+        projection: String?,
+        skip: Int = 0,
+        limit: Int? = nil
     ) -> AsyncThrowingStream<PluginStreamElement, Error> {
         #if canImport(CLibMongoc)
         let queue = self.queue
@@ -683,6 +768,12 @@ final class MongoDBConnection: @unchecked Sendable {
                        let obj = try? JSONSerialization.jsonObject(with: data) {
                         optsJson["projection"] = obj
                     }
+                    if skip > 0 {
+                        optsJson["skip"] = skip
+                    }
+                    if let limit, limit > 0 {
+                        optsJson["limit"] = limit
+                    }
                     let timeoutMS = queryTimeoutMS
                     if timeoutMS > 0 {
                         optsJson["maxTimeMS"] = timeoutMS

Plugins/MongoDBDriverPlugin/MongoDBFindLimitPolicy.swift

@@ -0,0 +1,28 @@
+//
+//  MongoDBFindLimitPolicy.swift
+//  MongoDBDriverPlugin
+//
+
+import Foundation
+import TableProPluginKit
+
+enum MongoDBFindLimitPolicy {
+    /// Server-side limit for a user query. One document beyond the row cap is requested so the
+    /// caller can tell "exactly a full page" from "there is more", without ever asking the server
+    /// for the whole collection.
+    static func fetchLimit(parsedLimit: Int?, rowCap: Int?) -> Int {
+        var candidates: [Int] = []
+        if let parsedLimit, parsedLimit > 0 {
+            candidates.append(parsedLimit)
+        }
+        if let rowCap, rowCap > 0 {
+            candidates.append(rowCap + 1)
+        }
+        return min(candidates.min() ?? PluginRowLimits.emergencyMax, PluginRowLimits.emergencyMax)
+    }
+
+    static func isTruncated(rowCount: Int, rowCap: Int?) -> Bool {
+        guard let rowCap, rowCap > 0 else { return false }
+        return rowCount > rowCap
+    }
+}