Skip starting a separate task for backup and restore operations (#25285)

microsoft · Feb 2, 2024 · 4269fe0 · 4269fe0
1 parent 97be1ff
commit 4269fe0
Show file tree

Hide file tree

Showing 5 changed files with 80 additions and 42 deletions.
diff --git a/extensions/mssql/src/objectManagement/ui/backupDatabaseDialog.ts b/extensions/mssql/src/objectManagement/ui/backupDatabaseDialog.ts
@@ -83,6 +83,10 @@ export class BackupDatabaseDialog extends ObjectManagementDialogBase<Database, D
 		return true;
 	}
 
+	protected override get startTaskOnApply(): boolean {
+		return false; // The underlying backup operation in the SQL Tools Service starts its own task separately
+	}
+
 	private get encryptionSupported(): boolean {
 		return this._encryptorOptions.length > 0;
 	}

diff --git a/extensions/mssql/src/objectManagement/ui/objectManagementDialogBase.ts b/extensions/mssql/src/objectManagement/ui/objectManagementDialogBase.ts
@@ -72,44 +72,67 @@ export abstract class ObjectManagementDialogBase<ObjectInfoType extends ObjectMa
 		return this.options.isNewObject ? TelemetryActions.CreateObject : TelemetryActions.UpdateObject;
 	}
 
-	protected override async initialize(): Promise<void> {
-		await super.initialize();
-		this.dialogObject.registerOperation({
-			displayName: this.saveChangesTaskLabel,
-			description: '',
-			isCancelable: false,
-			operation: async (operation: azdata.BackgroundOperation): Promise<void> => {
-				try {
-					if (this.isDirty) {
-						const startTime = Date.now();
-						await this.saveChanges(this._contextId, this.objectInfo);
-						if (this.options.objectExplorerContext) {
-							if (this.options.isNewObject) {
-								await refreshNode(this.options.objectExplorerContext);
-							} else {
-								// For edit mode, the node context is the object itself, we need to refresh the parent node to reflect the changes.
-								await refreshParentNode(this.options.objectExplorerContext);
-							}
-						}
-
-						TelemetryReporter.sendTelemetryEvent(this.actionName, {
-							objectType: this.options.objectType
-						}, {
-							elapsedTimeMs: Date.now() - startTime
-						});
-						operation.updateStatus(azdata.TaskStatus.Succeeded);
+	/**
+	 * Whether to start a background task after clicking OK in the dialog. Some operations, like Backup & Restore,
+	 * start their own background tasks, and so don't need one started directly from the dialog.
+	 */
+	protected get startTaskOnApply(): boolean {
+		return true;
+	}
+
+	private async saveChangesAndRefresh(operation?: azdata.BackgroundOperation): Promise<void> {
+		try {
+			if (this.isDirty) {
+				const startTime = Date.now();
+				await this.saveChanges(this._contextId, this.objectInfo);
+				if (this.options.objectExplorerContext) {
+					if (this.options.isNewObject) {
+						await refreshNode(this.options.objectExplorerContext);
+					} else {
+						// For edit mode, the node context is the object itself, we need to refresh the parent node to reflect the changes.
+						await refreshParentNode(this.options.objectExplorerContext);
 					}
 				}
-				catch (err) {
-					operation.updateStatus(azdata.TaskStatus.Failed, getErrorMessage(err));
-					TelemetryReporter.createErrorEvent2(ObjectManagementViewName, this.actionName, err).withAdditionalProperties({
-						objectType: this.options.objectType
-					}).send();
-				} finally {
-					await this.disposeView();
+
+				TelemetryReporter.sendTelemetryEvent(this.actionName, {
+					objectType: this.options.objectType
+				}, {
+					elapsedTimeMs: Date.now() - startTime
+				});
+				if (operation) {
+					operation.updateStatus(azdata.TaskStatus.Succeeded);
 				}
 			}
-		});
+		}
+		catch (err) {
+			if (operation) {
+				operation.updateStatus(azdata.TaskStatus.Failed, getErrorMessage(err));
+			}
+			TelemetryReporter.createErrorEvent2(ObjectManagementViewName, this.actionName, err).withAdditionalProperties({
+				objectType: this.options.objectType
+			}).send();
+		} finally {
+			await this.disposeView();
+		}
+	}
+
+	protected override async handleDialogClosed(reason: azdata.window.CloseReason): Promise<any> {
+		if (reason === 'ok') {
+			if (this.startTaskOnApply) {
+				azdata.tasks.startBackgroundOperation({
+					displayName: this.saveChangesTaskLabel,
+					description: '',
+					isCancelable: false,
+					operation: async (operation: azdata.BackgroundOperation): Promise<void> => {
+						await this.saveChangesAndRefresh(operation);
+					}
+				});
+			} else {
+				await this.saveChangesAndRefresh();
+			}
+		}
+		let result = await super.handleDialogClosed(reason);
+		return result;
 	}
 
 	protected get viewInfo(): ViewInfoType {

diff --git a/extensions/mssql/src/objectManagement/ui/restoreDatabaseDialog.ts b/extensions/mssql/src/objectManagement/ui/restoreDatabaseDialog.ts
@@ -83,6 +83,10 @@ export class RestoreDatabaseDialog extends ObjectManagementDialogBase<Database,
 		return true;
 	}
 
+	protected override get startTaskOnApply(): boolean {
+		return false; // The underlying restore operation in the SQL Tools Service starts its own task separately
+	}
+
 	private get restoreDialogDocUrl(): string {
 		let helpUrl = '';
 		switch (this.activeTabId) {

diff --git a/extensions/mssql/src/objectManagement/ui/serverPropertiesDialog.ts b/extensions/mssql/src/objectManagement/ui/serverPropertiesDialog.ts
@@ -92,12 +92,15 @@ export class ServerPropertiesDialog extends ObjectManagementDialogBase<Server, S
 
 	constructor(objectManagementService: IObjectManagementService, options: ObjectManagementDialogOptions) {
 		super(objectManagementService, options);
-		this.disposables.push(this.dialogObject.onClosed(async (reason: azdata.window.CloseReason) => {
-			if (reason === 'ok') {
-				// only show message if user apply changes
-				await this.notifyServerRestart();
-			}
-		}));
+	}
+
+	protected override async handleDialogClosed(reason: azdata.window.CloseReason): Promise<any> {
+		if (reason === 'ok') {
+			// only show message if user apply changes
+			await this.notifyServerRestart();
+		}
+		let result = await super.handleDialogClosed(reason);
+		return result;
 	}
 
 	protected override get helpUrl(): string {

diff --git a/extensions/mssql/src/ui/dialogBase.ts b/extensions/mssql/src/ui/dialogBase.ts
@@ -61,13 +61,17 @@ export abstract class DialogBase<DialogResult> {
 		});
 		this._closePromise = new Promise<DialogResult | undefined>(resolve => {
 			this.disposables.push(this.dialogObject.onClosed(async (reason: azdata.window.CloseReason) => {
-				await this.dispose(reason);
-				const result = reason === 'ok' ? this.dialogResult : undefined;
+				let result = await this.handleDialogClosed(reason);
 				resolve(result);
 			}));
 		});
 	}
 
+	protected async handleDialogClosed(reason: azdata.window.CloseReason): Promise<any> {
+		await this.dispose(reason);
+		return reason === 'ok' ? this.dialogResult : undefined;
+	}
+
 	public waitForClose(): Promise<DialogResult | undefined> {
 		return this._closePromise;
 	}