Fix bugbash SQL-Migration extension migration load issues (#16575)

* fix migration refresh issues

* merge latest

extensions/sql-migration/src/api/azure.ts

@@ -9,6 +9,7 @@ import * as azurecore from 'azurecore';
 import { azureResource } from 'azureResource';
 import * as constants from '../constants/strings';
 import { getSessionIdHeader } from './utils';
+import { ProvisioningState } from '../models/migrationLocalStorage';
 
 async function getAzureCoreAPI(): Promise<azurecore.IExtension> {
 	const api = (await vscode.extensions.getExtension(azurecore.extension.name)?.activate()) as azurecore.IExtension;
@@ -181,9 +182,9 @@ export async function createSqlMigrationService(account: azdata.Account, subscri
 	for (i = 0; i < maxRetry; i++) {
 		const asyncResponse = await api.makeAzureRestRequest(account, subscription, asyncUrl.replace('https://management.azure.com/', ''), azurecore.HttpRequestMethod.GET, undefined, true, undefined, getSessionIdHeader(sessionId));
 		const creationStatus = asyncResponse.response.data.status;
-		if (creationStatus === 'Succeeded') {
+		if (creationStatus === ProvisioningState.Succeeded) {
 			break;
-		} else if (creationStatus === 'Failed') {
+		} else if (creationStatus === ProvisioningState.Failed) {
 			throw new Error(asyncResponse.errors.toString());
 		}
 		await new Promise(resolve => setTimeout(resolve, 3000)); //adding  3 sec delay before getting creation status

extensions/sql-migration/src/api/utils.ts

@@ -3,11 +3,11 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { CategoryValue, DropDownComponent, IconPath } from 'azdata';
+import { window, CategoryValue, DropDownComponent, IconPath } from 'azdata';
 import { IconPathHelper } from '../constants/iconPathHelper';
 import { DAYS, HRS, MINUTE, SEC } from '../constants/strings';
 import { AdsMigrationStatus } from '../dialog/migrationStatus/migrationStatusDialogModel';
-import { MigrationContext } from '../models/migrationLocalStorage';
+import { MigrationStatus, MigrationContext, ProvisioningState } from '../models/migrationLocalStorage';
 import * as crypto from 'crypto';
 
 export function deepClone<T>(obj: T): T {
@@ -97,23 +97,26 @@ export function filterMigrations(databaseMigrations: MigrationContext[], statusF
 		filteredMigration = databaseMigrations.filter((value) => {
 			const status = value.migrationContext.properties.migrationStatus;
 			const provisioning = value.migrationContext.properties.provisioningState;
-			return status === 'InProgress' || status === 'Creating' || provisioning === 'Creating';
+			return status === MigrationStatus.InProgress
+				|| status === MigrationStatus.Creating
+				|| provisioning === MigrationStatus.Creating;
 		});
 	} else if (statusFilter === AdsMigrationStatus.SUCCEEDED) {
 		filteredMigration = databaseMigrations.filter((value) => {
 			const status = value.migrationContext.properties.migrationStatus;
-			return status === 'Succeeded';
+			return status === MigrationStatus.Succeeded;
 		});
 	} else if (statusFilter === AdsMigrationStatus.FAILED) {
 		filteredMigration = databaseMigrations.filter((value) => {
 			const status = value.migrationContext.properties.migrationStatus;
 			const provisioning = value.migrationContext.properties.provisioningState;
-			return status === 'Failed' || provisioning === 'Failed';
+			return status === MigrationStatus.Failed
+				|| provisioning === ProvisioningState.Failed;
 		});
 	} else if (statusFilter === AdsMigrationStatus.COMPLETING) {
 		filteredMigration = databaseMigrations.filter((value) => {
 			const status = value.migrationContext.properties.migrationStatus;
-			return status === 'Completing';
+			return status === MigrationStatus.Completing;
 		});
 	}
 	if (databaseNameFilter) {
@@ -203,17 +206,17 @@ export function getSessionIdHeader(sessionId: string): { [key: string]: string }
 
 export function getMigrationStatusImage(status: string): IconPath {
 	switch (status) {
-		case 'InProgress':
+		case MigrationStatus.InProgress:
 			return IconPathHelper.inProgressMigration;
-		case 'Succeeded':
+		case MigrationStatus.Succeeded:
 			return IconPathHelper.completedMigration;
-		case 'Creating':
+		case MigrationStatus.Creating:
 			return IconPathHelper.notStartedMigration;
-		case 'Completing':
+		case MigrationStatus.Completing:
 			return IconPathHelper.completingCutover;
-		case 'Canceling':
+		case MigrationStatus.Canceling:
 			return IconPathHelper.cancel;
-		case 'Failed':
+		case MigrationStatus.Failed:
 		default:
 			return IconPathHelper.error;
 	}
@@ -226,3 +229,9 @@ export function get12HourTime(date: Date | undefined): string {
 	};
 	return (date ? date : new Date()).toLocaleTimeString([], localeTimeStringOptions);
 }
+
+export function clearDialogMessage(dialog: window.Dialog): void {
+	dialog.message = {
+		text: ''
+	};
+}

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

@@ -6,6 +6,7 @@
 import { AzureAccount } from 'azurecore';
 import * as nls from 'vscode-nls';
 import { SupportedAutoRefreshIntervals } from '../api/utils';
+import { MigrationStatus } from '../models/migrationLocalStorage';
 import { MigrationSourceAuthenticationType } from '../models/stateMachine';
 const localize = nls.loadMessageBundle();
 
@@ -231,6 +232,7 @@ export const INVALID_REGION_ERROR = localize('sql.migration.invalid.region.error
 export const INVALID_SERVICE_NAME_ERROR = localize('sql.migration.invalid.service.name.error', "Please enter a valid name for the Migration Service.");
 export const SERVICE_NOT_FOUND = localize('sql.migration.service.not.found', "No Migration Services found. Please create a new one.");
 export const SERVICE_NOT_SETUP_ERROR = localize('sql.migration.service.not.setup', "Please add a Migration Service to proceed.");
+export const SERVICE_STATUS_REFRESH_ERROR = localize('sql.migration.service.status.refresh.error', 'An error occurred while refreshing the migration service creation status.');
 export const MANAGED_INSTANCE = localize('sql.migration.managed.instance', "Azure SQL managed instance");
 export const NO_MANAGED_INSTANCE_FOUND = localize('sql.migration.no.managedInstance.found', "No managed instance found");
 export const NO_VIRTUAL_MACHINE_FOUND = localize('sql.migration.no.virtualMachine.found', "No virtual machine found");
@@ -360,6 +362,8 @@ export const LAST_APPLIED_LSN = localize('sql.migration.last.applied.lsn', "Last
 export const LAST_APPLIED_BACKUP_FILES = localize('sql.migration.last.applied.backup.files', "Last applied backup files");
 export const LAST_APPLIED_BACKUP_FILES_TAKEN_ON = localize('sql.migration.last.applied.files.taken.on', "Last applied backup files taken on");
 export const ACTIVE_BACKUP_FILES = localize('sql.migration.active.backup.files', "Active Backup files");
+export const MIGRATION_STATUS_REFRESH_ERROR = localize('sql.migration.cutover.status.refresh.error', 'An error occurred while refreshing the migration status.');
+export const MIGRATION_CANCELLATION_ERROR = localize('sql.migration.cancel.error', 'An error occurred while canceling the migration.');
 export const STATUS = localize('sql.migration.status', "Status");
 export const BACKUP_START_TIME = localize('sql.migration.backup.start.time', "Backup start time");
 export const FIRST_LSN = localize('sql.migration.first.lsn', "First LSN");
@@ -446,7 +450,9 @@ export const StatusLookup: LookupTable<string | undefined> = {
 };
 
 export function STATUS_WARNING_COUNT(status: string, count: number): string | undefined {
-	if (status === 'InProgress' || status === 'Creating' || status === 'Completing') {
+	if (status === MigrationStatus.InProgress ||
+		status === MigrationStatus.Creating ||
+		status === MigrationStatus.Completing) {
 		switch (count) {
 			case 0:
 				return undefined;

extensions/sql-migration/src/dialog/createSqlMigrationService/createSqlMigrationServiceDialog.ts

@@ -13,6 +13,7 @@ import { azureResource } from 'azureResource';
 import { IconPathHelper } from '../../constants/iconPathHelper';
 import { CreateResourceGroupDialog } from '../createResourceGroup/createResourceGroupDialog';
 import * as EventEmitter from 'events';
+import { clearDialogMessage } from '../../api/utils';
 
 export class CreateSqlMigrationServiceDialog {
 
@@ -89,6 +90,7 @@ export class CreateSqlMigrationServiceDialog {
 				}
 
 				try {
+					clearDialogMessage(this._dialogObject);
 					this._selectedResourceGroup = resourceGroup;
 					this._createdMigrationService = await createSqlMigrationService(this._model._azureAccount, subscription, resourceGroup, location, serviceName!, this._model._sessionId);
 					if (this._createdMigrationService.error) {
@@ -97,9 +99,6 @@ export class CreateSqlMigrationServiceDialog {
 						this.setFormEnabledState(true);
 						return;
 					}
-					this._dialogObject.message = {
-						text: ''
-					};
 
 					if (this._isBlobContainerUsed) {
 						this._dialogObject.okButton.enabled = true;
@@ -511,9 +510,15 @@ export class CreateSqlMigrationServiceDialog {
 		let migrationServiceStatus!: SqlMigrationService;
 		for (let i = 0; i < maxRetries; i++) {
 			try {
+				clearDialogMessage(this._dialogObject);
 				migrationServiceStatus = await getSqlMigrationService(this._model._azureAccount, subscription, resourceGroup, location, this._createdMigrationService.name, this._model._sessionId);
 				break;
 			} catch (e) {
+				this._dialogObject.message = {
+					text: constants.SERVICE_STATUS_REFRESH_ERROR,
+					description: e.message,
+					level: azdata.window.MessageLevel.Error
+				};
 				console.log(e);
 			}
 			await new Promise(r => setTimeout(r, 5000));

extensions/sql-migration/src/dialog/migrationCutover/migrationCutoverDialog.ts

@@ -6,10 +6,10 @@
 import * as azdata from 'azdata';
 import * as vscode from 'vscode';
 import { IconPathHelper } from '../../constants/iconPathHelper';
-import { MigrationContext } from '../../models/migrationLocalStorage';
-import { MigrationCutoverDialogModel, MigrationStatus } from './migrationCutoverDialogModel';
+import { MigrationContext, MigrationStatus, ProvisioningState } from '../../models/migrationLocalStorage';
+import { MigrationCutoverDialogModel } from './migrationCutoverDialogModel';
 import * as loc from '../../constants/strings';
-import { convertByteSizeToReadableUnit, convertIsoTimeToLocalTime, getSqlServerName, getMigrationStatusImage, SupportedAutoRefreshIntervals } from '../../api/utils';
+import { convertByteSizeToReadableUnit, convertIsoTimeToLocalTime, getSqlServerName, getMigrationStatusImage, SupportedAutoRefreshIntervals, clearDialogMessage } from '../../api/utils';
 import { EOL } from 'os';
 import { ConfirmCutoverDialog } from './confirmCutoverDialog';
 import { MigrationMode } from '../../models/stateMachine';
@@ -445,7 +445,12 @@ export class MigrationCutoverDialog {
 	}
 
 	private setAutoRefresh(interval: SupportedAutoRefreshIntervals): void {
-		const shouldRefresh = (status: string | undefined) => !status || ['InProgress', 'Creating', 'Completing', 'Creating'].includes(status);
+		const shouldRefresh = (status: string | undefined) => !status
+			|| status === MigrationStatus.InProgress
+			|| status === MigrationStatus.Creating
+			|| status === MigrationStatus.Completing
+			|| status === MigrationStatus.Canceling;
+
 		if (shouldRefresh(this.getMigrationStatus())) {
 			const classVariable = this;
 			clearInterval(this._autoRefreshHandle);
@@ -474,6 +479,8 @@ export class MigrationCutoverDialog {
 		}
 
 		try {
+			clearDialogMessage(this._dialogObject);
+
 			if (this._isProvisioned() && this._isOnlineMigration()) {
 				this._cutoverButton.updateCssStyles({
 					'display': 'inline'
@@ -490,7 +497,10 @@ export class MigrationCutoverDialog {
 			errors.push(this._model.migrationStatus.properties.migrationStatusDetails?.restoreBlockingReason);
 			this._dialogObject.message = {
 				text: errors.filter(e => e !== undefined).join(EOL),
-				level: (this._model.migrationStatus.properties.migrationStatus === MigrationStatus.InProgress || this._model.migrationStatus.properties.migrationStatus === 'Completing') ? azdata.window.MessageLevel.Warning : azdata.window.MessageLevel.Error,
+				level: this._model.migrationStatus.properties.migrationStatus === MigrationStatus.InProgress
+					|| this._model.migrationStatus.properties.migrationStatus === MigrationStatus.Completing
+					? azdata.window.MessageLevel.Warning
+					: azdata.window.MessageLevel.Error,
 				description: this.getMigrationDetails()
 			};
 			const sqlServerInfo = await azdata.connection.getServerInfo((await azdata.connection.getCurrentConnection()).connectionId);
@@ -599,12 +609,21 @@ export class MigrationCutoverDialog {
 				if (restoredCount > 0 || isBlobMigration) {
 					this._cutoverButton.enabled = true;
 				}
-				this._cancelButton.enabled = true;
 			} else {
 				this._cutoverButton.enabled = false;
-				this._cancelButton.enabled = false;
 			}
+
+			this._cancelButton.enabled =
+				migrationStatusTextValue === MigrationStatus.Canceling ||
+				migrationStatusTextValue === MigrationStatus.Creating ||
+				migrationStatusTextValue === MigrationStatus.InProgress;
+
 		} catch (e) {
+			this._dialogObject.message = {
+				level: azdata.window.MessageLevel.Error,
+				text: loc.MIGRATION_STATUS_REFRESH_ERROR,
+				description: e.message
+			};
 			console.log(e);
 		} finally {
 			this.isRefreshing = false;
@@ -696,7 +715,9 @@ export class MigrationCutoverDialog {
 
 	private _isProvisioned(): boolean {
 		const { migrationStatus, provisioningState } = this._model._migration.migrationContext.properties;
-		return provisioningState === 'Succeeded' || migrationStatus === 'Completing' || migrationStatus === 'Canceling';
+		return provisioningState === ProvisioningState.Succeeded
+			|| migrationStatus === MigrationStatus.Completing
+			|| migrationStatus === MigrationStatus.Canceling;
 	}
 
 	private _isOnlineMigration(): boolean {
@@ -713,9 +734,11 @@ export class MigrationCutoverDialog {
 
 	private getMigrationStatus(): string {
 		if (this._model.migrationStatus) {
-			return this._model.migrationStatus.properties.migrationStatus ?? this._model.migrationStatus.properties.provisioningState;
+			return this._model.migrationStatus.properties.migrationStatus
+				?? this._model.migrationStatus.properties.provisioningState;
 		}
-		return this._model._migration.migrationContext.properties.migrationStatus;
+		return this._model._migration.migrationContext.properties.migrationStatus
+			?? this._model._migration.migrationContext.properties.provisioningState;
 	}
 }
 

extensions/sql-migration/src/dialog/migrationCutover/migrationCutoverDialogModel.ts

@@ -8,13 +8,6 @@ import { MigrationContext } from '../../models/migrationLocalStorage';
 import { sendSqlMigrationActionEvent, TelemetryAction, TelemetryViews } from '../../telemtery';
 import * as constants from '../../constants/strings';
 
-export enum MigrationStatus {
-	Failed = 'Failed',
-	Succeeded = 'Succeeded',
-	InProgress = 'InProgress',
-	Canceled = 'Canceled'
-}
-
 export class MigrationCutoverDialogModel {
 
 	public migrationStatus!: DatabaseMigration;

extensions/sql-migration/src/dialog/migrationStatus/migrationStatusDialog.ts

@@ -6,11 +6,11 @@
 import * as azdata from 'azdata';
 import * as vscode from 'vscode';
 import { IconPathHelper } from '../../constants/iconPathHelper';
-import { MigrationContext, MigrationLocalStorage } from '../../models/migrationLocalStorage';
+import { MigrationContext, MigrationLocalStorage, MigrationStatus, ProvisioningState } from '../../models/migrationLocalStorage';
 import { MigrationCutoverDialog } from '../migrationCutover/migrationCutoverDialog';
 import { AdsMigrationStatus, MigrationStatusDialogModel } from './migrationStatusDialogModel';
 import * as loc from '../../constants/strings';
-import { convertTimeDifferenceToDuration, filterMigrations, getMigrationStatusImage, SupportedAutoRefreshIntervals } from '../../api/utils';
+import { clearDialogMessage, convertTimeDifferenceToDuration, filterMigrations, getMigrationStatusImage, SupportedAutoRefreshIntervals } from '../../api/utils';
 import { SqlMigrationServiceDetailsDialog } from '../sqlMigrationService/sqlMigrationServiceDetailsDialog';
 import { ConfirmCutoverDialog } from '../migrationCutover/confirmCutoverDialog';
 import { MigrationCutoverDialogModel } from '../migrationCutover/migrationCutoverDialogModel';
@@ -96,8 +96,15 @@ export class MigrationStatusDialog {
 		azdata.window.openDialog(this._dialogObject);
 	}
 
-	private canCancelMigration = (status: string | undefined) => status && status in ['InProgress', 'Creating', 'Completing', 'Creating'];
-	private canCutoverMigration = (status: string | undefined) => status === 'InProgress';
+	private canCancelMigration = (status: string | undefined) => status &&
+		(
+			status === MigrationStatus.InProgress ||
+			status === MigrationStatus.Creating ||
+			status === MigrationStatus.Completing ||
+			status === MigrationStatus.Canceling
+		);
+
+	private canCutoverMigration = (status: string | undefined) => status === MigrationStatus.InProgress;
 
 	private createSearchAndRefreshContainer(): azdata.FlexContainer {
 		this._searchBox = this._view.modelBuilder.inputBox().withProps({
@@ -177,6 +184,7 @@ export class MigrationStatusDialog {
 			'sqlmigration.cutover',
 			async (migrationId: string) => {
 				try {
+					clearDialogMessage(this._dialogObject);
 					const migration = this._model._migrations.find(migration => migration.migrationContext.id === migrationId);
 					if (this.canCutoverMigration(migration?.migrationContext.properties.migrationStatus)) {
 						const cutoverDialogModel = new MigrationCutoverDialogModel(migration!);
@@ -187,6 +195,12 @@ export class MigrationStatusDialog {
 						await vscode.window.showInformationMessage(loc.MIGRATION_CANNOT_CUTOVER);
 					}
 				} catch (e) {
+					this._dialogObject.message = {
+						text: loc.MIGRATION_STATUS_REFRESH_ERROR,
+						description: e.message,
+						level: azdata.window.MessageLevel.Error
+					};
+
 					console.log(e);
 				}
 			}));
@@ -231,6 +245,7 @@ export class MigrationStatusDialog {
 			'sqlmigration.copy.migration',
 			async (migrationId: string) => {
 				try {
+					clearDialogMessage(this._dialogObject);
 					const migration = this._model._migrations.find(migration => migration.migrationContext.id === migrationId);
 					const cutoverDialogModel = new MigrationCutoverDialogModel(migration!);
 					await cutoverDialogModel.fetchStatus();
@@ -245,6 +260,12 @@ export class MigrationStatusDialog {
 
 					await vscode.window.showInformationMessage(loc.DETAILS_COPIED);
 				} catch (e) {
+					this._dialogObject.message = {
+						text: loc.MIGRATION_STATUS_REFRESH_ERROR,
+						description: e.message,
+						level: azdata.window.MessageLevel.Error
+					};
+
 					console.log(e);
 				}
 			}));
@@ -253,6 +274,7 @@ export class MigrationStatusDialog {
 			'sqlmigration.cancel.migration',
 			async (migrationId: string) => {
 				try {
+					clearDialogMessage(this._dialogObject);
 					const migration = this._model._migrations.find(migration => migration.migrationContext.id === migrationId);
 					if (this.canCancelMigration(migration?.migrationContext.properties.migrationStatus)) {
 						vscode.window.showInformationMessage(loc.CANCEL_MIGRATION_CONFIRMATION, loc.YES, loc.NO).then(async (v) => {
@@ -266,6 +288,12 @@ export class MigrationStatusDialog {
 						await vscode.window.showInformationMessage(loc.MIGRATION_CANNOT_CANCEL);
 					}
 				} catch (e) {
+					this._dialogObject.message = {
+						text: loc.MIGRATION_CANCELLATION_ERROR,
+						description: e.message,
+						level: azdata.window.MessageLevel.Error
+					};
+
 					console.log(e);
 				}
 			}));
@@ -371,7 +399,7 @@ export class MigrationStatusDialog {
 	}
 
 	private _getMigrationMode(migration: MigrationContext): string {
-		if (migration.migrationContext.properties.provisioningState === 'Creating') {
+		if (migration.migrationContext.properties.provisioningState === ProvisioningState.Creating) {
 			return '---';
 		}
 		return migration.migrationContext.properties.autoCutoverConfiguration?.autoCutover?.valueOf() ? loc.OFFLINE : loc.ONLINE;
@@ -466,11 +494,18 @@ export class MigrationStatusDialog {
 
 		this.isRefreshing = true;
 		try {
+			clearDialogMessage(this._dialogObject);
 			this._refreshLoader.loading = true;
 			const currentConnection = await azdata.connection.getCurrentConnection();
 			this._model._migrations = await MigrationLocalStorage.getMigrationsBySourceConnections(currentConnection, true);
 			await this.populateMigrationTable();
 		} catch (e) {
+			this._dialogObject.message = {
+				text: loc.MIGRATION_STATUS_REFRESH_ERROR,
+				description: e.message,
+				level: azdata.window.MessageLevel.Error
+			};
+
 			console.log(e);
 		} finally {
 			this.isRefreshing = false;
@@ -583,13 +618,10 @@ export class MigrationStatusDialog {
 	}
 
 	private _statusInfoMap(status: string): azdata.IconPath {
-		switch (status) {
-			case 'InProgress':
-			case 'Creating':
-			case 'Completing':
-				return IconPathHelper.warning;
-			default:
-				return IconPathHelper.error;
-		}
+		return status === MigrationStatus.InProgress
+			|| status === MigrationStatus.Creating
+			|| status === MigrationStatus.Completing
+			? IconPathHelper.warning
+			: IconPathHelper.error;
 	}
 }

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

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 import * as vscode from 'vscode';
 import { azureResource } from 'azureResource';
-import { DatabaseMigration, SqlMigrationService, SqlManagedInstance, getMigrationStatus, AzureAsyncOperationResource, getMigrationAsyncOperationDetails, SqlVMServer } from '../api/azure';
+import { DatabaseMigration, SqlMigrationService, SqlManagedInstance, getMigrationStatus, AzureAsyncOperationResource, getMigrationAsyncOperationDetails, SqlVMServer, getSubscriptions } from '../api/azure';
 import * as azdata from 'azdata';
 
 export class MigrationLocalStorage {
@@ -23,16 +23,20 @@ export class MigrationLocalStorage {
 		const migrationMementos: MigrationContext[] = this.context.globalState.get(this.mementoToken) || [];
 		for (let i = 0; i < migrationMementos.length; i++) {
 			const migration = migrationMementos[i];
+			migration.sessionId = migration.sessionId ?? undefinedSessionId;
 			if (migration.sourceConnectionProfile.serverName === connectionProfile.serverName) {
 				if (refreshStatus) {
 					try {
 						const backupConfiguration = migration.migrationContext.properties.backupConfiguration;
 						const sourceDatabase = migration.migrationContext.properties.sourceDatabaseName;
+
+						await this.refreshMigrationAzureAccount(migration);
+
 						migration.migrationContext = await getMigrationStatus(
 							migration.azureAccount,
 							migration.subscription,
 							migration.migrationContext,
-							migration.sessionId ?? undefinedSessionId
+							migration.sessionId!
 						);
 						migration.migrationContext.properties.sourceDatabaseName = sourceDatabase;
 						migration.migrationContext.properties.backupConfiguration = backupConfiguration;
@@ -41,7 +45,7 @@ export class MigrationLocalStorage {
 								migration.azureAccount,
 								migration.subscription,
 								migration.asyncUrl,
-								migration.sessionId ?? undefinedSessionId
+								migration.sessionId!
 							);
 						}
 					}
@@ -62,6 +66,20 @@ export class MigrationLocalStorage {
 		return result;
 	}
 
+	public static async refreshMigrationAzureAccount(migration: MigrationContext): Promise<void> {
+		if (migration.azureAccount.isStale) {
+			const accounts = await azdata.accounts.getAllAccounts();
+			const account = accounts.find(a => !a.isStale && a.key.accountId === migration.azureAccount.key.accountId);
+			if (account) {
+				const subscriptions = await getSubscriptions(account);
+				const subscription = subscriptions.find(s => s.id === migration.subscription.id);
+				if (subscription) {
+					migration.azureAccount = account;
+				}
+			}
+		}
+	}
+
 	public static saveMigration(
 		connectionProfile: azdata.connection.ConnectionProfile,
 		migrationContext: DatabaseMigration,
@@ -106,3 +124,19 @@ export interface MigrationContext {
 	asyncOperationResult?: AzureAsyncOperationResource,
 	sessionId?: string
 }
+
+export enum MigrationStatus {
+	Failed = 'Failed',
+	Succeeded = 'Succeeded',
+	InProgress = 'InProgress',
+	Canceled = 'Canceled',
+	Completing = 'Completing',
+	Creating = 'Creating',
+	Canceling = 'Canceling'
+}
+
+export enum ProvisioningState {
+	Failed = 'Failed',
+	Succeeded = 'Succeeded',
+	Creating = 'Creating'
+}