fix(plugin-mongodb): parse db.getCollection and align exported columns (#1971)

CHANGELOG.md

@@ -18,6 +18,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Exporting a MongoDB collection works again. Every format failed with "Unsupported MongoDB method: getCollection" before any data was read. You can now also write `db.getCollection("name")` in the query editor, which is the only way to reach collections whose names contain dots or spaces, start with a digit, or match a database method such as `stats`. (#1971)
+- Exported MongoDB collections keep every field. TablePro took the column list from the first document alone, so a field that showed up later was dropped from JSON exports and pushed CSV rows out of line with their header. (#1971)
 - The `postgres` database shows in the database list again. It was marked as a system database, which hid it from the sidebar, Cmd+K, the database filter, and the Backup and Restore Dump pickers. PostgreSQL creates it for users and applications, so nothing about it is internal. CockroachDB's `defaultdb` and Redshift's `dev` were hidden the same way and now show too. (#1967)
 - The database a connection is using always shows in the sidebar and the database switchers, even when it is a system database or the database filter excludes it. (#1967)
 - The license activation sheet now opens when you click **Activate License**. It was built and then failed to appear, and once that happened further clicks did nothing at all.

Plugins/MongoDBDriverPlugin/MongoDBConnection+SyncHelpers.swift

@@ -365,9 +365,8 @@ extension MongoDBConnection {
         streamState: MongoStreamState
     ) {
         var docPtr: OpaquePointer?
-        var headerSent = false
-        var columns: [String] = []
-        var columnTypeNames: [String] = []
+        var sample: [[String: Any]] = []
+        var projection: MongoStreamProjection?
 
         while mongoc_cursor_next(cursor, &docPtr) {
             if Task.isCancelled {
@@ -379,31 +378,16 @@ extension MongoDBConnection {
             guard let doc = docPtr else { continue }
             let dict = bsonToDict(doc)
 
-            if !headerSent {
-                columns = BsonDocumentFlattener.unionColumns(from: [dict])
-                let bsonTypes = BsonDocumentFlattener.columnTypes(for: columns, documents: [dict])
-                columnTypeNames = bsonTypes.map { bsonTypeToStreamString($0) }
-                continuation.yield(.header(PluginStreamHeader(
-                    columns: columns,
-                    columnTypeNames: columnTypeNames
-                )))
-                headerSent = true
-            } else {
-                for key in dict.keys.sorted() where !columns.contains(key) {
-                    columns.append(key)
-                    let type = BsonDocumentFlattener.columnTypes(for: [key], documents: [dict])
-                    columnTypeNames.append(bsonTypeToStreamString(type.first ?? 2))
-                }
+            if let projection {
+                continuation.yield(.rows([projection.row(for: dict, convert: streamCellValue)]))
+                continue
             }
 
-            let row: [PluginCellValue] = columns.map { column in
-                guard let value = dict[column] else { return .null }
-                if let data = value as? Data {
-                    return .bytes(data)
-                }
-                return PluginCellValue.fromOptional(BsonDocumentFlattener.stringValue(for: value))
+            sample.append(dict)
+            if sample.count >= MongoStreamProjection.sampleSize {
+                projection = openStream(sample: sample, continuation: continuation)
+                sample = []
             }
-            continuation.yield(.rows([row]))
         }
 
         var error = bson_error_t()
@@ -413,17 +397,40 @@ extension MongoDBConnection {
             return
         }
 
-        if !headerSent {
-            continuation.yield(.header(PluginStreamHeader(
-                columns: ["_id"],
-                columnTypeNames: ["VARCHAR"]
-            )))
+        if projection == nil {
+            _ = openStream(sample: sample, continuation: continuation)
         }
 
         cleanup(streamState)
         continuation.finish()
     }
 
+    private func openStream(
+        sample: [[String: Any]],
+        continuation: AsyncThrowingStream<PluginStreamElement, Error>.Continuation
+    ) -> MongoStreamProjection {
+        let columns = BsonDocumentFlattener.unionColumns(from: sample)
+        let columnTypeNames = BsonDocumentFlattener
+            .columnTypes(for: columns, documents: sample)
+            .map { bsonTypeToStreamString($0) }
+        let projection = MongoStreamProjection(columns: columns, columnTypeNames: columnTypeNames)
+
+        continuation.yield(.header(projection.header))
+
+        if !sample.isEmpty {
+            continuation.yield(.rows(sample.map { projection.row(for: $0, convert: streamCellValue) }))
+        }
+
+        return projection
+    }
+
+    private func streamCellValue(_ value: Any) -> PluginCellValue {
+        if let data = value as? Data {
+            return .bytes(data)
+        }
+        return PluginCellValue.fromOptional(BsonDocumentFlattener.stringValue(for: value))
+    }
+
     private func cleanup(_ state: MongoStreamState) {
         state.lock.lock()
         let cur = state.cursor

Plugins/MongoDBDriverPlugin/MongoDBPluginDriver.swift

@@ -27,7 +27,7 @@ final class MongoDBPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
     }
 
     func defaultExportQuery(table: String) -> String? {
-        "db.getCollection(\"\(table)\").find({})"
+        MongoDBQueryBuilder().buildExportQuery(collection: table)
     }
 
     init(config: DriverConnectionConfig) {

Plugins/MongoDBDriverPlugin/MongoDBQueryBuilder.swift

@@ -66,6 +66,12 @@ struct MongoDBQueryBuilder {
         "\(Self.mongoCollectionAccessor(collection)).countDocuments(\(filterJson))"
     }
 
+    // MARK: - Export Query
+
+    func buildExportQuery(collection: String) -> String {
+        "\(Self.mongoCollectionAccessor(collection)).find({})"
+    }
+
     // MARK: - Filter Document
 
     /// Convert filter tuples to MongoDB filter document string

Plugins/MongoDBDriverPlugin/MongoStreamProjection.swift

@@ -0,0 +1,40 @@
+//
+//  MongoStreamProjection.swift
+//  MongoDBDriverPlugin
+//
+//  Keeps streamed documents aligned to the column set announced in the stream header.
+//
+
+import Foundation
+import TableProPluginKit
+
+struct MongoStreamProjection {
+    static let sampleSize = 200
+
+    let columns: [String]
+    let columnTypeNames: [String]
+
+    init(columns: [String], columnTypeNames: [String]) {
+        guard !columns.isEmpty else {
+            self.columns = ["_id"]
+            self.columnTypeNames = ["VARCHAR"]
+            return
+        }
+
+        self.columns = columns
+        self.columnTypeNames = columns.indices.map { index in
+            index < columnTypeNames.count ? columnTypeNames[index] : "VARCHAR"
+        }
+    }
+
+    var header: PluginStreamHeader {
+        PluginStreamHeader(columns: columns, columnTypeNames: columnTypeNames)
+    }
+
+    func row(for document: [String: Any], convert: (Any) -> PluginCellValue) -> [PluginCellValue] {
+        columns.map { column in
+            guard let value = document[column] else { return .null }
+            return convert(value)
+        }
+    }
+}

Plugins/TableProPluginKit/MongoShellParser.swift

@@ -123,65 +123,19 @@ public struct MongoShellParser {
     // MARK: - Private Parsing
 
     /// Parse db["collection"].method(args) bracket notation.
-    /// Supports both double and single quotes around the collection name.
+    /// Supports double quotes, single quotes and backticks around the collection name.
     private static func parseBracketExpression(_ input: String) throws -> MongoOperation {
-        // input starts with db[
-        let afterBracket = String(input.dropFirst(3)) // drop "db["
+        let nameStart = input.index(input.startIndex, offsetBy: 3)
+        let literal = try readStringLiteral(in: input, startingAt: nameStart)
 
-        // Determine quote character (" or ')
-        guard let quoteChar = afterBracket.first, quoteChar == "\"" || quoteChar == "'" else {
-            throw MongoShellParseError.invalidSyntax("Expected quoted collection name in db[...]")
+        guard literal.endIndex < input.endIndex, input[literal.endIndex] == "]" else {
+            throw MongoShellParseError.invalidSyntax(
+                String(localized: "Expected ']' after the collection name")
+            )
         }
 
-        // Find closing quote (handle escaped quotes)
-        var collectionName = ""
-        var i = afterBracket.index(after: afterBracket.startIndex)
-        var escapeNext = false
-        while i < afterBracket.endIndex {
-            let ch = afterBracket[i]
-            if escapeNext {
-                collectionName.append(ch)
-                escapeNext = false
-                i = afterBracket.index(after: i)
-                continue
-            }
-            if ch == "\\" {
-                escapeNext = true
-                i = afterBracket.index(after: i)
-                continue
-            }
-            if ch == quoteChar {
-                break
-            }
-            collectionName.append(ch)
-            i = afterBracket.index(after: i)
-        }
-
-        guard i < afterBracket.endIndex else {
-            throw MongoShellParseError.invalidSyntax("Unterminated string in db[...]")
-        }
-
-        // Move past closing quote and expect "]"
-        i = afterBracket.index(after: i)
-        guard i < afterBracket.endIndex, afterBracket[i] == "]" else {
-            throw MongoShellParseError.invalidSyntax("Expected ']' after collection name in db[...]")
-        }
-        i = afterBracket.index(after: i)
-
-        let remaining = String(afterBracket[i...]).trimmingCharacters(in: .whitespacesAndNewlines)
-
-        // No method chain — treat as find all
-        if remaining.isEmpty {
-            return .find(collection: collectionName, filter: "{}", options: MongoFindOptions())
-        }
-
-        // Expect ".method(args)" after db["collection"]
-        guard remaining.hasPrefix(".") else {
-            throw MongoShellParseError.invalidSyntax("Expected '.method()' after db[\"...\"]")
-        }
-
-        let methodChain = String(remaining.dropFirst())
-        return try parseMethodChain(collection: collectionName, chain: methodChain)
+        let chain = String(input[input.index(after: literal.endIndex)...])
+        return try parseAccessorChain(collection: literal.value, chain: chain)
     }
 
     private static func parseDbExpression(_ input: String) throws -> MongoOperation {
@@ -202,6 +156,9 @@ public struct MongoShellParser {
         // This correctly handles dotted collection names like "system.version".
         let beforeParen = afterDb[afterDb.startIndex..<firstParen]
         guard let lastDot = beforeParen.lastIndex(of: ".") else {
+            if beforeParen.trimmingCharacters(in: .whitespacesAndNewlines) == "getCollection" {
+                return try parseGetCollectionExpression(afterDb, openParen: firstParen)
+            }
             // No dot before paren — db-level method call like db.getCollectionNames()
             return try parseDbLevelMethod(afterDb)
         }
@@ -212,6 +169,51 @@ public struct MongoShellParser {
         return try parseMethodChain(collection: collection, chain: remainder)
     }
 
+    private static func parseGetCollectionExpression(
+        _ input: String,
+        openParen: String.Index
+    ) throws -> MongoOperation {
+        let argAndRest = try extractParenthesizedArgAndRemainder(from: input, startingAt: openParen)
+        let argument = argAndRest.arg
+
+        guard !argument.isEmpty else {
+            throw MongoShellParseError.missingArgument(
+                String(localized: "getCollection requires a collection name")
+            )
+        }
+
+        let literal = try readStringLiteral(in: argument, startingAt: argument.startIndex)
+        let trailing = argument[literal.endIndex...].trimmingCharacters(in: .whitespacesAndNewlines)
+
+        guard trailing.isEmpty else {
+            throw MongoShellParseError.invalidSyntax(
+                String(localized: "getCollection takes a single quoted collection name")
+            )
+        }
+        guard !literal.value.isEmpty else {
+            throw MongoShellParseError.missingArgument(
+                String(localized: "getCollection requires a collection name")
+            )
+        }
+
+        return try parseAccessorChain(collection: literal.value, chain: argAndRest.remainder)
+    }
+
+    private static func parseAccessorChain(collection: String, chain: String) throws -> MongoOperation {
+        let trimmed = chain.trimmingCharacters(in: .whitespacesAndNewlines)
+
+        guard !trimmed.isEmpty else {
+            return .find(collection: collection, filter: "{}", options: MongoFindOptions())
+        }
+        guard trimmed.hasPrefix(".") else {
+            throw MongoShellParseError.invalidSyntax(
+                String(localized: "Expected '.method()' after the collection reference")
+            )
+        }
+
+        return try parseMethodChain(collection: collection, chain: String(trimmed.dropFirst()))
+    }
+
     /// Parse a db-level method call like db.getCollectionNames(), db.stats(), etc.
     /// Input is the string after "db." — e.g. "getCollectionNames()" or "createCollection(\"test\")"
     private static func parseDbLevelMethod(_ input: String) throws -> MongoOperation {
@@ -382,6 +384,53 @@ public struct MongoShellParser {
 
     // MARK: - Argument Extraction Helpers
 
+    private struct StringLiteral {
+        let value: String
+        let endIndex: String.Index
+    }
+
+    /// Read a quoted JavaScript string literal, returning its unescaped value and the index after
+    /// the closing quote. Accepts double quotes, single quotes and backticks.
+    private static func readStringLiteral(in text: String, startingAt start: String.Index) throws -> StringLiteral {
+        guard start < text.endIndex, isStringDelimiter(text[start]) else {
+            throw MongoShellParseError.invalidSyntax(
+                String(localized: "Expected a quoted collection name")
+            )
+        }
+
+        let quote = text[start]
+        var value = ""
+        var index = text.index(after: start)
+        var escapeNext = false
+
+        while index < text.endIndex {
+            let character = text[index]
+            index = text.index(after: index)
+
+            if escapeNext {
+                value.append(character)
+                escapeNext = false
+                continue
+            }
+            if character == "\\" {
+                escapeNext = true
+                continue
+            }
+            if character == quote {
+                return StringLiteral(value: value, endIndex: index)
+            }
+            value.append(character)
+        }
+
+        throw MongoShellParseError.invalidSyntax(
+            String(localized: "Unterminated string in the collection reference")
+        )
+    }
+
+    private static func isStringDelimiter(_ character: Character) -> Bool {
+        character == "\"" || character == "'" || character == "`"
+    }
+
     /// Extract content inside balanced parentheses starting at the given index
     private static func extractParenthesizedArg(from str: String, startingAt openParen: String.Index) throws -> String {
         let result = try extractParenthesizedArgAndRemainder(from: str, startingAt: openParen)