From e507817b841b84c3d1193601cfe02b4ceb504123 Mon Sep 17 00:00:00 2001
From: Shadowfacts <me@shadowfacts.net>
Date: Mon, 27 Jan 2020 22:20:02 -0500
Subject: [PATCH] Improve Mongo output handling

---
 MongoView/MongoEvaluator.swift                | 54 +++++--------------
 MongoView/Node.swift                          |  9 +++-
 .../QueryViewController.swift                 |  2 +-
 3 files changed, 23 insertions(+), 42 deletions(-)

diff --git a/MongoView/MongoEvaluator.swift b/MongoView/MongoEvaluator.swift
index 82b4424..d3bb83a 100644
--- a/MongoView/MongoEvaluator.swift
+++ b/MongoView/MongoEvaluator.swift
@@ -7,63 +7,37 @@
 //
 
 import Foundation
-import JavaScriptCore
 import MongoSwift
 
 struct MongoEvaluator {
     
-    static let context: JSContext = {
-        let context = JSContext()!
-        let date: @convention(block) (JSValue) -> Dictionary<String, Any> = { (input) in
-            ["$date": input]
-        }
-        let functions: [NSString: @convention(block) (JSValue) -> Dictionary<String, Any>] = [
-            "ObjectId": { (input) in
-                ["$oid": input]
-            },
-            "NumberLong": { (input) in
-                ["$numberLong": input]
-            },
-            "NumberInt": { (input) in
-                ["$numberInt": input]
-            },
-            "NumberDecimal": { (input) in
-                ["$numberDecimal": input]
-            },
-            "Date": date,
-            "ISODate": date,
-        ]
-        functions.forEach { (key, value) in
-            context.setObject(value, forKeyedSubscript: key)
-        }
-        let dbRef: @convention(block) (String, String) -> Dictionary<String, Any> = { (ref, id) in
-            ["$ref": ref, "$id": ["$oid": id]]
-        }
-        context.setObject(dbRef, forKeyedSubscript: "DBRef" as NSString)
-        return context
-    }()
+    static let decoder = BSONDecoder()
 
-    static func parse(shellJSON: String) -> Document? {
-        guard !shellJSON.isEmpty,
-            let result = context.evaluateScript("JSON.stringify(\(shellJSON))") else {
-                return nil
+    static func parse(extendedJSON: String) -> BSON? {
+        guard !extendedJSON.isEmpty else {
+            return nil
         }
         do {
-            return try Document(fromJSON: result.toString())
+            return try decoder.decode(BSON.self, from: extendedJSON)
         } catch {
-            print("Unable to parse Document: \(error)")
+            print("Unable to parse BSON: \(error)")
             return nil
         }
     }
     
-    static func eval(command: String, connectingTo connStr: String) -> [Document] {
+    static func eval(command: String, connectingTo connStr: String) -> [BSON] {
+        let realCommand = """
+Object.prototype.printExtJSON = function() { print(JSON.stringify(this)); };
+Array.prototype.printExtJSON = function() { this.map(JSON.stringify).forEach(it => print(it)); };
+\(command).printExtJSON()
+"""
         print("Running command \"\(command)\" against \(connStr)")
         
         let mongoOutputPipe = Pipe()
 
         let mongoProcess = Process()
         mongoProcess.launchPath = "/usr/local/bin/mongo"
-        mongoProcess.arguments = [connStr, "--quiet", "--norc", "--eval", command]
+        mongoProcess.arguments = [connStr, "--quiet", "--norc", "--eval", realCommand]
         mongoProcess.standardOutput = mongoOutputPipe
         mongoProcess.launch()
         
@@ -78,7 +52,7 @@ struct MongoEvaluator {
         mongoOutputHandle.closeFile()
         
         let lines = strs.joined(separator: "").components(separatedBy: "\n")
-        return lines.compactMap(parse(shellJSON:))
+        return lines.compactMap(parse(extendedJSON:))
     }
     
 }
diff --git a/MongoView/Node.swift b/MongoView/Node.swift
index 875267c..b53cf72 100644
--- a/MongoView/Node.swift
+++ b/MongoView/Node.swift
@@ -37,8 +37,15 @@ class Node {
     }
 
     init(key: Key? = nil, value: BSON) {
-        self.key = key
         self.value = value
+        
+        if key == nil,
+            case let .document(doc) = value,
+            case let .objectId(id) = doc["_id"] {
+            self.key = .objectId(id)
+        } else {
+            self.key = key
+        }
     }
     
     convenience init(document: Document) {
diff --git a/MongoView/View Controllers/QueryViewController.swift b/MongoView/View Controllers/QueryViewController.swift
index 1201673..2d0c649 100644
--- a/MongoView/View Controllers/QueryViewController.swift	
+++ b/MongoView/View Controllers/QueryViewController.swift	
@@ -20,7 +20,7 @@ class QueryViewController: NSViewController {
     let collection: DatabaseCollection
     
     var defaultQuery: String {
-        "db.getCollection('\(collection.name)').find({})"
+        "db.getCollection('\(collection.name)').find({}).toArray()"
     }
     
     var hasQueryChanged: Bool {