Save And Close Functionality (#17000)

* save and close

* wip

* working save and close

* cleanup

* pr changes

* pr changes

* fix capitalization

* fix build

* pr fix

extensions/sql-migration/src/constants/strings.ts

@@ -16,6 +16,11 @@ export function WIZARD_TITLE(instanceName: string): string {
 }
 // //#endregion
 
+// Resume Migration Dialog
+export const RESUME_TITLE = localize('sql.migration.resume.title', "Run migration workflow again");
+export const START_MIGRATION = localize('sql.migration.resume.start', "Start with migration assessment again (recommended)");
+export const CONTINUE_MIGRATION = localize('sql.migration.resume.contine', "Continue last migration attempt...");
+
 // Assessments Progress Page
 export const ASSESSMENT_BLOCKING_ISSUE_TITLE = localize('sql.migration.assessments.blocking.issue', 'This is a blocking issue that will prevent the database migration from succeeding.');
 export const ASSESSMENT_IN_PROGRESS = localize('sql.migration.assessment.in.progress', "Assessment in progress");
@@ -455,6 +460,7 @@ export const SQL_MIGRATION_SERVICE_DETAILS_AUTH_KEYS_TITLE = localize('sql.migra
 export const SQL_MIGRATION_SERVICE_DETAILS_STATUS_UNAVAILABLE = localize('sql.migration.service.details.status.unavailable', "-- unavailable --");
 
 //Source Credentials page.
+export const SAVE_AND_CLOSE = localize('sql.migration.save.close', "Save and close");
 export const SOURCE_CONFIGURATION = localize('sql.migration.source.configuration', "Source configuration");
 export const SOURCE_CREDENTIALS = localize('sql.migration.source.credentials', "Source credentials");
 export const ENTER_YOUR_SQL_CREDS = localize('sql.migration.enter.your.sql.cred', "Enter the credentials for the source SQL Server instance. These credentials will be used while migrating databases to Azure SQL.");
@@ -507,3 +513,9 @@ export function WARNINGS_COUNT(totalCount: number): string {
 export const AUTHENTICATION_TYPE = localize('sql.migration.authentication.type', "Authentication type");
 
 export const REFRESH_BUTTON_LABEL = localize('sql.migration.status.refresh.label', 'Refresh');
+
+// Saved Assessment Dialog
+
+export const NEXT_LABEL = localize('sql.migration.saved.assessment.next', "Next");
+export const CANCEL_LABEL = localize('sql.migration.saved.assessment.cancel', "Cancel");
+export const SAVED_ASSESSMENT_RESULT = localize('sql.migration.saved.assessment.result', "Saved assessment result");

extensions/sql-migration/src/dialog/assessmentResults/assessmentResultsDialog.ts

@@ -32,6 +32,9 @@ export class AssessmentResultsDialog {
 
 	constructor(public ownerUri: string, public model: MigrationStateModel, public title: string, private _skuRecommendationPage: SKURecommendationPage, private _targetType: MigrationTargetType) {
 		this._model = model;
+		if (this._model.resumeAssessment && this._model.savedInfo.closedPage >= 2) {
+			this._model._databaseAssessment = <string[]>this._model.savedInfo.databaseAssessment;
+		}
 		this._tree = new SqlDatabaseTree(this._model, this._targetType);
 	}
 

extensions/sql-migration/src/dialog/assessmentResults/savedAssessmentDialog.ts

@@ -0,0 +1,144 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as azdata from 'azdata';
+import * as vscode from 'vscode';
+import * as constants from '../../constants/strings';
+import { MigrationStateModel } from '../../models/stateMachine';
+import { WizardController } from '../../wizard/wizardController';
+
+export class SavedAssessmentDialog {
+
+	private static readonly OkButtonText: string = constants.NEXT_LABEL;
+	private static readonly CancelButtonText: string = constants.CANCEL_LABEL;
+
+	private _isOpen: boolean = false;
+	private dialog: azdata.window.Dialog | undefined;
+	private _rootContainer!: azdata.FlexContainer;
+	private stateModel: MigrationStateModel;
+	private context: vscode.ExtensionContext;
+	private _disposables: vscode.Disposable[] = [];
+
+	constructor(context: vscode.ExtensionContext, stateModel: MigrationStateModel) {
+		this.stateModel = stateModel;
+		this.context = context;
+	}
+
+	private async initializeDialog(dialog: azdata.window.Dialog): Promise<void> {
+		return new Promise<void>((resolve, reject) => {
+			dialog.registerContent(async (view) => {
+				try {
+					this._rootContainer = this.initializePageContent(view);
+					await view.initializeModel(this._rootContainer);
+					this._disposables.push(dialog.okButton.onClick(async e => {
+						await this.execute();
+					}));
+					this._disposables.push(dialog.cancelButton.onClick(e => {
+						this.cancel();
+					}));
+
+					this._disposables.push(view.onClosed(e => {
+						this._disposables.forEach(
+							d => { try { d.dispose(); } catch { } }
+						);
+					}));
+					resolve();
+				} catch (ex) {
+					reject(ex);
+				}
+			});
+		});
+	}
+
+	public async openDialog(dialogName?: string) {
+		if (!this._isOpen) {
+			this._isOpen = true;
+			this.dialog = azdata.window.createModelViewDialog(constants.SAVED_ASSESSMENT_RESULT, constants.SAVED_ASSESSMENT_RESULT, '60%');
+			this.dialog.okButton.label = SavedAssessmentDialog.OkButtonText;
+			this.dialog.cancelButton.label = SavedAssessmentDialog.CancelButtonText;
+			const dialogSetupPromises: Thenable<void>[] = [];
+			dialogSetupPromises.push(this.initializeDialog(this.dialog));
+			azdata.window.openDialog(this.dialog);
+			await Promise.all(dialogSetupPromises);
+		}
+	}
+
+	protected async execute() {
+		if (this.stateModel.resumeAssessment) {
+			const wizardController = new WizardController(this.context, this.stateModel);
+			await wizardController.openWizard(this.stateModel.sourceConnectionId);
+			console.log(this.stateModel.savedInfo.selectedDatabases);
+		} else {
+			// normal flow
+			const wizardController = new WizardController(this.context, this.stateModel);
+			await wizardController.openWizard(this.stateModel.sourceConnectionId);
+		}
+		this._isOpen = false;
+	}
+
+	protected cancel() {
+		this._isOpen = false;
+	}
+
+	public get isOpen(): boolean {
+		return this._isOpen;
+	}
+
+	public initializePageContent(view: azdata.ModelView): azdata.FlexContainer {
+		const buttonGroup = 'resumeMigration';
+
+		const text = view.modelBuilder.text().withProps({
+			CSSStyles: {
+				'font-size': '18px',
+				'font-weight': 'bold',
+				'margin': '100px 8px 0px 36px'
+			},
+			value: constants.RESUME_TITLE
+		}).component();
+
+		const radioStart = view.modelBuilder.radioButton().withProps({
+			label: constants.START_MIGRATION,
+			name: buttonGroup,
+			CSSStyles: {
+				'font-size': '14px',
+				'margin': '40px 8px 0px 36px'
+			},
+			checked: true
+		}).component();
+
+		radioStart.onDidChangeCheckedState((e) => {
+			if (e) {
+				this.stateModel.resumeAssessment = false;
+			}
+		});
+		const radioContinue = view.modelBuilder.radioButton().withProps({
+			label: constants.CONTINUE_MIGRATION,
+			name: buttonGroup,
+			CSSStyles: {
+				'font-size': '14px',
+				'margin': '10px 8px 0px 36px'
+			},
+			checked: false
+		}).component();
+
+		radioContinue.onDidChangeCheckedState((e) => {
+			if (e) {
+				this.stateModel.resumeAssessment = true;
+			}
+		});
+
+		const flex = view.modelBuilder.flexContainer().withLayout({
+			flexFlow: 'column',
+			height: '100%',
+			width: '100%'
+		}).component();
+		flex.addItem(text, { flex: '0 0 auto' });
+		flex.addItem(radioStart, { flex: '0 0 auto' });
+		flex.addItem(radioContinue, { flex: '0 0 auto' });
+
+		return flex;
+	}
+
+}

extensions/sql-migration/src/dialog/assessmentResults/sqlDatabasesTree.ts

@@ -187,6 +187,7 @@ export class SqlDatabaseTree {
 			await this._databaseCount.updateProperties({
 				'value': constants.DATABASES(this.selectedDbs().length, this._model._databaseAssessment.length)
 			});
+			this._model._databaseSelection = <azdata.DeclarativeTableCellValue[][]>this._databaseTable.dataValues;
 		}));
 
 		this._disposables.push(this._databaseTable.onRowSelected(async (e) => {
@@ -947,8 +948,12 @@ export class SqlDatabaseTree {
 			});
 		}
 		await this._instanceTable.setDataValues(instanceTableValues);
+		if (this._model.resumeAssessment && this._model.savedInfo.closedPage >= 2) {
+			await this._databaseTable.setDataValues(this._model.savedInfo.migrationDatabases);
+		} else {
 			await this._databaseTable.setDataValues(this._databaseTableValues);
 		}
+	}
 
 	// undo when bug #16445 is fixed
 	private createIconTextCell(icon: IconPath, text: string): string {

extensions/sql-migration/src/main.ts

@@ -6,15 +6,20 @@
 import * as vscode from 'vscode';
 import * as azdata from 'azdata';
 import { WizardController } from './wizard/wizardController';
+import * as mssql from '../../mssql';
 import { promises as fs } from 'fs';
 import * as loc from './constants/strings';
 import { MigrationNotebookInfo, NotebookPathHelper } from './constants/notebookPathHelper';
 import { IconPathHelper } from './constants/iconPathHelper';
 import { DashboardWidget } from './dashboard/sqlServerDashboard';
 import { MigrationLocalStorage } from './models/migrationLocalStorage';
+import { MigrationStateModel, SavedInfo } from './models/stateMachine';
+import { SavedAssessmentDialog } from './dialog/assessmentResults/savedAssessmentDialog';
 
 class SQLMigration {
 
+	public stateModel!: MigrationStateModel;
+
 	constructor(private readonly context: vscode.ExtensionContext) {
 		NotebookPathHelper.setExtensionContext(context);
 		IconPathHelper.setExtensionContext(context);
@@ -76,17 +81,48 @@ class SQLMigration {
 	async launchMigrationWizard(): Promise<void> {
 		let activeConnection = await azdata.connection.getCurrentConnection();
 		let connectionId: string = '';
+		let serverName: string = '';
 		if (!activeConnection) {
 			const connection = await azdata.connection.openConnectionDialog();
 			if (connection) {
 				connectionId = connection.connectionId;
+				serverName = connection.options.server;
 			}
 		} else {
 			connectionId = activeConnection.connectionId;
+			serverName = activeConnection.serverName;
 		}
-		const wizardController = new WizardController(this.context);
+		if (serverName) {
+			const api = (await vscode.extensions.getExtension(mssql.extension.name)?.activate()) as mssql.IExtension;
+			if (api) {
+				this.stateModel = new MigrationStateModel(this.context, connectionId, api.sqlMigration);
+				this.context.subscriptions.push(this.stateModel);
+				let savedInfo = this.checkSavedInfo(serverName);
+				if (savedInfo) {
+					this.stateModel.savedInfo = savedInfo;
+					this.stateModel.serverName = serverName;
+					let savedAssessmentDialog = new SavedAssessmentDialog(this.context, this.stateModel);
+					await savedAssessmentDialog.openDialog();
+				} else {
+					const wizardController = new WizardController(this.context, this.stateModel);
 					await wizardController.openWizard(connectionId);
 				}
+			}
+
+		}
+
+
+
+	}
+
+	private checkSavedInfo(serverName: string): SavedInfo | undefined {
+		let savedInfo: SavedInfo | undefined = this.context.globalState.get(`${this.stateModel.mementoString}.${serverName}`);
+		if (savedInfo) {
+			return savedInfo;
+		} else {
+			return undefined;
+		}
+	}
 
 	async launchNewSupportRequest(): Promise<void> {
 		await vscode.env.openExternal(vscode.Uri.parse(

extensions/sql-migration/src/models/stateMachine.ts

@@ -15,7 +15,7 @@ import { MigrationLocalStorage } from './migrationLocalStorage';
 import * as nls from 'vscode-nls';
 import { v4 as uuidv4 } from 'uuid';
 import { sendSqlMigrationActionEvent, TelemetryAction, TelemetryViews } from '../telemtery';
-import { hashString } from '../api/utils';
+import { hashString, deepClone } from '../api/utils';
 const localize = nls.loadMessageBundle();
 
 export enum State {
@@ -58,6 +58,16 @@ export enum NetworkContainerType {
 	NETWORK_SHARE
 }
 
+export enum Page {
+	AzureAccount,
+	DatabaseSelector,
+	SKURecommendation,
+	MigrationMode,
+	DatabaseBackup,
+	IntegrationRuntime,
+	Summary
+}
+
 export interface DatabaseBackupModel {
 	migrationMode: MigrationMode;
 	networkContainerType: NetworkContainerType;
@@ -97,10 +107,28 @@ export interface StateChangeEvent {
 	newState: State;
 }
 
+export interface SavedInfo {
+	closedPage: number;
+	serverAssessment: ServerAssessment | null;
+	azureAccount: azdata.Account | null;
+	azureTenant: azurecore.Tenant | null;
+	selectedDatabases: azdata.DeclarativeTableCellValue[][];
+	migrationTargetType: MigrationTargetType | null;
+	migrationDatabases: azdata.DeclarativeTableCellValue[][];
+	subscription: azureResource.AzureResourceSubscription | null;
+	location: azureResource.AzureLocation | null;
+	resourceGroup: azureResource.AzureResourceResourceGroup | null;
+	targetServerInstance: azureResource.AzureSqlManagedInstance | SqlVMServer | null;
+	migrationMode: MigrationMode | null;
+	databaseAssessment: string[] | null;
+}
+
+
 export class MigrationStateModel implements Model, vscode.Disposable {
 	public _azureAccounts!: azdata.Account[];
 	public _azureAccount!: azdata.Account;
 	public _accountTenants!: azurecore.Tenant[];
+	public _azureTenant!: azurecore.Tenant;
 
 	public _connecionProfile!: azdata.connection.ConnectionProfile;
 	public _authenticationType!: MigrationSourceAuthenticationType;
@@ -137,15 +165,20 @@ export class MigrationStateModel implements Model, vscode.Disposable {
 	private _gatheringInformationError: string | undefined;
 
 	private _skuRecommendations: SKURecommendations | undefined;
-	public _assessmentResults!: ServerAssessement;
+	public _assessmentResults!: ServerAssessment;
 	public _runAssessments: boolean = true;
 	private _assessmentApiResponse!: mssql.AssessmentResult;
+	public mementoString: string;
 
 	public _vmDbs: string[] = [];
 	public _miDbs: string[] = [];
 	public _targetType!: MigrationTargetType;
 	public refreshDatabaseBackupPage!: boolean;
 
+	public _databaseSelection!: azdata.DeclarativeTableCellValue[][];
+	public resumeAssessment!: boolean;
+	public savedInfo!: SavedInfo;
+	public closedPage!: number;
 	public _sessionId: string = uuidv4();
 
 	public excludeDbs: string[] = [
@@ -154,9 +187,11 @@ export class MigrationStateModel implements Model, vscode.Disposable {
 		'msdb',
 		'model'
 	];
+	public serverName!: string;
+	public databaseSelectorTableValues!: azdata.DeclarativeTableCellValue[][];
 
 	constructor(
-		private readonly _extensionContext: vscode.ExtensionContext,
+		public extensionContext: vscode.ExtensionContext,
 		private readonly _sourceConnectionId: string,
 		public readonly migrationService: mssql.ISqlMigrationService
 	) {
@@ -164,6 +199,7 @@ export class MigrationStateModel implements Model, vscode.Disposable {
 		this._databaseBackup = {} as DatabaseBackupModel;
 		this._databaseBackup.networkShare = {} as NetworkShare;
 		this._databaseBackup.blobs = [];
+		this.mementoString = 'sqlMigration.assessmentResults';
 	}
 
 	public get sourceConnectionId(): string {
@@ -185,7 +221,7 @@ export class MigrationStateModel implements Model, vscode.Disposable {
 		return finalResult;
 	}
 
-	public async getDatabaseAssessments(targetType: MigrationTargetType): Promise<ServerAssessement> {
+	public async getDatabaseAssessments(targetType: MigrationTargetType): Promise<ServerAssessment> {
 		const ownerUri = await azdata.connection.getUriForConnection(this.sourceConnectionId);
 		try {
 			const response = (await this.migrationService.getAssessments(ownerUri, this._databaseAssessment))!;
@@ -388,7 +424,7 @@ export class MigrationStateModel implements Model, vscode.Disposable {
 	}
 
 	public getExtensionPath(): string {
-		return this._extensionContext.extensionPath;
+		return this.extensionContext.extensionPath;
 	}
 
 	public async getAccountValues(): Promise<azdata.CategoryValue[]> {
@@ -972,9 +1008,53 @@ export class MigrationStateModel implements Model, vscode.Disposable {
 			await vscode.commands.executeCommand('sqlmigration.refreshMigrationTiles');
 		}
 	}
+
+	public async saveInfo(serverName: string, currentPage: Page): Promise<void> {
+		let saveInfo: SavedInfo;
+		saveInfo = {
+			closedPage: currentPage,
+			serverAssessment: null,
+			azureAccount: null,
+			azureTenant: null,
+			selectedDatabases: [],
+			migrationTargetType: null,
+			migrationDatabases: [],
+			subscription: null,
+			location: null,
+			resourceGroup: null,
+			targetServerInstance: null,
+			migrationMode: null,
+			databaseAssessment: null
+		};
+		switch (currentPage) {
+			case Page.Summary:
+
+			case Page.IntegrationRuntime:
+
+			case Page.DatabaseBackup:
+
+			case Page.MigrationMode:
+				saveInfo.migrationMode = this._databaseBackup.migrationMode;
+			case Page.SKURecommendation:
+				saveInfo.migrationTargetType = this._targetType;
+				saveInfo.databaseAssessment = this._databaseAssessment;
+				saveInfo.serverAssessment = this._assessmentResults;
+				saveInfo.migrationDatabases = this._databaseSelection;
+				saveInfo.subscription = this._targetSubscription;
+				saveInfo.location = this._location;
+				saveInfo.resourceGroup = this._resourceGroup;
+				saveInfo.targetServerInstance = this._targetServerInstance;
+			case Page.DatabaseSelector:
+				saveInfo.selectedDatabases = this.databaseSelectorTableValues;
+			case Page.AzureAccount:
+				saveInfo.azureAccount = deepClone(this._azureAccount);
+				saveInfo.azureTenant = deepClone(this._azureTenant);
+				await this.extensionContext.globalState.update(`${this.mementoString}.${serverName}`, saveInfo);
+		}
+	}
 }
 
-export interface ServerAssessement {
+export interface ServerAssessment {
 	issues: mssql.SqlMigrationAssessmentResultItem[];
 	databaseAssessments: {
 		name: string;

extensions/sql-migration/src/wizard/accountsSelectionPage.ts

@@ -24,6 +24,7 @@ export class AccountsSelectionPage extends MigrationWizardPage {
 	}
 
 	protected async registerContent(view: azdata.ModelView): Promise<void> {
+		this.wizard.customButtons[0].enabled = true;
 		const form = view.modelBuilder.formContainer()
 			.withFormItems(
 				[
@@ -95,6 +96,14 @@ export class AccountsSelectionPage extends MigrationWizardPage {
 					await this._accountTenantFlexContainer.updateCssStyles({
 						'display': 'none'
 					});
+					if (this.migrationStateModel.resumeAssessment && this.migrationStateModel.savedInfo.closedPage >= 0) {
+						(<azdata.CategoryValue[]>this._azureAccountsDropdown.values)?.forEach((account, index) => {
+							if (account.name === this.migrationStateModel.savedInfo.azureAccount?.displayInfo.userId) {
+								selectDropDownIndex(this._azureAccountsDropdown, index);
+							}
+						});
+					}
+
 				}
 				this.migrationStateModel._subscriptions = undefined!;
 				this.migrationStateModel._targetSubscription = undefined!;
@@ -162,12 +171,17 @@ export class AccountsSelectionPage extends MigrationWizardPage {
 			 * All azure requests will only run on this tenant from now on
 			 */
 			const selectedIndex = findDropDownItemIndex(this._accountTenantDropdown, value);
+			const selectedTenant = this.migrationStateModel.getTenant(selectedIndex);
+			this.migrationStateModel._azureTenant = deepClone(selectedTenant);
 			if (selectedIndex > -1) {
 				this.migrationStateModel._azureAccount.properties.tenants = [this.migrationStateModel.getTenant(selectedIndex)];
 				this.migrationStateModel._subscriptions = undefined!;
 				this.migrationStateModel._targetSubscription = undefined!;
 				this.migrationStateModel._databaseBackup.subscription = undefined!;
 			}
+			const selectedAzureAccount = this.migrationStateModel.getAccount(selectedIndex);
+			this.migrationStateModel._azureAccount = deepClone(selectedAzureAccount);
+
 		}));
 
 		this._accountTenantFlexContainer = view.modelBuilder.flexContainer()

extensions/sql-migration/src/wizard/databaseSelectorPage.ts

@@ -105,7 +105,6 @@ export class DatabaseSelectorPage extends MigrationWizardPage {
 			|| assessedDatabases.length !== selectedDatabases.length
 			|| assessedDatabases.some(db => selectedDatabases.indexOf(db) < 0);
 
-		this.migrationStateModel._databaseAssessment = selectedDatabases;
 		this.wizard.message = {
 			text: '',
 			level: azdata.window.MessageLevel.Error
@@ -287,11 +286,17 @@ export class DatabaseSelectorPage extends MigrationWizardPage {
 			}
 		).component();
 
+		if (this.migrationStateModel.resumeAssessment && this.migrationStateModel.savedInfo.closedPage >= 1) {
+			await this._databaseSelectorTable.setDataValues(this.migrationStateModel.savedInfo.selectedDatabases);
+		} else {
 			await this._databaseSelectorTable.setDataValues(this._databaseTableValues);
+		}
 		this._disposables.push(this._databaseSelectorTable.onDataChanged(async () => {
 			await this._dbCount.updateProperties({
 				'value': constants.DATABASES_SELECTED(this.selectedDbs().length, this._databaseTableValues.length)
 			});
+			this.migrationStateModel._databaseAssessment = this.selectedDbs();
+			this.migrationStateModel.databaseSelectorTableValues = <azdata.DeclarativeTableCellValue[][]>this._databaseSelectorTable.dataValues;
 		}));
 		const flex = view.modelBuilder.flexContainer().withLayout({
 			flexFlow: 'column',

extensions/sql-migration/src/wizard/migrationModePage.ts

@@ -97,6 +97,15 @@ export class MigrationModePage extends MigrationWizardPage {
 			}
 		}).component();
 
+		if (this.migrationStateModel.resumeAssessment && this.migrationStateModel.savedInfo.closedPage >= 3) {
+			if (this.migrationStateModel.savedInfo.migrationMode === MigrationMode.ONLINE) {
+				onlineButton.checked = true;
+				offlineButton.checked = false;
+			} else {
+				onlineButton.checked = false;
+				offlineButton.checked = true;
+			}
+		}
 
 		this._disposables.push(offlineButton.onDidChangeCheckedState((e) => {
 			if (e) {

extensions/sql-migration/src/wizard/skuRecommendationPage.ts

@@ -6,7 +6,7 @@
 import * as azdata from 'azdata';
 import * as vscode from 'vscode';
 import { MigrationWizardPage } from '../models/migrationWizardPage';
-import { MigrationStateModel, MigrationTargetType, StateChangeEvent } from '../models/stateMachine';
+import { MigrationStateModel, MigrationTargetType, ServerAssessment, StateChangeEvent } from '../models/stateMachine';
 import { AssessmentResultsDialog } from '../dialog/assessmentResults/assessmentResultsDialog';
 import * as constants from '../constants/strings';
 import { EOL } from 'os';
@@ -259,7 +259,13 @@ export class SKURecommendationPage extends MigrationWizardPage {
 			width: 100
 		}).component();
 
-		const serverName = (await this.migrationStateModel.getSourceConnectionProfile()).serverName;
+		let serverName = '';
+		if (this.migrationStateModel.resumeAssessment && this.migrationStateModel.serverName) {
+			serverName = this.migrationStateModel.serverName;
+		} else {
+			serverName = (await this.migrationStateModel.getSourceConnectionProfile()).serverName;
+		}
+
 		let miDialog = new AssessmentResultsDialog('ownerUri', this.migrationStateModel, constants.ASSESSMENT_TILE(serverName), this, MigrationTargetType.SQLMI);
 		let vmDialog = new AssessmentResultsDialog('ownerUri', this.migrationStateModel, constants.ASSESSMENT_TILE(serverName), this, MigrationTargetType.SQLVM);
 
@@ -467,7 +473,11 @@ export class SKURecommendationPage extends MigrationWizardPage {
 		const serverName = (await this.migrationStateModel.getSourceConnectionProfile()).serverName;
 		this._igComponent.value = constants.ASSESSMENT_COMPLETED(serverName);
 		try {
+			if (this.migrationStateModel.resumeAssessment && this.migrationStateModel.savedInfo.closedPage) {
+				this.migrationStateModel._assessmentResults = <ServerAssessment>this.migrationStateModel.savedInfo.serverAssessment;
+			} else {
 				await this.migrationStateModel.getDatabaseAssessments(MigrationTargetType.SQLMI);
+			}
 			this._detailsComponent.value = constants.SKU_RECOMMENDATION_ALL_SUCCESSFUL(this.migrationStateModel._assessmentResults.databaseAssessments.length);
 
 			const errors: string[] = [];
@@ -506,18 +516,29 @@ errorId: ${e.errorId}
 	}
 
 	private async populateSubscriptionDropdown(): Promise<void> {
+		if (this.migrationStateModel.resumeAssessment && this.migrationStateModel.savedInfo.closedPage >= 0) {
+			this.migrationStateModel._azureAccount = <azdata.Account>this.migrationStateModel.savedInfo.azureAccount;
+		}
 		if (!this.migrationStateModel._targetSubscription) {
 			this._managedInstanceSubscriptionDropdown.loading = true;
 			this._resourceDropdown.loading = true;
 			try {
 				this._managedInstanceSubscriptionDropdown.values = await this.migrationStateModel.getSubscriptionsDropdownValues();
-				selectDropDownIndex(this._managedInstanceSubscriptionDropdown, 0);
 			} catch (e) {
 				console.log(e);
 			} finally {
 				this._managedInstanceSubscriptionDropdown.loading = false;
 				this._resourceDropdown.loading = false;
 			}
+			if (this.migrationStateModel.resumeAssessment && this.migrationStateModel.savedInfo.closedPage >= 2 && this._managedInstanceSubscriptionDropdown.values) {
+				this._managedInstanceSubscriptionDropdown.values.forEach((subscription, index) => {
+					if ((<azdata.CategoryValue>subscription).name === this.migrationStateModel.savedInfo?.subscription?.id) {
+						selectDropDownIndex(this._managedInstanceSubscriptionDropdown, index);
+					}
+				});
+			} else {
+				selectDropDownIndex(this._managedInstanceSubscriptionDropdown, 0);
+			}
 		}
 	}
 
@@ -526,9 +547,25 @@ errorId: ${e.errorId}
 		this._azureLocationDropdown.loading = true;
 		try {
 			this._azureResourceGroupDropdown.values = await this.migrationStateModel.getAzureResourceGroupDropdownValues(this.migrationStateModel._targetSubscription);
+			if (this.migrationStateModel.resumeAssessment && this.migrationStateModel.savedInfo.closedPage >= 2 && this._azureResourceGroupDropdown.values) {
+				this._azureResourceGroupDropdown.values.forEach((resourceGroup, index) => {
+					if (resourceGroup.name === this.migrationStateModel.savedInfo?.resourceGroup?.id) {
+						selectDropDownIndex(this._azureResourceGroupDropdown, index);
+					}
+				});
+			} else {
 				selectDropDownIndex(this._azureResourceGroupDropdown, 0);
+			}
 			this._azureLocationDropdown.values = await this.migrationStateModel.getAzureLocationDropdownValues(this.migrationStateModel._targetSubscription);
+			if (this.migrationStateModel.resumeAssessment && this.migrationStateModel.savedInfo.closedPage >= 2 && this._azureLocationDropdown.values) {
+				this._azureLocationDropdown.values.forEach((location, index) => {
+					if (location.displayName === this.migrationStateModel.savedInfo?.location?.displayName) {
+						selectDropDownIndex(this._azureLocationDropdown, index);
+					}
+				});
+			} else {
 				selectDropDownIndex(this._azureLocationDropdown, 0);
+			}
 		} catch (e) {
 			console.log(e);
 		} finally {
@@ -555,8 +592,15 @@ errorId: ${e.errorId}
 					this.migrationStateModel._location,
 					this.migrationStateModel._resourceGroup);
 			}
-
+			if (this.migrationStateModel.resumeAssessment && this.migrationStateModel.savedInfo.closedPage >= 2 && this._resourceDropdown.values) {
+				this._resourceDropdown.values.forEach((resource, index) => {
+					if (resource.displayName === this.migrationStateModel.savedInfo?.targetServerInstance?.name) {
+						selectDropDownIndex(this._resourceDropdown, index);
+					}
+				});
+			} else {
 				selectDropDownIndex(this._resourceDropdown, 0);
+			}
 		} catch (e) {
 			console.log(e);
 		} finally {

extensions/sql-migration/src/wizard/sqlSourceConfigurationPage.ts

@@ -21,6 +21,7 @@ export class SqlSourceConfigurationPage extends MigrationWizardPage {
 	}
 
 	protected async registerContent(view: azdata.ModelView): Promise<void> {
+		this.wizard.customButtons[0].enabled = true;
 		this._view = view;
 		const form = view.modelBuilder.formContainer()
 			.withFormItems(

extensions/sql-migration/src/wizard/wizardController.ts

@@ -21,14 +21,13 @@ export const WIZARD_INPUT_COMPONENT_WIDTH = '600px';
 export class WizardController {
 	private _wizardObject!: azdata.window.Wizard;
 	private _model!: MigrationStateModel;
-	constructor(private readonly extensionContext: vscode.ExtensionContext) {
-
+	constructor(private readonly extensionContext: vscode.ExtensionContext, model: MigrationStateModel) {
+		this._model = model;
 	}
 
 	public async openWizard(connectionId: string): Promise<void> {
 		const api = (await vscode.extensions.getExtension(mssql.extension.name)?.activate()) as mssql.IExtension;
 		if (api) {
-			this._model = new MigrationStateModel(this.extensionContext, connectionId, api.sqlMigration);
 			this.extensionContext.subscriptions.push(this._model);
 			await this.createWizard(this._model);
 		}
@@ -39,6 +38,8 @@ export class WizardController {
 		this._wizardObject = azdata.window.createWizard(loc.WIZARD_TITLE(serverName), 'MigrationWizard', 'wide');
 		this._wizardObject.generateScriptButton.enabled = false;
 		this._wizardObject.generateScriptButton.hidden = true;
+		const saveAndCloseButton = azdata.window.createButton(loc.SAVE_AND_CLOSE);
+		this._wizardObject.customButtons = [saveAndCloseButton];
 		const skuRecommendationPage = new SKURecommendationPage(this._wizardObject, stateModel);
 		const migrationModePage = new MigrationModePage(this._wizardObject, stateModel);
 		const databaseSelectorPage = new DatabaseSelectorPage(this._wizardObject, stateModel);
@@ -62,8 +63,11 @@ export class WizardController {
 		const wizardSetupPromises: Thenable<void>[] = [];
 		wizardSetupPromises.push(...pages.map(p => p.registerWizardContent()));
 		wizardSetupPromises.push(this._wizardObject.open());
+		if (this._model.resumeAssessment) {
+			wizardSetupPromises.push(this._wizardObject.setCurrentPage(this._model.savedInfo.closedPage));
+		}
 
-		this.extensionContext.subscriptions.push(this._wizardObject.onPageChanged(async (pageChangeInfo: azdata.window.WizardPageChangeInfo) => {
+		this._model.extensionContext.subscriptions.push(this._wizardObject.onPageChanged(async (pageChangeInfo: azdata.window.WizardPageChangeInfo) => {
 			const newPage = pageChangeInfo.newPage;
 			const lastPage = pageChangeInfo.lastPage;
 			this.sendPageButtonClickEvent(pageChangeInfo).catch(e => console.log(e));
@@ -82,13 +86,17 @@ export class WizardController {
 		});
 
 		await Promise.all(wizardSetupPromises);
-		this.extensionContext.subscriptions.push(this._wizardObject.onPageChanged(async (pageChangeInfo: azdata.window.WizardPageChangeInfo) => {
+		this._model.extensionContext.subscriptions.push(this._wizardObject.onPageChanged(async (pageChangeInfo: azdata.window.WizardPageChangeInfo) => {
 			await pages[0].onPageEnter(pageChangeInfo);
 		}));
 
-		this.extensionContext.subscriptions.push(this._wizardObject.doneButton.onClick(async (e) => {
+		this._model.extensionContext.subscriptions.push(this._wizardObject.doneButton.onClick(async (e) => {
 			await stateModel.startMigration();
 		}));
+		saveAndCloseButton.onClick(async () => {
+			await stateModel.saveInfo(serverName, this._wizardObject.currentPage);
+			await this._wizardObject.close();
+		});
 
 		this._wizardObject.cancelButton.onClick(e => {
 			sendSqlMigrationActionEvent(