Make mongo connections simpler and more robust

Rather than trying to parse the connection string, just use `.db(databaseName)` after calling `MongoClient.connect`

src/mongo/editors/MongoFindResultEditor.ts

@@ -33,7 +33,7 @@ export class MongoFindResultEditor implements ICosmosEditor<IMongoDocument[]> {
         const collection: Collection = db.collection(this._command.collection);
         // NOTE: Intentionally creating a _new_ tree item rather than searching for a cached node in the tree because
         // the executed 'find' command could have a filter or projection that is not handled by a cached tree node
-        this._collectionTreeItem = new MongoCollectionTreeItem(dbTreeItem.connectionString, collection, dbTreeItem.id, this._command.arguments);
+        this._collectionTreeItem = new MongoCollectionTreeItem(collection, dbTreeItem.id, this._command.arguments);
         const documents: MongoDocumentTreeItem[] = <MongoDocumentTreeItem[]>await this._collectionTreeItem.loadMoreChildren(undefined, true);
         return documents.map((docTreeItem) => docTreeItem.document);
     }

src/mongo/languageClient.ts

@@ -49,8 +49,8 @@ export default class MongoDBLanguageClient {
 		context.subscriptions.push(disposable);
 	}
 
-	async connect(databaseConnectionString: string) {
-		await this.client.sendRequest('connect', { connectionString: databaseConnectionString });
+	async connect(connectionString: string, databaseName: string) {
+		await this.client.sendRequest('connect', { connectionString: connectionString, databaseName: databaseName });
 	}
 
 	disconnect(): void {

src/mongo/registerMongoCommands.ts

@@ -48,7 +48,7 @@ export function registerMongoCommands(context: vscode.ExtensionContext, actionHa
             await connectedDb.refresh();
         }
         connectedDb = node;
-        await languageClient.connect(connectedDb.treeItem.connectionString);
+        await languageClient.connect(connectedDb.treeItem.connectionString, connectedDb.treeItem.databaseName);
         connectedDb.treeItem.isConnected = true;
         await node.refresh();
     });

src/mongo/services/languageService.ts

@@ -44,8 +44,8 @@ export class LanguageService {
 
 		connection.onRequest('connect', (connectionParams) => {
 			MongoClient.connect(connectionParams.connectionString)
-				.then(db => {
-					this.db = db;
+				.then(account => {
+					this.db = account.db(connectionParams.databaseName);
 					this.schemaService.registerSchemas(this.db)
 						.then(schemas => {
 							this.configureSchemas(schemas);
@@ -83,4 +83,4 @@ export class LanguageService {
 			schemas
 		})
 	}
-}
+}

src/mongo/tree/MongoCollectionTreeItem.ts

@@ -20,17 +20,15 @@ export class MongoCollectionTreeItem implements IAzureParentTreeItem {
 
 	private readonly collection: Collection;
 	private readonly _query: string | undefined;
-	private readonly _databaseConnectionString: string;
 	private _cursor: Cursor | undefined;
 	private _hasMoreChildren: boolean = true;
 	private _parentId: string;
 	private _batchSize: number = DefaultBatchSize;
 
-	constructor(databaseConnectionString: string, collection: Collection, parentId: string, query?: string) {
+	constructor(collection: Collection, parentId: string, query?: string) {
 		this.collection = collection;
 		this._parentId = parentId;
 		this._query = query;
-		this._databaseConnectionString = databaseConnectionString;
 	}
 
 	public async update(documents: IMongoDocument[]): Promise<IMongoDocument[]> {
@@ -149,8 +147,7 @@ export class MongoCollectionTreeItem implements IAzureParentTreeItem {
 	}
 
 	private async drop(): Promise<string> {
-		const db: Db = await MongoClient.connect(this._databaseConnectionString);
-		await db.dropCollection(this.collection.collectionName);
+		await this.collection.drop();
 		return `Dropped collection ${this.collection.collectionName}.`;
 	}
 

src/mongo/tree/MongoDatabaseTreeItem.ts

@@ -17,40 +17,26 @@ export class MongoDatabaseTreeItem implements IAzureParentTreeItem {
 	public readonly contextValue: string = MongoDatabaseTreeItem.contextValue;
 	public readonly childTypeLabel: string = "Collection";
 	public readonly connectionString: string;
+	public readonly databaseName: string;
 
-	private readonly _databaseName: string;
-	private readonly _accountConnectionString: string;
 	private readonly _parentId: string;
 	public isConnected: boolean = false;
 
-	constructor(databaseName: string, accountConnectionString: string, parentId: string) {
-		this._databaseName = databaseName;
-		this._accountConnectionString = accountConnectionString;
+	constructor(databaseName: string, connectionString: string, parentId: string) {
+		this.databaseName = databaseName;
+		this.connectionString = connectionString;
 		this._parentId = parentId;
-		const uri: vscode.Uri = vscode.Uri.parse(accountConnectionString);
-		let path: string = uri.path;
-		if (path.startsWith('/')) {
-			path = path.slice(1);
-		}
-		if (path.endsWith('/')) {
-			path = path.slice(0, -1);
-		}
-		if (path) {
-			this.connectionString = `${uri.scheme}://${uri.authority}/${path}/${this._databaseName}?${uri.query}`;
-		} else {
-			this.connectionString = `${uri.scheme}://${uri.authority}/${this._databaseName}?${uri.query}`;
-		}
 	}
 
 	public get label(): string {
 		if (this.isConnected) {
-			return this._databaseName + " (Connected)";
+			return this.databaseName + " (Connected)";
 		}
-		return this._databaseName;
+		return this.databaseName;
 	}
 
 	public get id(): string {
-		return `${this._parentId}/${this._databaseName}`;
+		return `${this._parentId}/${this.databaseName}`;
 	}
 
 	get iconPath(): any {
@@ -67,7 +53,7 @@ export class MongoDatabaseTreeItem implements IAzureParentTreeItem {
 	public async loadMoreChildren(_node: IAzureNode, _clearCache: boolean): Promise<IAzureTreeItem[]> {
 		const db: Db = await this.getDb();
 		const collections: Collection[] = await db.collections();
-		return collections.map(collection => new MongoCollectionTreeItem(this.connectionString, collection, this.id));
+		return collections.map(collection => new MongoCollectionTreeItem(collection, this.id));
 	}
 
 	public async createChild(_node: IAzureNode, showCreatingNode: (label: string) => void): Promise<IAzureTreeItem> {
@@ -98,7 +84,8 @@ export class MongoDatabaseTreeItem implements IAzureParentTreeItem {
 	}
 
 	public async getDb(): Promise<Db> {
-		return await MongoClient.connect(this.connectionString);
+		const accountConnection = await MongoClient.connect(this.connectionString);
+		return accountConnection.db(this.databaseName);
 	}
 
 	executeCommand(command: MongoCommand): Thenable<string> {
@@ -107,7 +94,7 @@ export class MongoDatabaseTreeItem implements IAzureParentTreeItem {
 				.then(db => {
 					const collection = db.collection(command.collection);
 					if (collection) {
-						const result = new MongoCollectionTreeItem(this.connectionString, collection, command.arguments).executeCommand(command.name, command.arguments);
+						const result = new MongoCollectionTreeItem(collection, command.arguments).executeCommand(command.name, command.arguments);
 						if (result) {
 							return result;
 						}
@@ -130,7 +117,7 @@ export class MongoDatabaseTreeItem implements IAzureParentTreeItem {
 		// However, we can 'insert' and then 'delete' a document, which has the side-effect of creating an empty collection
 		const result = await newCollection.insertOne({});
 		await newCollection.deleteOne({ _id: result.insertedId });
-		return new MongoCollectionTreeItem(this.connectionString, newCollection, this.id);
+		return new MongoCollectionTreeItem(newCollection, this.id);
 	}
 
 	executeCommandInShell(command: MongoCommand): Thenable<string> {
@@ -151,9 +138,9 @@ export class MongoDatabaseTreeItem implements IAzureParentTreeItem {
 	}
 
 	private async createShell(shellPath: string): Promise<Shell> {
-		return <Promise<null>>Shell.create(shellPath, this._accountConnectionString)
+		return <Promise<null>>Shell.create(shellPath, this.connectionString)
 			.then(shell => {
-				return shell.useDatabase(this._databaseName).then(() => shell);
+				return shell.useDatabase(this.databaseName).then(() => shell);
 			}, error => vscode.window.showErrorMessage(error));
 	}
 }