From ae7e176b8aab0ab2e5c387c7605c8f7308b866f3 Mon Sep 17 00:00:00 2001
From: Prashanth <prsri@microsoft.com>
Date: Mon, 30 Apr 2018 13:49:28 -0700
Subject: [PATCH] Force refresh a document's tree item post-update. Fixes #145

---
 src/docdb/editors/DocDBDocumentNodeEditor.ts   |  4 +++-
 src/mongo/editors/MongoCollectionNodeEditor.ts |  4 +++-
 src/mongo/editors/MongoDocumentNodeEditor.ts   |  4 +++-
 src/mongo/editors/MongoFindOneResultEditor.ts  |  7 +++++--
 src/utils/vscodeUtils.ts                       | 11 +++++++----
 5 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/src/docdb/editors/DocDBDocumentNodeEditor.ts b/src/docdb/editors/DocDBDocumentNodeEditor.ts
index bd93b194..33b1b21d 100644
--- a/src/docdb/editors/DocDBDocumentNodeEditor.ts
+++ b/src/docdb/editors/DocDBDocumentNodeEditor.ts
@@ -24,7 +24,9 @@ export class DocDBDocumentNodeEditor implements ICosmosEditor<RetrievedDocument>
     }
 
     public async update(document: RetrievedDocument): Promise<RetrievedDocument> {
-        return await this._documentNode.treeItem.update(document);
+        const updatedDoc = await this._documentNode.treeItem.update(document);
+        this._documentNode.refresh();
+        return updatedDoc;
     }
 
     public get id(): string {
diff --git a/src/mongo/editors/MongoCollectionNodeEditor.ts b/src/mongo/editors/MongoCollectionNodeEditor.ts
index d0f1277a..07ee3c57 100644
--- a/src/mongo/editors/MongoCollectionNodeEditor.ts
+++ b/src/mongo/editors/MongoCollectionNodeEditor.ts
@@ -7,7 +7,7 @@ import { IAzureParentNode, IAzureNode } from "vscode-azureextensionui";
 import { IMongoDocument, MongoDocumentTreeItem } from "../tree/MongoDocumentTreeItem";
 import { ICosmosEditor } from "../../CosmosEditorManager";
 import { MongoCollectionTreeItem } from "../tree/MongoCollectionTreeItem";
-import { getNodeEditorLabel } from '../../utils/vscodeUtils';
+import { getNodeEditorLabel, getDocumentTreeItemLabel } from '../../utils/vscodeUtils';
 // tslint:disable:no-var-requires
 const EJSON = require("mongodb-extended-json");
 
@@ -38,6 +38,8 @@ export class MongoCollectionNodeEditor implements ICosmosEditor<IMongoDocument[]
             const documentNode = documentNodes.find((node) => node.treeItem.document._id.toString() === updatedDoc._id.toString());
             if (documentNode) {
                 documentNode.treeItem.document = updatedDoc;
+                documentNode.treeItem.label = getDocumentTreeItemLabel(documentNode.treeItem.document);
+                documentNode.refresh();
             }
         }
     }
diff --git a/src/mongo/editors/MongoDocumentNodeEditor.ts b/src/mongo/editors/MongoDocumentNodeEditor.ts
index 49ce0436..2bc434e7 100644
--- a/src/mongo/editors/MongoDocumentNodeEditor.ts
+++ b/src/mongo/editors/MongoDocumentNodeEditor.ts
@@ -25,7 +25,9 @@ export class MongoDocumentNodeEditor implements ICosmosEditor<IMongoDocument> {
     }
 
     public async update(document: IMongoDocument): Promise<IMongoDocument> {
-        return await this._documentNode.treeItem.update(document);
+        const updatedDoc = await this._documentNode.treeItem.update(document);
+        this._documentNode.refresh();
+        return updatedDoc;
     }
 
     public get id(): string {
diff --git a/src/mongo/editors/MongoFindOneResultEditor.ts b/src/mongo/editors/MongoFindOneResultEditor.ts
index 8bde4616..db54da13 100644
--- a/src/mongo/editors/MongoFindOneResultEditor.ts
+++ b/src/mongo/editors/MongoFindOneResultEditor.ts
@@ -34,12 +34,15 @@ export class MongoFindOneResultEditor implements ICosmosEditor<IMongoDocument> {
 
     public async update(newDocument: IMongoDocument): Promise<IMongoDocument> {
         const node = await this._tree.findNode(this.id);
+        let result: IMongoDocument;
         if (node) {
-            return (<IAzureNode<MongoDocumentTreeItem>>node).treeItem.update(newDocument);
+            result = await (<IAzureNode<MongoDocumentTreeItem>>node).treeItem.update(newDocument);
+            node.refresh();
         }
         // If the node isn't cached already, just update it to Mongo directly (without worrying about updating the tree)
         const db = await this._databaseNode.treeItem.getDb();
-        return await MongoDocumentTreeItem.update(db.collection(this._collectionName), newDocument);
+        result = await MongoDocumentTreeItem.update(db.collection(this._collectionName), newDocument);
+        return result;
     }
 
     public get id(): string {
diff --git a/src/utils/vscodeUtils.ts b/src/utils/vscodeUtils.ts
index 278f6763..0a46f12f 100644
--- a/src/utils/vscodeUtils.ts
+++ b/src/utils/vscodeUtils.ts
@@ -108,12 +108,15 @@ function isAccountTreeItem(treeItem: IAzureTreeItem): boolean {
 export function getDocumentTreeItemLabel(document: IMongoDocument | RetrievedDocument): string {
     for (let field of documentDefaultFields) {
         if (document.hasOwnProperty(field)) {
-            if (document[field]) { //ignore if false-y value: null, undefined, "".
-                if (typeof document[field] === "string") {
-                    return document[field];
+            let value = document[field];
+            if (value) { //ignore if false-y value: null, undefined, "".
+                if (typeof value === "string") {
+                    return value;
+                } else if (typeof value === "number") {
+                    return value.toString();
                 } else { // A simple instanceOf check doesn't seem to work. I run into this issue detailed here :
                     //https://libertyseeds.ca/2015/01/03/Mongo-ObjectID-fails-instanceof-type-test-in-nodeunit/
-                    const result = document[field].toString();
+                    const result = value.toString();
                     if (!result.startsWith("[object ")) {
                         return result;
                     }