From 6631f8e2d9760af9a28c509f84139c7a2e7c4e7b Mon Sep 17 00:00:00 2001
From: siyang yao <SIYAO@MICROSOFT.COM>
Date: Tue, 24 Jan 2023 11:53:12 -0800
Subject: [PATCH] ADS: Make select all tables by default instead of having to
 click edit and do select all (#21698)

* table are selected as default

* refactor

* fix save and close bug
---
 .../tableMigrationSelectionDialog.ts          | 43 ++++++++++++-------
 .../src/wizard/targetSelectionPage.ts         |  5 +++
 .../src/wizard/wizardController.ts            | 10 +++--
 3 files changed, 40 insertions(+), 18 deletions(-)

diff --git a/extensions/sql-migration/src/dialog/tableMigrationSelection/tableMigrationSelectionDialog.ts b/extensions/sql-migration/src/dialog/tableMigrationSelection/tableMigrationSelectionDialog.ts
index 364f2495f3..2c5e9a45d5 100644
--- a/extensions/sql-migration/src/dialog/tableMigrationSelection/tableMigrationSelectionDialog.ts
+++ b/extensions/sql-migration/src/dialog/tableMigrationSelection/tableMigrationSelectionDialog.ts
@@ -48,19 +48,6 @@ export class TableMigrationSelectionDialog {
 					this._model.sourceConnectionId,
 					this._sourceDatabaseName);
 
-				this._tableSelectionMap = new Map();
-				sourceTableList.forEach(table => {
-					const sourceTable = targetDatabaseInfo.sourceTables.get(table.tableName);
-					const isSelected = sourceTable?.selectedForMigration === true;
-					const tableInfo: TableInfo = {
-						databaseName: table.databaseName,
-						rowCount: table.rowCount,
-						selectedForMigration: isSelected,
-						tableName: table.tableName,
-					};
-					this._tableSelectionMap.set(table.tableName, tableInfo);
-				});
-
 				const targetTableList: TableInfo[] = await collectTargetDatabaseTableInfo(
 					this._model._targetServerInstance as AzureSqlDatabaseServer,
 					targetDatabaseInfo.databaseName,
@@ -77,6 +64,28 @@ export class TableMigrationSelectionDialog {
 						selectedForMigration: false,
 						tableName: table.tableName,
 					}));
+
+				this._tableSelectionMap = new Map();
+				sourceTableList.forEach(table => {
+					// If the source table doesn't exist in the target, set isSelected to false.
+					// Otherwise, set it to true as default.
+					var isSelected = false;
+					var sourceTable = targetDatabaseInfo.sourceTables.get(table.tableName);
+					if (sourceTable === null || sourceTable === undefined) {
+						sourceTable = this._targetTableMap.get(table.tableName);
+						isSelected = sourceTable === null || sourceTable === undefined ? false : true;
+					} else {
+						isSelected = sourceTable.selectedForMigration;
+					}
+
+					const tableInfo: TableInfo = {
+						databaseName: table.databaseName,
+						rowCount: table.rowCount,
+						selectedForMigration: isSelected,
+						tableName: table.tableName,
+					};
+					this._tableSelectionMap.set(table.tableName, tableInfo);
+				});
 			}
 		} catch (error) {
 			this._dialog!.message = {
@@ -97,7 +106,9 @@ export class TableMigrationSelectionDialog {
 		let tableRow = 0;
 		this._missingTableCount = 0;
 		this._tableSelectionMap.forEach(sourceTable => {
-			if (filterText?.length === 0 || sourceTable.tableName.indexOf(filterText) > -1) {
+			const tableName = sourceTable.tableName.toLocaleLowerCase();
+			const searchText = filterText.toLocaleLowerCase();
+			if (filterText?.length === 0 || tableName.indexOf(searchText) > -1) {
 				const targetTable = this._targetTableMap.get(sourceTable.tableName);
 				if (targetTable) {
 					const targetTableRowCount = targetTable?.rowCount ?? 0;
@@ -113,9 +124,11 @@ export class TableMigrationSelectionDialog {
 					if (sourceTable.selectedForMigration) {
 						selectedItems.push(tableRow);
 					}
+
+					tableRow++;
 				}
+
 				this._missingTableCount += targetTable ? 0 : 1;
-				tableRow++;
 			}
 		});
 		await this._tableSelectionTable.updateProperty('data', data);
diff --git a/extensions/sql-migration/src/wizard/targetSelectionPage.ts b/extensions/sql-migration/src/wizard/targetSelectionPage.ts
index 3632b80adc..2a75daf53b 100644
--- a/extensions/sql-migration/src/wizard/targetSelectionPage.ts
+++ b/extensions/sql-migration/src/wizard/targetSelectionPage.ts
@@ -76,6 +76,11 @@ export class TargetSelectionPage extends MigrationWizardPage {
 				this._disposables.forEach(
 					d => { try { d.dispose(); } catch { } });
 			}));
+
+		if (this.migrationStateModel.resumeAssessment) {
+			await this.populateAzureAccountsDropdown();
+		}
+
 		await this._view.initializeModel(form);
 	}
 
diff --git a/extensions/sql-migration/src/wizard/wizardController.ts b/extensions/sql-migration/src/wizard/wizardController.ts
index fe633658ed..eeb77e27ad 100644
--- a/extensions/sql-migration/src/wizard/wizardController.ts
+++ b/extensions/sql-migration/src/wizard/wizardController.ts
@@ -112,10 +112,14 @@ export class WizardController {
 				this._model.refreshDatabaseBackupPage = true;
 			}
 
-			// if the user selected network share and selected save & close afterwards, it should always return to the database backup page so that
-			// the user can input their password again
-			if (this._model.savedInfo.closedPage >= Page.IntegrationRuntime &&
+			if (this._model.savedInfo.closedPage >= Page.IntegrationRuntime && this._model.isSqlDbTarget) {
+				// if the user selected the tables and selected save & close afterwards in SQLDB scenario,
+				// it should always return to the target database selection page so that the user can input their password again
+				wizardSetupPromises.push(this._wizardObject.setCurrentPage(Page.TargetSelection));
+			} else if (this._model.savedInfo.closedPage >= Page.IntegrationRuntime &&
 				this._model.savedInfo.networkContainerType === NetworkContainerType.NETWORK_SHARE) {
+				// if the user selected network share and selected save & close afterwards, it should always return to the database backup page so that
+				// the user can input their password again
 				wizardSetupPromises.push(this._wizardObject.setCurrentPage(Page.IntegrationRuntime));
 			} else {
 				wizardSetupPromises.push(this._wizardObject.setCurrentPage(this._model.savedInfo.closedPage));