From af8611455c4cdddf20645089f981cc9017b2e078 Mon Sep 17 00:00:00 2001
From: Alan Ren <alanren@microsoft.com>
Date: Thu, 14 Oct 2021 13:36:46 -0700
Subject: [PATCH] use STS to process the requests (#17358)

---
 extensions/mssql/src/sqlToolsServer.ts              |  3 ++-
 extensions/mssql/src/tableDesigner/tableDesigner.ts |  8 ++++++--
 src/sql/base/browser/ui/designer/designer.ts        | 10 ++++++++++
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/extensions/mssql/src/sqlToolsServer.ts b/extensions/mssql/src/sqlToolsServer.ts
index b7929375aa..e4b6738958 100644
--- a/extensions/mssql/src/sqlToolsServer.ts
+++ b/extensions/mssql/src/sqlToolsServer.ts
@@ -11,7 +11,7 @@ import * as path from 'path';
 import { getCommonLaunchArgsAndCleanupOldLogFiles, getOrDownloadServer } from './utils';
 import { Telemetry, LanguageClientErrorHandler } from './telemetry';
 import { SqlOpsDataClient, ClientOptions } from 'dataprotocol-client';
-import { TelemetryFeature, AgentServicesFeature, SerializationFeature, AccountFeature, SqlAssessmentServicesFeature, ProfilerFeature } from './features';
+import { TelemetryFeature, AgentServicesFeature, SerializationFeature, AccountFeature, SqlAssessmentServicesFeature, ProfilerFeature, TableDesignerFeature } from './features';
 import { CredentialStore } from './credentialstore/credentialstore';
 import { AzureResourceProvider } from './resourceProvider/resourceProvider';
 import { SchemaCompareService } from './schemaCompare/schemaCompareService';
@@ -163,6 +163,7 @@ function getClientOptions(context: AppContext): ClientOptions {
 			NotebookConvertService.asFeature(context),
 			ProfilerFeature,
 			SqlMigrationService.asFeature(context),
+			TableDesignerFeature
 		],
 		outputChannel: new CustomOutputChannel()
 	};
diff --git a/extensions/mssql/src/tableDesigner/tableDesigner.ts b/extensions/mssql/src/tableDesigner/tableDesigner.ts
index a63bb37af9..ab33daee98 100644
--- a/extensions/mssql/src/tableDesigner/tableDesigner.ts
+++ b/extensions/mssql/src/tableDesigner/tableDesigner.ts
@@ -10,10 +10,12 @@ import { sqlProviderName } from '../constants';
 
 export function registerTableDesignerCommands(appContext: AppContext) {
 	appContext.extensionContext.subscriptions.push(vscode.commands.registerCommand('mssql.newTable', async (context: azdata.ObjectExplorerContext) => {
+		const connectionUri = await azdata.connection.getUriForConnection(context.connectionProfile.id);
 		await azdata.designers.openTableDesigner(sqlProviderName, {
 			server: context.connectionProfile.serverName,
 			database: context.connectionProfile.databaseName,
-			isNewTable: true
+			isNewTable: true,
+			connectionUri: connectionUri
 		});
 	}));
 
@@ -22,13 +24,15 @@ export function registerTableDesignerCommands(appContext: AppContext) {
 		const database = context.connectionProfile.databaseName;
 		const schema = context.nodeInfo.metadata.schema;
 		const name = context.nodeInfo.metadata.name;
+		const connectionUri = await azdata.connection.getUriForConnection(context.connectionProfile.id);
 		await azdata.designers.openTableDesigner(sqlProviderName, {
 			server: server,
 			database: database,
 			isNewTable: false,
 			name: name,
 			schema: schema,
-			id: `${server}|${database}|${schema}|${name}`
+			id: `${server}|${database}|${schema}|${name}`,
+			connectionUri: connectionUri
 		});
 	}));
 
diff --git a/src/sql/base/browser/ui/designer/designer.ts b/src/sql/base/browser/ui/designer/designer.ts
index 8a0164e967..a762509f0a 100644
--- a/src/sql/base/browser/ui/designer/designer.ts
+++ b/src/sql/base/browser/ui/designer/designer.ts
@@ -290,6 +290,9 @@ export class Designer extends Disposable implements IThemable {
 			case DesignerEditType.Update:
 				if (typeof edit.property === 'string') {
 					// if the type of the property is string then the property is a top level property
+					if (!data[edit.property]) {
+						data[edit.property] = {};
+					}
 					const componentData = data[edit.property];
 					const componentType = this._componentMap.get(edit.property).defintion.componentType;
 					this.setComponentData(componentType, componentData, edit.value);
@@ -297,6 +300,9 @@ export class Designer extends Disposable implements IThemable {
 					const columnPropertyName = edit.property.property;
 					const tableInfo = this._componentMap.get(edit.property.parentProperty).defintion.componentProperties as DesignerTableProperties;
 					const tableProperties = data[edit.property.parentProperty] as DesignerTableProperties;
+					if (!tableProperties.data[edit.property.index][columnPropertyName]) {
+						tableProperties.data[edit.property.index][columnPropertyName] = {};
+					}
 					const componentData = tableProperties.data[edit.property.index][columnPropertyName];
 					const itemProperty = tableInfo.itemProperties.find(property => property.propertyName === columnPropertyName);
 					if (itemProperty) {
@@ -335,6 +341,10 @@ export class Designer extends Disposable implements IThemable {
 	}
 
 	private setComponentValue(definition: DesignerDataPropertyInfo, component: DesignerUIComponent, data: DesignerData): void {
+		// Skip the property if it is not in the data model
+		if (!data[definition.propertyName]) {
+			return;
+		}
 		this._supressEditProcessing = true;
 		switch (definition.componentType) {
 			case 'input':