From 3e0e16ba7a9bf224b692316386c6eba46592e546 Mon Sep 17 00:00:00 2001
From: Sandeep Somavarapu <sasomava@microsoft.com>
Date: Thu, 11 May 2017 16:30:52 +0200
Subject: [PATCH] Time out execution of script in shell after 5s

---
 src/mongo/commands.ts |  2 +-
 src/mongo/mongo.ts    | 61 ++++++++++++++-----------------------------
 src/mongo/shell.ts    | 12 ++++++++-
 3 files changed, 31 insertions(+), 44 deletions(-)

diff --git a/src/mongo/commands.ts b/src/mongo/commands.ts
index 4b9a68c..5613e60 100644
--- a/src/mongo/commands.ts
+++ b/src/mongo/commands.ts
@@ -40,7 +40,7 @@ export class MongoCommands {
 							}
 							editorBuilder.insert(new vscode.Position(0, 0), result);
 						});
-					});
+					}, error => vscode.window.showErrorMessage(error));
 			});
 		return script;
 	}
diff --git a/src/mongo/mongo.ts b/src/mongo/mongo.ts
index b6b3bff..5ace606 100644
--- a/src/mongo/mongo.ts
+++ b/src/mongo/mongo.ts
@@ -195,7 +195,6 @@ export class Server implements IMongoResource {
 export class Database implements IMongoResource {
 
 	readonly contextKey: string = 'mongoDb';
-	private shell: Shell;
 
 	private _onChange: EventEmitter<void> = new EventEmitter<void>();
 	readonly onChange: Event<void> = this._onChange.event;
@@ -237,11 +236,11 @@ export class Database implements IMongoResource {
 							return result;
 						}
 					}
-					return reportProgress(this.getShell().then(() => this.shell.exec(script.script)), 'Executing script');
+					return reportProgress(this.executeScriptInShell(script), 'Executing script');
 				});
 		}
 		const result = this.executeCommand(script.command, script.arguments);
-		return result ? result : reportProgress(this.getShell().then(() => this.shell.exec(script.script)), 'Executing script');
+		return result ? result : reportProgress(this.executeScriptInShell(script), 'Executing script');
 	}
 
 	updateDocuments(documentOrDocuments: any, collectionName: string): Thenable<string> {
@@ -283,10 +282,22 @@ export class Database implements IMongoResource {
 		return this.getDb().then(db => new Collection(db.collection(collection), this));
 	}
 
-	private getShell(): Promise<void> {
-		if (this.shell) {
-			return Promise.resolve();
+	executeScriptInShell(script: MongoScript): Thenable<string> {
+		return this.getShell().then(shell => shell.exec(script.script));
+	}
+
+	executeCommand(command: string, args?: string): Thenable<string> {
+		try {
+			if (command === 'createCollection') {
+				return reportProgress(this.createCollection(stripQuotes(args)).then(() => JSON.stringify({ 'Created': 'Ok' })), 'Creating collection');
+			}
+			return null;
+		} catch (error) {
+			return Promise.resolve(error);
 		}
+	}
+
+	private getShell(): Promise<Shell> {
 		const shellPath = <string>vscode.workspace.getConfiguration().get('mongo.shell.path')
 		if (!shellPath) {
 			return <Promise<null>>vscode.window.showInputBox({
@@ -299,46 +310,12 @@ export class Database implements IMongoResource {
 		}
 	}
 
-	private createShell(shellPath: string): Promise<void> {
+	private createShell(shellPath: string): Promise<Shell> {
 		return <Promise<null>>Shell.create(shellPath, this.server.id)
 			.then(shell => {
-				this.shell = shell;
-				return this.shell.useDatabase(this.id).then(() => null);
+				return shell.useDatabase(this.id).then(() => shell);
 			}, error => vscode.window.showErrorMessage(error));
 	}
-
-	_executeScript(script: string): Promise<string> {
-		return this.getDb().then(db => {
-			return db.eval(new Code(`function() {
-				var result = ${script};
-				if (result.hasNext) {
-					let results = [];
-					for (let counter = 0; counter < 20 && result.hasNext(); counter++) {
-						results.push(result.next());
-					}
-					return results;
-				} else {
-					return result;
-				}
-			}`), [], { readPreference: ReadPreference.PRIMARY }).then(result => {
-					db.close();
-					return JSON.stringify(result, null, '\t')
-				}, error => {
-					console.log(error);
-				});
-		});
-	}
-
-	executeCommand(command: string, args?: string): Thenable<string> {
-		try {
-			if (command === 'createCollection') {
-				return reportProgress(this.createCollection(stripQuotes(args)).then(() => JSON.stringify({ 'Created': 'Ok' })), 'Creating collection');
-			}
-			return null;
-		} catch (error) {
-			return Promise.resolve(error);
-		}
-	}
 }
 
 export class Collection implements IMongoResource {
diff --git a/src/mongo/shell.ts b/src/mongo/shell.ts
index eadd757..5ace841 100644
--- a/src/mongo/shell.ts
+++ b/src/mongo/shell.ts
@@ -85,13 +85,23 @@ export class Shell {
 		};
 
 		return await new Promise<string>((c, e) => {
+			let executed = false;
+			const handler = setTimeout(() => {
+				if (!executed) {
+					e('Timed out executing ' + script);
+				}
+			}, 5000);
 			const disposable = this.onResult.event(result => {
 				disposable.dispose();
 				let lines = (<string>result.result).split(os.EOL).filter(line => !!line && line !== 'Type "it" for more');
 				lines = lines[lines.length - 1] === 'Type "it" for more' ? lines.splice(lines.length - 1, 1) : lines;
 				let value = lines.join(os.EOL);
+				executed = true;
 				c(lines.join(os.EOL));
-			})
+				if (handler) {
+					clearTimeout(handler);
+				}
+			});
 		});
 	}