Added multiple location option to package management dialog (#9790)

* Added multiple location option to package management dialog

extensions/machine-learning-services/src/common/constants.ts

@@ -47,7 +47,8 @@ export const msgYes = localize('msgYes', "Yes");
 export const msgNo = localize('msgNo', "No");
 export const managePackageCommandError = localize('mls.managePackages.error', "Either no connection is available or the server does not have external script enabled.");
 export function taskFailedError(taskName: string, err: string): string { return localize('mls.taskFailedError.error', "Failed to complete task '{0}'. Error: {1}", taskName, err); }
-export const installDependenciesMsgTaskName = localize('mls.installDependencies.msgTaskName', "Installing Machine Learning extension dependencies");
+export const installPackageMngDependenciesMsgTaskName = localize('mls.installPackageMngDependencies.msgTaskName', "Installing package management dependencies");
+export const installModelMngDependenciesMsgTaskName = localize('mls.installModelMngDependencies.msgTaskName', "Installing model management dependencies");
 export const noResultError = localize('mls.noResultError', "No Result returned");
 export const requiredPackagesNotInstalled = localize('mls.requiredPackagesNotInstalled', "The required dependencies are not installed");
 export const confirmEnableExternalScripts = localize('mls.confirmEnableExternalScripts', "External script is required for package management. Are you sure you want to enable that.");

extensions/machine-learning-services/src/common/queryRunner.ts

@@ -7,22 +7,32 @@ import * as azdata from 'azdata';
 import * as nbExtensionApis from '../typings/notebookServices';
 import { ApiWrapper } from './apiWrapper';
 import * as constants from '../common/constants';
+import * as utils from '../common/utils';
 
-const maxNumberOfRetries = 3;
+const maxNumberOfRetries = 2;
 
 const listPythonPackagesQuery = `
+Declare @tablevar table(name NVARCHAR(MAX), version NVARCHAR(MAX))
+insert into @tablevar(name, version)
 EXEC sp_execute_external_script
 @language=N'Python',
 @script=N'import pkg_resources
 import pandas
 OutputDataSet = pandas.DataFrame([(d.project_name, d.version) for d in pkg_resources.working_set])'
+select e.name, version from sys.external_libraries e join @tablevar t on e.name = t.name
+where [language] = 'PYTHON'
 `;
 
 const listRPackagesQuery = `
+Declare @tablevar table(name NVARCHAR(MAX), version NVARCHAR(MAX))
+insert into @tablevar(name, version)
 EXEC sp_execute_external_script
 @language=N'R',
 @script=N'
 OutputDataSet <- as.data.frame(installed.packages()[,c(1,3)])'
+
+select e.name, version from sys.external_libraries e join @tablevar t on e.name = t.name
+where [language] = 'R'
 `;
 
 const checkMlInstalledQuery = `
@@ -63,24 +73,24 @@ export class QueryRunner {
 	 * Returns python packages installed in SQL server instance
 	 * @param connection SQL Connection
 	 */
-	public async getPythonPackages(connection: azdata.connection.ConnectionProfile): Promise<nbExtensionApis.IPackageDetails[]> {
-		return this.getPackages(connection, listPythonPackagesQuery);
+	public async getPythonPackages(connection: azdata.connection.ConnectionProfile, databaseName: string): Promise<nbExtensionApis.IPackageDetails[]> {
+		return this.getPackages(connection, databaseName, listPythonPackagesQuery);
 	}
 
 	/**
 	 * Returns python packages installed in SQL server instance
 	 * @param connection SQL Connection
 	 */
-	public async getRPackages(connection: azdata.connection.ConnectionProfile): Promise<nbExtensionApis.IPackageDetails[]> {
-		return this.getPackages(connection, listRPackagesQuery);
+	public async getRPackages(connection: azdata.connection.ConnectionProfile, databaseName: string): Promise<nbExtensionApis.IPackageDetails[]> {
+		return this.getPackages(connection, databaseName, listRPackagesQuery);
 	}
 
-	private async getPackages(connection: azdata.connection.ConnectionProfile, script: string): Promise<nbExtensionApis.IPackageDetails[]> {
+	private async getPackages(connection: azdata.connection.ConnectionProfile, databaseName: string, script: string): Promise<nbExtensionApis.IPackageDetails[]> {
 		let packages: nbExtensionApis.IPackageDetails[] = [];
 		let result: azdata.SimpleExecuteResult | undefined = undefined;
 
 		for (let index = 0; index < maxNumberOfRetries; index++) {
-			result = await this.runQuery(connection, script);
+			result = await this.runQuery(connection, utils.getScriptWithDBChange(connection.databaseName, databaseName, script));
 			if (result && result.rowCount > 0) {
 				break;
 			}

extensions/machine-learning-services/src/modelManagement/modelPythonClient.ts

@@ -40,7 +40,7 @@ export class ModelPythonClient {
 	 * Installs dependencies for python client
 	 */
 	private async installDependencies(): Promise<void> {
-		await utils.executeTasks(this._apiWrapper, constants.installDependenciesMsgTaskName, [
+		await utils.executeTasks(this._apiWrapper, constants.installModelMngDependenciesMsgTaskName, [
 			this._packageManager.installRequiredPythonPackages(this._config.modelsRequiredPythonPackages)], true);
 	}
 

extensions/machine-learning-services/src/packageManagement/packageManageProviderBase.ts

@@ -5,7 +5,6 @@
 
 import * as azdata from 'azdata';
 import { ApiWrapper } from '../common/apiWrapper';
-import * as constants from '../common/constants';
 import * as nbExtensionApis from '../typings/notebookServices';
 import * as utils from '../common/utils';
 
@@ -23,14 +22,17 @@ export abstract class SqlPackageManageProviderBase {
 	}
 
 	/**
-	 * Returns location title
+	 * Returns database names
 	 */
-	public async getLocationTitle(): Promise<string> {
+	public async getLocations(): Promise<nbExtensionApis.IPackageLocation[]> {
 		let connection = await this.getCurrentConnection();
 		if (connection) {
-			return `${connection.serverName} ${connection.databaseName ? connection.databaseName : ''}`;
+			let databases = await this._apiWrapper.listDatabases(connection.connectionId);
+			return databases.map(x => {
+				return { displayName: x, name: x };
+			});
 		}
-		return constants.noConnectionError;
+		return [];
 	}
 
 	protected async getCurrentConnection(): Promise<azdata.connection.ConnectionProfile> {
@@ -42,16 +44,16 @@ export abstract class SqlPackageManageProviderBase {
 	 * @param packages Packages to install
 	 * @param useMinVersion minimum version
 	 */
-	public async installPackages(packages: nbExtensionApis.IPackageDetails[], useMinVersion: boolean): Promise<void> {
+	public async installPackages(packages: nbExtensionApis.IPackageDetails[], useMinVersion: boolean, databaseName: string): Promise<void> {
 
 		if (packages) {
-			await Promise.all(packages.map(x => this.installPackage(x, useMinVersion)));
+			await Promise.all(packages.map(x => this.installPackage(x, useMinVersion, databaseName)));
 		}
 		//TODO: use useMinVersion
 		console.log(useMinVersion);
 	}
 
-	private async installPackage(packageDetail: nbExtensionApis.IPackageDetails, useMinVersion: boolean): Promise<void> {
+	private async installPackage(packageDetail: nbExtensionApis.IPackageDetails, useMinVersion: boolean, databaseName: string): Promise<void> {
 		if (useMinVersion) {
 			let packageOverview = await this.getPackageOverview(packageDetail.name);
 			if (packageOverview && packageOverview.versions) {
@@ -60,16 +62,16 @@ export abstract class SqlPackageManageProviderBase {
 			}
 		}
 
-		await this.executeScripts(ScriptMode.Install, packageDetail);
+		await this.executeScripts(ScriptMode.Install, packageDetail, databaseName);
 	}
 
 	/**
 	 * Uninstalls given packages
 	 * @param packages Packages to uninstall
 	 */
-	public async uninstallPackages(packages: nbExtensionApis.IPackageDetails[]): Promise<void> {
+	public async uninstallPackages(packages: nbExtensionApis.IPackageDetails[], databaseName: string): Promise<void> {
 		if (packages) {
-			await Promise.all(packages.map(x => this.executeScripts(ScriptMode.Uninstall, x)));
+			await Promise.all(packages.map(x => this.executeScripts(ScriptMode.Uninstall, x, databaseName)));
 		}
 	}
 
@@ -88,8 +90,8 @@ export abstract class SqlPackageManageProviderBase {
 	/**
 	 * Returns list of packages
 	 */
-	public async listPackages(): Promise<nbExtensionApis.IPackageDetails[]> {
-		let packages = await this.fetchPackages();
+	public async listPackages(databaseName: string): Promise<nbExtensionApis.IPackageDetails[]> {
+		let packages = await this.fetchPackages(databaseName);
 		if (packages) {
 			packages = packages.sort((a, b) => this.comparePackages(a, b));
 		} else {
@@ -110,6 +112,6 @@ export abstract class SqlPackageManageProviderBase {
 	}
 
 	protected abstract fetchPackage(packageName: string): Promise<nbExtensionApis.IPackageOverview>;
-	protected abstract fetchPackages(): Promise<nbExtensionApis.IPackageDetails[]>;
-	protected abstract executeScripts(scriptMode: ScriptMode, packageDetails: nbExtensionApis.IPackageDetails): Promise<void>;
+	protected abstract fetchPackages(databaseName: string): Promise<nbExtensionApis.IPackageDetails[]>;
+	protected abstract executeScripts(scriptMode: ScriptMode, packageDetails: nbExtensionApis.IPackageDetails, databaseName: string): Promise<void>;
 }

extensions/machine-learning-services/src/packageManagement/packageManagementService.ts

@@ -103,15 +103,15 @@ export class PackageManagementService {
 	 * Returns python packages installed in SQL server instance
 	 * @param connection SQL Connection
 	 */
-	public async getPythonPackages(connection: azdata.connection.ConnectionProfile): Promise<nbExtensionApis.IPackageDetails[]> {
-		return this._queryRunner.getPythonPackages(connection);
+	public async getPythonPackages(connection: azdata.connection.ConnectionProfile, databaseName: string): Promise<nbExtensionApis.IPackageDetails[]> {
+		return this._queryRunner.getPythonPackages(connection, databaseName);
 	}
 
 	/**
 	 * Returns python packages installed in SQL server instance
 	 * @param connection SQL Connection
 	 */
-	public async getRPackages(connection: azdata.connection.ConnectionProfile): Promise<nbExtensionApis.IPackageDetails[]> {
-		return this._queryRunner.getRPackages(connection);
+	public async getRPackages(connection: azdata.connection.ConnectionProfile, databaseName: string): Promise<nbExtensionApis.IPackageDetails[]> {
+		return this._queryRunner.getRPackages(connection, databaseName);
 	}
 }

extensions/machine-learning-services/src/packageManagement/packageManager.ts

@@ -93,7 +93,6 @@ export class PackageManager {
 				// Execute the command
 				//
 				this._apiWrapper.executeCommand(constants.managePackagesCommand, {
-					multiLocations: false,
 					defaultLocation: defaultProvider.packageTarget.location,
 					defaultProviderId: defaultProvider.providerId
 				});
@@ -116,7 +115,7 @@ export class PackageManager {
 	 * Installs dependencies for the extension
 	 */
 	public async installDependencies(): Promise<void> {
-		await utils.executeTasks(this._apiWrapper, constants.installDependenciesMsgTaskName, [
+		await utils.executeTasks(this._apiWrapper, constants.installPackageMngDependenciesMsgTaskName, [
 			this.installRequiredPythonPackages(this._config.requiredSqlPythonPackages),
 			this.installRequiredRPackages()], true);
 	}
@@ -130,7 +129,7 @@ export class PackageManager {
 		}
 
 		await utils.createFolder(utils.getRPackagesFolderPath(this._rootFolder));
-		await Promise.all(this._config.requiredSqlPythonPackages.map(x => this.installRPackage(x)));
+		await Promise.all(this._config.requiredSqlRPackages.map(x => this.installRPackage(x)));
 	}
 
 	/**
@@ -151,7 +150,8 @@ export class PackageManager {
 		let fileContent = '';
 		requiredPackages.forEach(packageDetails => {
 			let hasVersion = ('version' in packageDetails) && !isNullOrUndefined(packageDetails['version']) && packageDetails['version'].length > 0;
-			if (!installedPackages.find(x => x.name === packageDetails['name'] && (!hasVersion || packageDetails['version'] === x.version))) {
+			if (!installedPackages.find(x => x.name === packageDetails['name']
+				&& (!hasVersion || utils.comparePackageVersions(packageDetails['version'] || '', x.version) <= 0))) {
 				let packageNameDetail = hasVersion ? `${packageDetails.name}==${packageDetails.version}` : `${packageDetails.name}`;
 				fileContent = `${fileContent}${packageNameDetail}\n`;
 			}
@@ -177,7 +177,7 @@ export class PackageManager {
 	private async getInstalledPipPackages(): Promise<nbExtensionApis.IPackageDetails[]> {
 		try {
 			let cmd = `"${this.pythonExecutable}" -m pip list --format=json`;
-			let packagesInfo = await this._processService.executeBufferedCommand(cmd, this._outputChannel);
+			let packagesInfo = await this._processService.executeBufferedCommand(cmd, undefined);
 			let packagesResult: nbExtensionApis.IPackageDetails[] = [];
 			if (packagesInfo) {
 				packagesResult = <nbExtensionApis.IPackageDetails[]>JSON.parse(packagesInfo);

extensions/machine-learning-services/src/packageManagement/sqlPythonPackageManageProvider.ts

@@ -9,7 +9,7 @@ import * as nbExtensionApis from '../typings/notebookServices';
 import { ApiWrapper } from '../common/apiWrapper';
 import { ProcessService } from '../common/processService';
 import { Config } from '../configurations/config';
-import { SqlPackageManageProviderBase, ScriptMode } from './SqlPackageManageProviderBase';
+import { SqlPackageManageProviderBase, ScriptMode } from './packageManageProviderBase';
 import { HttpClient } from '../common/httpClient';
 import * as utils from '../common/utils';
 import { PackageManagementService } from './packageManagementService';
@@ -50,8 +50,8 @@ export class SqlPythonPackageManageProvider extends SqlPackageManageProviderBase
 	/**
 	 * Returns list of packages
 	 */
-	protected async fetchPackages(): Promise<nbExtensionApis.IPackageDetails[]> {
-		return await this._service.getPythonPackages(await this.getCurrentConnection());
+	protected async fetchPackages(databaseName: string): Promise<nbExtensionApis.IPackageDetails[]> {
+		return await this._service.getPythonPackages(await this.getCurrentConnection(), databaseName);
 	}
 
 	/**
@@ -59,14 +59,14 @@ export class SqlPythonPackageManageProvider extends SqlPackageManageProviderBase
 	 * @param packageDetails Packages to install or uninstall
 	 * @param scriptMode can be 'install' or 'uninstall'
 	 */
-	protected async executeScripts(scriptMode: ScriptMode, packageDetails: nbExtensionApis.IPackageDetails): Promise<void> {
+	protected async executeScripts(scriptMode: ScriptMode, packageDetails: nbExtensionApis.IPackageDetails, databaseName: string): Promise<void> {
 		let connection = await this.getCurrentConnection();
 		let credentials = await this._apiWrapper.getCredentials(connection.connectionId);
 
 		if (connection) {
 			let port = '1433';
 			let server = connection.serverName;
-			let database = connection.databaseName ? `, database="${connection.databaseName}"` : '';
+			let database = databaseName ? `, database="${databaseName}"` : '';
 			let index = connection.serverName.indexOf(',');
 			if (index > 0) {
 				port = connection.serverName.substring(index + 1);

extensions/machine-learning-services/src/packageManagement/sqlRPackageManageProvider.ts

@@ -10,7 +10,7 @@ import * as nbExtensionApis from '../typings/notebookServices';
 import { ApiWrapper } from '../common/apiWrapper';
 import { ProcessService } from '../common/processService';
 import { Config } from '../configurations/config';
-import { SqlPackageManageProviderBase, ScriptMode } from './SqlPackageManageProviderBase';
+import { SqlPackageManageProviderBase, ScriptMode } from './packageManageProviderBase';
 import { HttpClient } from '../common/httpClient';
 import * as constants from '../common/constants';
 import { PackageManagementService } from './packageManagementService';
@@ -54,8 +54,8 @@ export class SqlRPackageManageProvider extends SqlPackageManageProviderBase impl
 	/**
 	 * Returns list of packages
 	 */
-	protected async fetchPackages(): Promise<nbExtensionApis.IPackageDetails[]> {
-		return await this._service.getRPackages(await this.getCurrentConnection());
+	protected async fetchPackages(databaseName: string): Promise<nbExtensionApis.IPackageDetails[]> {
+		return await this._service.getRPackages(await this.getCurrentConnection(), databaseName);
 	}
 
 	/**
@@ -63,12 +63,12 @@ export class SqlRPackageManageProvider extends SqlPackageManageProviderBase impl
 	 * @param packageDetails Packages to install or uninstall
 	 * @param scriptMode can be 'install' or 'uninstall'
 	 */
-	protected async executeScripts(scriptMode: ScriptMode, packageDetails: nbExtensionApis.IPackageDetails): Promise<void> {
+	protected async executeScripts(scriptMode: ScriptMode, packageDetails: nbExtensionApis.IPackageDetails, databaseName: string): Promise<void> {
 		let connection = await this.getCurrentConnection();
 		let credentials = await this._apiWrapper.getCredentials(connection.connectionId);
 
 		if (connection) {
-			let database = connection.databaseName ? `, database="${connection.databaseName}"` : '';
+			let database = databaseName ? `, database="${databaseName}"` : '';
 			let connectionParts = `server="${connection.serverName}", uid="${connection.userName}", pwd="${credentials[azdata.ConnectionOptionSpecialType.password]}"${database}`;
 			let rCommandScript = scriptMode === ScriptMode.Install ? 'sql_install.packages' : 'sql_remove.packages';
 

extensions/machine-learning-services/src/test/packageManagement/packageManager.test.ts

@@ -116,7 +116,7 @@ describe('Package Manager', () => {
 
 	it('installDependencies Should install packages that are not already installed', async function (): Promise<void> {
 		let testContext = createContext();
-		//let packagesInstalled = false;
+		let packagesInstalled = false;
 		let installedPackages = `[
 			{"name":"pymssql","version":"2.1.4"}
 		]`;
@@ -128,15 +128,67 @@ describe('Package Manager', () => {
 		});
 		testContext.processService.setup(x => x.executeBufferedCommand(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns((command) => {
 			if (command.indexOf('pip install') > 0) {
-				//packagesInstalled = true;
+				packagesInstalled = true;
 			}
 			return Promise.resolve(installedPackages);
 		});
 
 		let packageManager = createPackageManager(testContext);
 		await packageManager.installDependencies();
-		//should.equal(testContext.getOpStatus(), azdata.TaskStatus.Succeeded);
-		//should.equal(packagesInstalled, true);
+		should.equal(testContext.getOpStatus(), azdata.TaskStatus.Succeeded);
+		should.equal(packagesInstalled, true);
+	});
+
+	it('installDependencies Should not install packages if runtime is disabled in setting', async function (): Promise<void> {
+		let testContext = createContext();
+		testContext.config.setup(x => x.rEnabled).returns(() => false);
+		testContext.config.setup(x => x.pythonEnabled).returns(() => false);
+		let packagesInstalled = false;
+		let installedPackages = `[
+			{"name":"pymssql","version":"2.1.4"}
+		]`;
+		testContext.apiWrapper.setup(x => x.showQuickPick(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve({
+			label: 'Yes'
+		}));
+		testContext.apiWrapper.setup(x => x.startBackgroundOperation(TypeMoq.It.isAny())).returns((operationInfo: azdata.BackgroundOperationInfo) => {
+			operationInfo.operation(testContext.op);
+		});
+		testContext.processService.setup(x => x.executeBufferedCommand(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns((command) => {
+			if (command.indexOf('pip install') > 0 || command.indexOf('install.packages') > 0) {
+				packagesInstalled = true;
+			}
+			return Promise.resolve(installedPackages);
+		});
+
+		let packageManager = createPackageManager(testContext);
+		await packageManager.installDependencies();
+		should.equal(testContext.getOpStatus(), azdata.TaskStatus.Succeeded);
+		should.equal(packagesInstalled, false);
+	});
+
+	it('installDependencies Should install packages that have older version installed', async function (): Promise<void> {
+		let testContext = createContext();
+		let packagesInstalled = false;
+		let installedPackages = `[
+			{"name":"sqlmlutils","version":"0.1.1"}
+		]`;
+		testContext.apiWrapper.setup(x => x.showQuickPick(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve({
+			label: 'Yes'
+		}));
+		testContext.apiWrapper.setup(x => x.startBackgroundOperation(TypeMoq.It.isAny())).returns((operationInfo: azdata.BackgroundOperationInfo) => {
+			operationInfo.operation(testContext.op);
+		});
+		testContext.processService.setup(x => x.executeBufferedCommand(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns((command) => {
+			if (command.indexOf('pip install') > 0) {
+				packagesInstalled = true;
+			}
+			return Promise.resolve(installedPackages);
+		});
+
+		let packageManager = createPackageManager(testContext);
+		await packageManager.installDependencies();
+		should.equal(testContext.getOpStatus(), azdata.TaskStatus.Succeeded);
+		should.equal(packagesInstalled, true);
 	});
 
 	it('installDependencies Should install packages if list packages fails', async function (): Promise<void> {
@@ -197,7 +249,7 @@ describe('Package Manager', () => {
 			{ name: 'pymssql', version: '2.1.4' },
 			{ name: 'sqlmlutils', version: '' }
 		]);
-		testContext.config.setup(x => x.requiredSqlPythonPackages).returns( () => [
+		testContext.config.setup(x => x.requiredSqlRPackages).returns( () => [
 			{ name: 'RODBCext', repository: 'https://cran.microsoft.com' },
 			{ name: 'sqlmlutils', fileName: 'sqlmlutils_0.7.1.zip', downloadUrl: 'https://github.com/microsoft/sqlmlutils/blob/master/R/dist/sqlmlutils_0.7.1.zip?raw=true'}
 		]);

extensions/machine-learning-services/src/test/packageManagement/sqlPythonPackageManageProvider.test.ts

@@ -7,7 +7,6 @@ import * as azdata from 'azdata';
 import * as should from 'should';
 import 'mocha';
 import * as TypeMoq from 'typemoq';
-import * as constants from '../../common/constants';
 import { SqlPythonPackageManageProvider } from '../../packageManagement/sqlPythonPackageManageProvider';
 import { createContext, TestContext } from './utils';
 import * as nbExtensionApis from '../../typings/notebookServices';
@@ -40,10 +39,10 @@ describe('SQL Python Package Manager', () => {
 
 		let connection = new azdata.connection.ConnectionProfile();
 		testContext.apiWrapper.setup(x => x.getCurrentConnection()).returns(() => { return Promise.resolve(connection); });
-		testContext.serverConfigManager.setup(x => x.getPythonPackages(TypeMoq.It.isAny())).returns(() => Promise.resolve(packages));
+		testContext.serverConfigManager.setup(x => x.getPythonPackages(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(packages));
 
 		let provider = createProvider(testContext);
-		let actual = await provider.listPackages();
+		let actual = await provider.listPackages(connection.databaseName);
 		let expected = [
 			{
 				'name': 'a-name',
@@ -72,10 +71,10 @@ describe('SQL Python Package Manager', () => {
 
 		let connection = new azdata.connection.ConnectionProfile();
 		testContext.apiWrapper.setup(x => x.getCurrentConnection()).returns(() => { return Promise.resolve(connection); });
-		testContext.serverConfigManager.setup(x => x.getPythonPackages(TypeMoq.It.isAny())).returns(() => Promise.resolve(packages));
+		testContext.serverConfigManager.setup(x => x.getPythonPackages(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(packages));
 
 		let provider = createProvider(testContext);
-		let actual = await provider.listPackages();
+		let actual = await provider.listPackages(connection.databaseName);
 		let expected = [
 			{
 				'name': 'b-name',
@@ -95,10 +94,10 @@ describe('SQL Python Package Manager', () => {
 		let connection = new azdata.connection.ConnectionProfile();
 		let packages: nbExtensionApis.IPackageDetails[];
 		testContext.apiWrapper.setup(x => x.getCurrentConnection()).returns(() => { return Promise.resolve(connection); });
-		testContext.serverConfigManager.setup(x => x.getPythonPackages(TypeMoq.It.isAny())).returns(() => Promise.resolve(packages));
+		testContext.serverConfigManager.setup(x => x.getPythonPackages(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(packages));
 
 		let provider = createProvider(testContext);
-		let actual = await provider.listPackages();
+		let actual = await provider.listPackages(connection.databaseName);
 		let expected: nbExtensionApis.IPackageDetails[] = [];
 		should.deepEqual(actual, expected);
 	});
@@ -108,10 +107,10 @@ describe('SQL Python Package Manager', () => {
 
 		let connection = new azdata.connection.ConnectionProfile();
 		testContext.apiWrapper.setup(x => x.getCurrentConnection()).returns(() => { return Promise.resolve(connection); });
-		testContext.serverConfigManager.setup(x => x.getPythonPackages(TypeMoq.It.isAny())).returns(() => Promise.resolve([]));
+		testContext.serverConfigManager.setup(x => x.getPythonPackages(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve([]));
 
 		let provider = createProvider(testContext);
-		let actual = await provider.listPackages();
+		let actual = await provider.listPackages(connection.databaseName);
 		let expected: nbExtensionApis.IPackageDetails[] = [];
 		should.deepEqual(actual, expected);
 	});
@@ -152,7 +151,7 @@ describe('SQL Python Package Manager', () => {
 		});
 
 		let provider = createProvider(testContext);
-		await provider.installPackages(packages, false);
+		await provider.installPackages(packages, false, connection.databaseName);
 
 		should.deepEqual(packagesUpdated, true);
 	});
@@ -192,7 +191,7 @@ describe('SQL Python Package Manager', () => {
 		});
 
 		let provider = createProvider(testContext);
-		await provider.uninstallPackages(packages);
+		await provider.uninstallPackages(packages, connection.databaseName);
 
 		should.deepEqual(packagesUpdated, true);
 	});
@@ -233,7 +232,7 @@ describe('SQL Python Package Manager', () => {
 		});
 
 		let provider = createProvider(testContext);
-		await provider.installPackages(packages, false);
+		await provider.installPackages(packages, false, connection.databaseName);
 
 		should.deepEqual(packagesUpdated, true);
 	});
@@ -255,7 +254,7 @@ describe('SQL Python Package Manager', () => {
 
 
 		let provider = createProvider(testContext);
-		await provider.installPackages(packages, false);
+		await provider.installPackages(packages, false, connection.databaseName);
 
 		should.deepEqual(packagesUpdated, false);
 	});
@@ -277,7 +276,7 @@ describe('SQL Python Package Manager', () => {
 
 
 		let provider = createProvider(testContext);
-		await provider.uninstallPackages(packages);
+		await provider.uninstallPackages(packages, connection.databaseName);
 
 		should.deepEqual(packagesUpdated, false);
 	});
@@ -346,42 +345,44 @@ describe('SQL Python Package Manager', () => {
 		should.deepEqual(actual, packagePreview);
 	});
 
-	it('getLocationTitle Should default string for no connection', async function (): Promise<void> {
+	it('getLocations Should return empty array for no connection', async function (): Promise<void> {
 		let testContext = createContext();
 		let connection: azdata.connection.ConnectionProfile;
 		testContext.apiWrapper.setup(x => x.getCurrentConnection()).returns(() => { return Promise.resolve(connection); });
 
 		let provider = createProvider(testContext);
-		let actual = await provider.getLocationTitle();
+		let actual = await provider.getLocations();
 
-		should.deepEqual(actual, constants.noConnectionError);
+		should.deepEqual(actual, []);
 	});
 
-	it('getLocationTitle Should return connection title string for valid connection', async function (): Promise<void> {
+	it('getLocations Should return database names for valid connection', async function (): Promise<void> {
 		let testContext = createContext();
 
 		let connection = new azdata.connection.ConnectionProfile();
 		connection.serverName = 'serverName';
 		connection.databaseName = 'databaseName';
+		const databaseNames = [
+			'db1',
+			'db2'
+		];
+		const expected = [
+			{
+				displayName: 'db1',
+				name: 'db1'
+			},
+			{
+				displayName: 'db2',
+				name: 'db2'
+			}
+		];
 		testContext.apiWrapper.setup(x => x.getCurrentConnection()).returns(() => { return Promise.resolve(connection); });
+		testContext.apiWrapper.setup(x => x.listDatabases(connection.connectionId)).returns(() => { return Promise.resolve(databaseNames); });
 
 		let provider = createProvider(testContext);
-		let actual = await provider.getLocationTitle();
+		let actual = await provider.getLocations();
 
-		should.deepEqual(actual, `${connection.serverName} ${connection.databaseName}`);
-	});
-
-	it('getLocationTitle Should return server name as connection title if there is not database name', async function (): Promise<void> {
-		let testContext = createContext();
-
-		let connection = new azdata.connection.ConnectionProfile();
-		connection.serverName = 'serverName';
-		testContext.apiWrapper.setup(x => x.getCurrentConnection()).returns(() => { return Promise.resolve(connection); });
-
-		let provider = createProvider(testContext);
-		let actual = await provider.getLocationTitle();
-
-		should.deepEqual(actual, `${connection.serverName} `);
+		should.deepEqual(actual, expected);
 	});
 
 	function createProvider(testContext: TestContext): SqlPythonPackageManageProvider {

extensions/machine-learning-services/src/test/packageManagement/sqlRPackageManageProvider.test.ts

@@ -7,7 +7,6 @@ import * as azdata from 'azdata';
 import * as should from 'should';
 import 'mocha';
 import * as TypeMoq from 'typemoq';
-import * as constants from '../../common/constants';
 import { SqlRPackageManageProvider } from '../../packageManagement/sqlRPackageManageProvider';
 import { createContext, TestContext } from './utils';
 import * as nbExtensionApis from '../../typings/notebookServices';
@@ -40,10 +39,10 @@ describe('SQL R Package Manager', () => {
 
 		let connection = new azdata.connection.ConnectionProfile();
 		testContext.apiWrapper.setup(x => x.getCurrentConnection()).returns(() => { return Promise.resolve(connection); });
-		testContext.serverConfigManager.setup(x => x.getRPackages(TypeMoq.It.isAny())).returns(() => Promise.resolve(packages));
+		testContext.serverConfigManager.setup(x => x.getRPackages(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(packages));
 
 		let provider = createProvider(testContext);
-		let actual = await provider.listPackages();
+		let actual = await provider.listPackages(connection.databaseName);
 		let expected = [
 			{
 				'name': 'a-name',
@@ -63,10 +62,10 @@ describe('SQL R Package Manager', () => {
 		let connection = new azdata.connection.ConnectionProfile();
 		let packages: nbExtensionApis.IPackageDetails[];
 		testContext.apiWrapper.setup(x => x.getCurrentConnection()).returns(() => { return Promise.resolve(connection); });
-		testContext.serverConfigManager.setup(x => x.getRPackages(TypeMoq.It.isAny())).returns(() => Promise.resolve(packages));
+		testContext.serverConfigManager.setup(x => x.getRPackages(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(packages));
 
 		let provider = createProvider(testContext);
-		let actual = await provider.listPackages();
+		let actual = await provider.listPackages(connection.databaseName);
 		let expected: nbExtensionApis.IPackageDetails[] = [];
 		should.deepEqual(actual, expected);
 	});
@@ -76,10 +75,10 @@ describe('SQL R Package Manager', () => {
 
 		let connection = new azdata.connection.ConnectionProfile();
 		testContext.apiWrapper.setup(x => x.getCurrentConnection()).returns(() => { return Promise.resolve(connection); });
-		testContext.serverConfigManager.setup(x => x.getRPackages(TypeMoq.It.isAny())).returns(() => Promise.resolve([]));
+		testContext.serverConfigManager.setup(x => x.getRPackages(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve([]));
 
 		let provider = createProvider(testContext);
-		let actual = await provider.listPackages();
+		let actual = await provider.listPackages(connection.databaseName);
 		let expected: nbExtensionApis.IPackageDetails[] = [];
 		should.deepEqual(actual, expected);
 	});
@@ -118,7 +117,7 @@ describe('SQL R Package Manager', () => {
 		});
 
 		let provider = createProvider(testContext);
-		await provider.installPackages(packages, false);
+		await provider.installPackages(packages, false, connection.databaseName);
 
 		should.deepEqual(packagesUpdated, true);
 	});
@@ -157,7 +156,7 @@ describe('SQL R Package Manager', () => {
 		});
 
 		let provider = createProvider(testContext);
-		await provider.uninstallPackages(packages);
+		await provider.uninstallPackages(packages, connection.databaseName);
 
 		should.deepEqual(packagesUpdated, true);
 	});
@@ -179,7 +178,7 @@ describe('SQL R Package Manager', () => {
 
 
 		let provider = createProvider(testContext);
-		await provider.installPackages(packages, false);
+		await provider.installPackages(packages, false, connection.databaseName);
 
 		should.deepEqual(packagesUpdated, false);
 	});
@@ -201,7 +200,7 @@ describe('SQL R Package Manager', () => {
 
 
 		let provider = createProvider(testContext);
-		await provider.uninstallPackages(packages);
+		await provider.uninstallPackages(packages, connection.databaseName);
 
 		should.deepEqual(packagesUpdated, false);
 	});
@@ -271,42 +270,44 @@ describe('SQL R Package Manager', () => {
 		should.deepEqual(actual, packagePreview);
 	});
 
-	it('getLocationTitle Should default string for no connection', async function (): Promise<void> {
+	it('getLocations Should return empty array for no connection', async function (): Promise<void> {
 		let testContext = createContext();
 		let connection: azdata.connection.ConnectionProfile;
 		testContext.apiWrapper.setup(x => x.getCurrentConnection()).returns(() => { return Promise.resolve(connection); });
 
 		let provider = createProvider(testContext);
-		let actual = await provider.getLocationTitle();
+		let actual = await provider.getLocations();
 
-		should.deepEqual(actual, constants.noConnectionError);
+		should.deepEqual(actual, []);
 	});
 
-	it('getLocationTitle Should return connection title string for valid connection', async function (): Promise<void> {
+	it('getLocations Should return database names for valid connection', async function (): Promise<void> {
 		let testContext = createContext();
 
 		let connection = new azdata.connection.ConnectionProfile();
 		connection.serverName = 'serverName';
 		connection.databaseName = 'databaseName';
+		const databaseNames = [
+			'db1',
+			'db2'
+		];
+		const expected = [
+			{
+				displayName: 'db1',
+				name: 'db1'
+			},
+			{
+				displayName: 'db2',
+				name: 'db2'
+			}
+		];
 		testContext.apiWrapper.setup(x => x.getCurrentConnection()).returns(() => { return Promise.resolve(connection); });
+		testContext.apiWrapper.setup(x => x.listDatabases(connection.connectionId)).returns(() => { return Promise.resolve(databaseNames); });
 
 		let provider = createProvider(testContext);
-		let actual = await provider.getLocationTitle();
+		let actual = await provider.getLocations();
 
-		should.deepEqual(actual, `${connection.serverName} ${connection.databaseName}`);
-	});
-
-	it('getLocationTitle Should return server name as connection title if there is not database name', async function (): Promise<void> {
-		let testContext = createContext();
-
-		let connection = new azdata.connection.ConnectionProfile();
-		connection.serverName = 'serverName';
-		testContext.apiWrapper.setup(x => x.getCurrentConnection()).returns(() => { return Promise.resolve(connection); });
-
-		let provider = createProvider(testContext);
-		let actual = await provider.getLocationTitle();
-
-		should.deepEqual(actual, `${connection.serverName} `);
+		should.deepEqual(actual, expected);
 	});
 
 	function createProvider(testContext: TestContext): SqlRPackageManageProvider {

extensions/machine-learning-services/src/test/queryRunner.test.ts

@@ -59,7 +59,7 @@ describe('Query Runner', () => {
 		let queryProvider: azdata.QueryProvider;
 		testContext.apiWrapper.setup(x => x.getProvider<azdata.QueryProvider>(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => queryProvider);
 
-		let actual = await queryRunner.getPythonPackages(connection);
+		let actual = await queryRunner.getPythonPackages(connection, connection.databaseName);
 		should.deepEqual(actual, []);
 	});
 
@@ -70,7 +70,7 @@ describe('Query Runner', () => {
 		testContext.queryProvider.runQueryAndReturn = () => { return Promise.reject(); };
 		testContext.apiWrapper.setup(x => x.getProvider<azdata.QueryProvider>(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => testContext.queryProvider);
 
-		let actual = await queryRunner.getPythonPackages(connection);
+		let actual = await queryRunner.getPythonPackages(connection, connection.databaseName);
 		should.deepEqual(actual, []);
 	});
 
@@ -117,7 +117,7 @@ describe('Query Runner', () => {
 		testContext.queryProvider.runQueryAndReturn = () => { return Promise.resolve(result); };
 		testContext.apiWrapper.setup(x => x.getProvider<azdata.QueryProvider>(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => testContext.queryProvider);
 
-		let actual = await queryRunner.getPythonPackages(connection);
+		let actual = await queryRunner.getPythonPackages(connection, connection.databaseName);
 
 		should.deepEqual(actual, expected);
 	});
@@ -138,7 +138,7 @@ describe('Query Runner', () => {
 		testContext.queryProvider.runQueryAndReturn = () => { return Promise.resolve(result); };
 		testContext.apiWrapper.setup(x => x.getProvider<azdata.QueryProvider>(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => testContext.queryProvider);
 
-		let actual = await queryRunner.getPythonPackages(connection);
+		let actual = await queryRunner.getPythonPackages(connection, connection.databaseName);
 
 		should.deepEqual(actual, expected);
 	});

extensions/machine-learning-services/src/typings/notebookServices.d.ts

@@ -51,13 +51,56 @@ export interface IPackageOverview {
 	summary: string;
 }
 
-export interface IPackageManageProvider {
-	providerId: string;
-	packageTarget: IPackageTarget;
-	listPackages(): Promise<IPackageDetails[]>
-	installPackages(package: IPackageDetails[], useMinVersion: boolean): Promise<void>;
-	uninstallPackages(package: IPackageDetails[]): Promise<void>;
-	canUseProvider(): Promise<boolean>;
-	getLocationTitle(): Promise<string>;
-	getPackageOverview(packageName: string): Promise<IPackageOverview>
+export interface IPackageLocation {
+	name: string;
+	displayName: string;
+}
+
+/**
+ * Package manage provider interface
+ */
+export interface IPackageManageProvider {
+	/**
+	 * Provider id
+	 */
+	providerId: string;
+
+	/**
+	 * package target
+	 */
+	packageTarget: IPackageTarget;
+
+	/**
+	 * Returns list of installed packages
+	 */
+	listPackages(location?: string): Promise<IPackageDetails[]>;
+
+	/**
+	 * Installs give packages
+	 * @param package Packages to install
+	 * @param useMinVersion if true, minimal version will be used
+	 */
+	installPackages(package: IPackageDetails[], useMinVersion: boolean, location?: string): Promise<void>;
+
+	/**
+	 * Uninstalls given packages
+	 * @param package package to uninstall
+	 */
+	uninstallPackages(package: IPackageDetails[], location?: string): Promise<void>;
+
+	/**
+	 * Returns true if the provider can be used in current context
+	 */
+	canUseProvider(): Promise<boolean>;
+
+	/**
+	 * Returns location title
+	 */
+	getLocations(): Promise<IPackageLocation[]>;
+
+	/**
+	 * Returns Package Overview
+	 * @param packageName package name
+	 */
+	getPackageOverview(packageName: string): Promise<IPackageOverview>;
 }