VSCODE-56: Display document count in tree view

package.json

@@ -264,6 +264,10 @@
         "command": "mdb.viewCollectionDocuments",
         "title": "View Documents"
       },
+      {
+        "command": "mdb.refreshDocumentList",
+        "title": "Refresh"
+      },
       {
         "command": "mdb.copyCollectionName",
         "title": "Copy Collection Name"
@@ -423,6 +427,10 @@
           "command": "mdb.viewCollectionDocuments",
           "when": "view == mongoDB && viewItem == documentListTreeItem"
         },
+        {
+          "command": "mdb.refreshDocumentList",
+          "when": "view == mongoDB && viewItem == documentListTreeItem"
+        },
         {
           "command": "mdb.refreshSchema",
           "when": "view == mongoDB && viewItem == schemaTreeItem"
@@ -504,6 +512,10 @@
           "command": "mdb.viewCollectionDocuments",
           "when": "false"
         },
+        {
+          "command": "mdb.refreshDocumentList",
+          "when": "false"
+        },
         {
           "command": "mdb.copyCollectionName",
           "when": "false"
@@ -647,6 +659,7 @@
     "mongodb-data-service": "^16.8.1",
     "mongodb-ns": "^2.2.0",
     "mongodb-schema": "^8.2.5",
+    "numeral": "^2.0.6",
     "react": "^16.13.1",
     "react-dom": "^16.13.1",
     "react-redux": "^7.2.0",

src/explorer/collectionTreeItem.ts

@@ -1,6 +1,7 @@
 import * as vscode from 'vscode';
 const path = require('path');
 
+import { createLogger } from '../logging';
 import DocumentListTreeItem, {
   CollectionTypes,
   MAX_DOCUMENTS_VISIBLE
@@ -10,6 +11,8 @@ import TreeItemParent from './treeItemParentInterface';
 import SchemaTreeItem from './schemaTreeItem';
 import { getImagesPath } from '../extensionConstants';
 
+const log = createLogger('tree view collection folder');
+
 type CollectionModelType = {
   name: string;
   type: CollectionTypes;
@@ -36,9 +39,11 @@ export default class CollectionTreeItem extends vscode.TreeItem
   collection: CollectionModelType;
   collectionName: string;
   databaseName: string;
+  namespace: string;
 
   private _dataService: any;
   private _type: CollectionTypes;
+  documentCount: number | null;
 
   isExpanded: boolean;
 
@@ -50,6 +55,7 @@ export default class CollectionTreeItem extends vscode.TreeItem
     dataService: any,
     isExpanded: boolean,
     cacheIsUpToDate: boolean,
+    cachedDocumentCount: number | null,
     existingDocumentListChild?: DocumentListTreeItem,
     existingSchemaChild?: SchemaTreeItem,
     existingIndexListChild?: IndexListTreeItem
@@ -64,12 +70,15 @@ export default class CollectionTreeItem extends vscode.TreeItem
     this.collection = collection;
     this.collectionName = collection.name;
     this.databaseName = databaseName;
+    this.namespace = `${this.databaseName}.${this.collectionName}`;
 
     this._type = collection.type; // Type can be `collection` or `view`.
     this._dataService = dataService;
 
     this.isExpanded = isExpanded;
 
+    this.documentCount = cachedDocumentCount;
+
     this.cacheIsUpToDate = cacheIsUpToDate;
 
     this._documentListChild = existingDocumentListChild
@@ -81,7 +90,8 @@ export default class CollectionTreeItem extends vscode.TreeItem
         this._dataService,
         false, // Collapsed.
         MAX_DOCUMENTS_VISIBLE,
-        false, // No more documents to show.
+        this.documentCount,
+        this.refreshDocumentCount,
         false, // Cache is not up to date.
         [] // Empty cache.
       );
@@ -119,9 +129,13 @@ export default class CollectionTreeItem extends vscode.TreeItem
     return element;
   }
 
-  getChildren(): Thenable<any[]> {
+  async getChildren(): Promise<any[]> {
     if (!this.isExpanded) {
-      return Promise.resolve([]);
+      return [];
+    }
+
+    if (this.documentCount === null) {
+      await this.refreshDocumentCount();
     }
 
     // Update cache if one of the children has been expanded/collapsed.
@@ -130,11 +144,7 @@ export default class CollectionTreeItem extends vscode.TreeItem
     }
 
     if (this.cacheIsUpToDate) {
-      return Promise.resolve([
-        this._documentListChild,
-        this._schemaChild,
-        this._indexListChild
-      ]);
+      return [this._documentListChild, this._schemaChild, this._indexListChild];
     }
 
     this.cacheIsUpToDate = true;
@@ -143,11 +153,7 @@ export default class CollectionTreeItem extends vscode.TreeItem
     // is ensure to be set by vscode.
     this.rebuildChildrenCache();
 
-    return Promise.resolve([
-      this._documentListChild,
-      this._schemaChild,
-      this._indexListChild
-    ]);
+    return [this._documentListChild, this._schemaChild, this._indexListChild];
   }
 
   rebuildDocumentListTreeItem(): void {
@@ -158,7 +164,8 @@ export default class CollectionTreeItem extends vscode.TreeItem
       this._dataService,
       this._documentListChild.isExpanded,
       this._documentListChild.getMaxDocumentsToShow(),
-      this._documentListChild.hasMoreDocumentsToShow,
+      this.documentCount,
+      this.refreshDocumentCount,
       this._documentListChild.cacheIsUpToDate,
       this._documentListChild.getChildrenCache()
     );
@@ -209,15 +216,18 @@ export default class CollectionTreeItem extends vscode.TreeItem
     this.cacheIsUpToDate = false;
   }
 
-  onDidExpand(): Promise<boolean> {
+  async onDidExpand(): Promise<boolean> {
     this.isExpanded = true;
     this.cacheIsUpToDate = false;
 
-    return Promise.resolve(true);
+    await this.refreshDocumentCount();
+
+    return true;
   }
 
   resetCache(): void {
     this.cacheIsUpToDate = false;
+    this.documentCount = null;
 
     this._documentListChild = new DocumentListTreeItem(
       this.collectionName,
@@ -226,7 +236,8 @@ export default class CollectionTreeItem extends vscode.TreeItem
       this._dataService,
       false, // Collapsed.
       MAX_DOCUMENTS_VISIBLE,
-      false, // No more documents to show.
+      this.documentCount,
+      this.refreshDocumentCount,
       false, // Cache is not up to date.
       [] // Empty cache.
     );
@@ -268,6 +279,43 @@ export default class CollectionTreeItem extends vscode.TreeItem
     return this._documentListChild.getMaxDocumentsToShow();
   }
 
+  getCount(): Promise<number> {
+    log.info(`fetching document count from namespace ${this.namespace}`);
+
+    return new Promise((resolve, reject) => {
+      this._dataService.estimatedCount(
+        this.namespace,
+        {}, // No options.
+        (err: Error | undefined, count: number) => {
+          if (err) {
+            return reject(
+              new Error(
+                `Unable to get collection document count: ${err.message}`
+              )
+            );
+          }
+
+          return resolve(count);
+        }
+      );
+    });
+  }
+
+  refreshDocumentCount = async (): Promise<boolean> => {
+    try {
+      // We fetch the document when we expand in order to show
+      // the document count in the document list tree item `description`.
+      this.documentCount = await this.getCount();
+    } catch (err) {
+      vscode.window.showInformationMessage(
+        `Unable to fetch document count: ${err}`
+      );
+      return false;
+    }
+
+    return true;
+  };
+
   async onDropCollectionClicked(): Promise<boolean> {
     const collectionName = this.collectionName;
 

src/explorer/databaseTreeItem.ts

@@ -85,6 +85,7 @@ export default class DatabaseTreeItem extends vscode.TreeItem
           this._dataService,
           pastChildrenCache[collectionName].isExpanded,
           pastChildrenCache[collectionName].cacheIsUpToDate,
+          pastChildrenCache[collectionName].documentCount,
           pastChildrenCache[collectionName].getDocumentListChild(),
           pastChildrenCache[collectionName].getSchemaChild(),
           pastChildrenCache[collectionName].getIndexListChild()
@@ -122,6 +123,7 @@ export default class DatabaseTreeItem extends vscode.TreeItem
               this._dataService,
               pastChildrenCache[collection.name].isExpanded,
               pastChildrenCache[collection.name].cacheIsUpToDate,
+              pastChildrenCache[collection.name].documentCount,
               pastChildrenCache[collection.name].getDocumentListChild(),
               pastChildrenCache[collection.name].getSchemaChild(),
               pastChildrenCache[collection.name].getIndexListChild()
@@ -132,7 +134,8 @@ export default class DatabaseTreeItem extends vscode.TreeItem
               this.databaseName,
               this._dataService,
               false, // Not expanded.
-              false // No cache.
+              false, // No cache.
+              null // No document count yet.
             );
           }
         });