From a3efb193e00b346dd63c20007c8670ab75086e95 Mon Sep 17 00:00:00 2001
From: Vasu Bhog <vabhog@microsoft.com>
Date: Wed, 27 Apr 2022 15:14:39 -0700
Subject: [PATCH] Fix connection profile prompt for Create Function with SQL
 Binding (#19211)

* fix no connection profile and progress report

* show connection profile prompt if user exits object

* address comments
---
 .../sql-bindings/src/common/constants.ts      |  1 +
 .../src/services/azureFunctionsService.ts     | 48 ++++++++++++++++++-
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/extensions/sql-bindings/src/common/constants.ts b/extensions/sql-bindings/src/common/constants.ts
index af737dde9a..b461c6f706 100644
--- a/extensions/sql-bindings/src/common/constants.ts
+++ b/extensions/sql-bindings/src/common/constants.ts
@@ -87,3 +87,4 @@ export const userPasswordLater = localize('userPasswordLater', 'In order to user
 export const openFile = localize('openFile', "Open File");
 export const closeButton = localize('closeButton', "Close");
 export function addSqlBinding(functionName: string): string { return localize('addSqlBinding', 'Adding SQL Binding to function "{0}"...'), functionName; }
+export const connectionProgressTitle = localize('connectionProgressTitle', "Testing SQL Server connection...");
diff --git a/extensions/sql-bindings/src/services/azureFunctionsService.ts b/extensions/sql-bindings/src/services/azureFunctionsService.ts
index 46adf8a0fa..958b0ae763 100644
--- a/extensions/sql-bindings/src/services/azureFunctionsService.ts
+++ b/extensions/sql-bindings/src/services/azureFunctionsService.ts
@@ -106,7 +106,6 @@ export async function createAzureFunction(node?: ITreeNodeInfo): Promise<void> {
 		TelemetryReporter.createActionEvent(TelemetryViews.CreateAzureFunctionWithSqlBinding, telemetryStep)
 			.withAdditionalProperties(propertyBag).send();
 
-
 		// Get connection string parameters and construct object name from prompt or connectionInfo given
 		let objectName: string | undefined;
 		const vscodeMssqlApi = await utils.getVscodeMssqlApi();
@@ -115,6 +114,53 @@ export async function createAzureFunction(node?: ITreeNodeInfo): Promise<void> {
 			telemetryStep = CreateAzureFunctionStep.launchFromCommandPalette;
 
 			// prompt user for connection profile to get connection info
+			while (true) {
+				connectionInfo = await vscodeMssqlApi.promptForConnection(true);
+				if (!connectionInfo) {
+					// User cancelled
+					return;
+				}
+				telemetryStep = 'getConnectionInfo';
+				let connectionURI: string = '';
+				try {
+					await vscode.window.withProgress(
+						{
+							location: vscode.ProgressLocation.Notification,
+							title: constants.connectionProgressTitle,
+							cancellable: false
+						}, async (_progress, _token) => {
+							// list databases based on connection profile selected
+							connectionURI = await vscodeMssqlApi.connect(connectionInfo!);
+						}
+					);
+				} catch (e) {
+					// mssql connection error will be shown to the user
+					// we will then prompt user to choose a connection profile again
+					continue;
+				}
+				// list databases based on connection profile selected
+				let listDatabases = await vscodeMssqlApi.listDatabases(connectionURI);
+				const selectedDatabase = (await vscode.window.showQuickPick(listDatabases, {
+					canPickMany: false,
+					title: constants.selectDatabase,
+					ignoreFocusOut: true
+				}));
+
+				if (!selectedDatabase) {
+					// User cancelled
+					// we will then prompt user to choose a connection profile again
+					continue;
+				}
+				connectionInfo.database = selectedDatabase;
+
+				// prompt user for object name to create function from
+				objectName = await azureFunctionsUtils.promptForObjectName(selectedBinding);
+				if (!objectName) {
+					// user cancelled
+					return;
+				}
+				break;
+			}
 			telemetryStep = CreateAzureFunctionStep.getConnectionProfile;
 			connectionInfo = await vscodeMssqlApi.promptForConnection(true);
 			if (!connectionInfo) {