From 5c5b73525de3e009e161d27ef63a7771eccf7bd6 Mon Sep 17 00:00:00 2001
From: Cory Rivera <corivera@microsoft.com>
Date: Fri, 28 Jul 2023 14:26:33 -0700
Subject: [PATCH] Connect query editor when producing script from object
 management dialogs. (#24015)

---
 .../ui/objectManagementDialogBase.ts          |  1 -
 .../mssql/src/ui/scriptableDialogBase.ts      | 26 ++++++++++++++++++-
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/extensions/mssql/src/objectManagement/ui/objectManagementDialogBase.ts b/extensions/mssql/src/objectManagement/ui/objectManagementDialogBase.ts
index 904bcd9aeb..a28e00f383 100644
--- a/extensions/mssql/src/objectManagement/ui/objectManagementDialogBase.ts
+++ b/extensions/mssql/src/objectManagement/ui/objectManagementDialogBase.ts
@@ -26,7 +26,6 @@ export interface ObjectManagementDialogOptions extends ScriptableDialogOptions {
 	isNewObject: boolean;
 	parentUrn: string;
 	objectUrn?: string;
-	objectExplorerContext?: azdata.ObjectExplorerContext;
 	objectName?: string;
 }
 
diff --git a/extensions/mssql/src/ui/scriptableDialogBase.ts b/extensions/mssql/src/ui/scriptableDialogBase.ts
index e49db066e2..ef94c408c2 100644
--- a/extensions/mssql/src/ui/scriptableDialogBase.ts
+++ b/extensions/mssql/src/ui/scriptableDialogBase.ts
@@ -15,6 +15,10 @@ export interface ScriptableDialogOptions {
 	 * The width of the dialog, defaults to narrow if not set
 	 */
 	width?: azdata.window.DialogWidth;
+	/**
+	 * The object explorer context in which this dialog was opened.
+	 */
+	objectExplorerContext?: azdata.ObjectExplorerContext;
 }
 
 /**
@@ -90,7 +94,27 @@ export abstract class ScriptableDialogBase<OptionsType extends ScriptableDialogO
 			const script = await this.generateScript();
 			if (script) {
 				message = localizedConstants.ScriptGeneratedText;
-				await azdata.queryeditor.openQueryDocument({ content: script }, providerId);
+				let doc = await azdata.queryeditor.openQueryDocument({ content: script }, providerId);
+				if (this.options.objectExplorerContext?.connectionProfile) {
+					let profile = this.options.objectExplorerContext?.connectionProfile;
+					let convertedProfile: azdata.connection.ConnectionProfile = {
+						providerId: profile.providerName,
+						connectionId: profile.id,
+						connectionName: profile.connectionName,
+						serverName: profile.serverName,
+						databaseName: profile.databaseName,
+						userName: profile.userName,
+						password: profile.password,
+						authenticationType: profile.authenticationType,
+						savePassword: profile.savePassword,
+						groupFullName: profile.groupFullName,
+						groupId: profile.groupId,
+						saveProfile: profile.savePassword,
+						azureTenantId: profile.azureTenantId,
+						options: profile.options
+					};
+					await doc.connect(convertedProfile);
+				}
 			} else {
 				message = localizedConstants.NoActionScriptedMessage;
 			}