Add a setting to configure the Mongo shell batch size (#1598)

* Generalize batch size setting

package.json

@@ -960,6 +960,11 @@
                     "default": true,
                     "description": "Show warning dialog when uploading a document to the cloud."
                 },
+                "azureDatabases.batchSize": {
+                    "type": "number",
+                    "description": "The batch size to be used when querying Azure Database resources.",
+                    "default": 50
+                },
                 "azureDatabases.enableOutputTimestamps": {
                     "type": "boolean",
                     "default": true,

src/constants.ts

@@ -41,8 +41,6 @@ export function getResourcesPath(): string {
     return ext.context.asAbsolutePath('resources');
 }
 
-export const defaultBatchSize: number = 50;
-
 export const doubleClickDebounceDelay = 500; //milliseconds
 
 export const defaultStoredProcedure =

src/docdb/tree/DocDBTreeItemBase.ts

@@ -5,7 +5,7 @@
 
 import { DocumentClient, FeedOptions, QueryError, QueryIterator } from 'documentdb';
 import { AzExtTreeItem, AzureParentTreeItem, AzureTreeItem } from 'vscode-azureextensionui';
-import { defaultBatchSize } from '../../constants';
+import { getBatchSizeSetting } from '../../utils/workspacUtils';
 import { IDocDBTreeRoot } from './IDocDBTreeRoot';
 
 /**
@@ -19,7 +19,7 @@ export abstract class DocDBTreeItemBase<T> extends AzureParentTreeItem<IDocDBTre
 
     private _hasMoreChildren: boolean = true;
     private _iterator: QueryIterator<T> | undefined;
-    private _batchSize: number = defaultBatchSize;
+    private _batchSize: number = getBatchSizeSetting();
 
     public hasMoreChildrenImpl(): boolean {
         return this._hasMoreChildren;
@@ -29,12 +29,15 @@ export abstract class DocDBTreeItemBase<T> extends AzureParentTreeItem<IDocDBTre
 
     public abstract getIterator(client: DocumentClient, feedOptions: FeedOptions): Promise<QueryIterator<T>>;
 
+    public async refreshImpl(): Promise<void> {
+        this._batchSize = getBatchSizeSetting();
+    }
+
     public async loadMoreChildrenImpl(clearCache: boolean): Promise<AzExtTreeItem[]> {
         if (clearCache || this._iterator === undefined) {
             this._hasMoreChildren = true;
             const client = this.root.getDocumentClient();
-            this._iterator = await this.getIterator(client, { maxItemCount: defaultBatchSize });
-            this._batchSize = defaultBatchSize;
+            this._iterator = await this.getIterator(client, { maxItemCount: this._batchSize });
         }
 
         const resources: T[] = [];

src/extensionVariables.ts

@@ -37,6 +37,7 @@ export namespace ext {
         export const mongoShellArgs = 'mongo.shell.args';
         export const documentLabelFields = 'cosmosDB.documentLabelFields';
         export const mongoShellTimeout = 'mongo.shell.timeout';
+        export const batchSize = 'azureDatabases.batchSize';
 
         export namespace vsCode {
             export const proxyStrictSSL = "http.proxyStrictSSL";

src/mongo/MongoShell.ts

@@ -8,6 +8,7 @@ import * as vscode from 'vscode';
 import { parseError } from 'vscode-azureextensionui';
 import { InteractiveChildProcess } from '../utils/InteractiveChildProcess';
 import { randomUtils } from '../utils/randomUtils';
+import { getBatchSizeSetting } from '../utils/workspacUtils';
 import { wrapError } from '../utils/wrapError';
 
 const timeoutMessage = "Timed out trying to execute the Mongo script. To use a longer timeout, modify the VS Code 'mongo.shell.timeout' setting.";
@@ -53,6 +54,9 @@ export class MongoShell extends vscode.Disposable {
             // to catch any errors related to the start-up of the process before trying to write to it.
             await shell.executeScript("");
 
+            // Configure the batch size
+            await shell.executeScript(`DBQuery.shellBatchSize = ${getBatchSizeSetting()}`);
+
             return shell;
         } catch (error) {
             throw wrapCheckOutputWindow(error);

src/mongo/tree/MongoCollectionTreeItem.ts

@@ -8,11 +8,12 @@ import { BulkWriteOpResultObject, Collection, CollectionInsertManyOptions, Curso
 import * as _ from 'underscore';
 import * as vscode from 'vscode';
 import { AzExtTreeItem, AzureParentTreeItem, DialogResponses, IActionContext, ICreateChildImplContext, UserCancelledError } from 'vscode-azureextensionui';
-import { defaultBatchSize, getThemeAgnosticIconPath } from '../../constants';
+import { getThemeAgnosticIconPath } from '../../constants';
 import { IEditableTreeItem } from '../../DatabasesFileSystem';
 import { ext } from '../../extensionVariables';
 import { nonNullValue } from '../../utils/nonNull';
 import { getDocumentTreeItemLabel } from '../../utils/vscodeUtils';
+import { getBatchSizeSetting } from '../../utils/workspacUtils';
 import { MongoCommand } from '../MongoCommand';
 import { IMongoTreeRoot } from './IMongoTreeRoot';
 import { IMongoDocument, MongoDocumentTreeItem } from './MongoDocumentTreeItem';
@@ -40,7 +41,7 @@ export class MongoCollectionTreeItem extends AzureParentTreeItem<IMongoTreeRoot>
     private readonly _projection: object | undefined;
     private _cursor: Cursor | undefined;
     private _hasMoreChildren: boolean = true;
-    private _batchSize: number = defaultBatchSize;
+    private _batchSize: number = getBatchSizeSetting();
 
     constructor(parent: AzureParentTreeItem, collection: Collection, findArgs?: {}[]) {
         super(parent);
@@ -106,6 +107,7 @@ export class MongoCollectionTreeItem extends AzureParentTreeItem<IMongoTreeRoot>
     }
 
     public async refreshImpl(): Promise<void> {
+        this._batchSize = getBatchSizeSetting();
         ext.fileSystem.fireChangedEvent(this);
     }
 
@@ -126,11 +128,10 @@ export class MongoCollectionTreeItem extends AzureParentTreeItem<IMongoTreeRoot>
 
     public async loadMoreChildrenImpl(clearCache: boolean): Promise<AzExtTreeItem[]> {
         if (clearCache || this._cursor === undefined) {
-            this._cursor = this.collection.find(this._query).batchSize(defaultBatchSize);
+            this._cursor = this.collection.find(this._query).batchSize(this._batchSize);
             if (this._projection) {
                 this._cursor = this._cursor.project(this._projection);
             }
-            this._batchSize = defaultBatchSize;
         }
 
         const documents: IMongoDocument[] = [];

src/utils/workspacUtils.ts

@@ -4,9 +4,16 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as vscode from 'vscode';
+import { ext } from '../extensionVariables';
+import { nonNullValue } from './nonNull';
 
 // tslint:disable-next-line: export-name
 export function getRootPath(): string | undefined {
     // if this is a multi-root workspace, return undefined
     return vscode.workspace.workspaceFolders && vscode.workspace.workspaceFolders.length === 1 ? vscode.workspace.workspaceFolders[0].uri.fsPath : undefined;
 }
+
+export function getBatchSizeSetting(): number {
+    const config: vscode.WorkspaceConfiguration = vscode.workspace.getConfiguration();
+    return nonNullValue(config.get<number>(ext.settingsKeys.batchSize), 'batchSize');
+}