Fix queryDocument.connect (#20682)

* Fix queryEditor.connect

* Remove decorator

* undo

* fix test

* Use connectionProfile for comparison

* remove unused

* Remove helper method

src/sql/platform/connection/common/connectionProfile.ts

@@ -40,20 +40,27 @@ export class ConnectionProfile extends ProviderConnectionInfo implements interfa
 
 	public constructor(
 		capabilitiesService: ICapabilitiesService,
-		model: string | azdata.IConnectionProfile | undefined) {
+		model: string | azdata.IConnectionProfile | azdata.connection.ConnectionProfile | undefined) {
 		super(capabilitiesService, model);
 		if (model && !isString(model)) {
 			this.groupId = model.groupId;
 			this.groupFullName = model.groupFullName;
 			this.savePassword = model.savePassword;
 			this.saveProfile = model.saveProfile;
-			this._id = model.id;
+
 			this.azureTenantId = model.azureTenantId;
-			this.azureAccount = model.azureAccount;
+
-			this.azureResourceId = model.azureResourceId;
+			// Special case setting properties to support both IConnectionProfile and azdata.connection.ConnectionProfile
-			this.azurePortalEndpoint = model.azurePortalEndpoint;
+			// It's not great that we have multiple definitions in azdata, but we can't break that right now so just
-			if (this.capabilitiesService && model.providerName) {
+			// support both at least for the time being
-				let capabilities = this.capabilitiesService.getCapabilities(model.providerName);
+			const isIConnectionProfile = 'id' in model;
+			this._id = isIConnectionProfile ? model.id : model.connectionId;
+			// TODO: @chgagnon - Should we add these properties to azdata.connection.ConnectionProfile?
+			this.azureAccount = isIConnectionProfile ? model.azureAccount : '';
+			this.azureResourceId = isIConnectionProfile ? model.azureResourceId : '';
+			this.azurePortalEndpoint = isIConnectionProfile ? model.azurePortalEndpoint : '';
+			if (this.capabilitiesService && this.providerName) {
+				let capabilities = this.capabilitiesService.getCapabilities(this.providerName);
 				if (capabilities && capabilities.connection && capabilities.connection.connectionOptions) {
 					const options = capabilities.connection.connectionOptions;
 					let appNameOption = options.find(option => option.specialValueType === interfaces.ConnectionOptionSpecialType.appName);

src/sql/platform/connection/common/providerConnectionInfo.ts

@@ -24,12 +24,12 @@ export class ProviderConnectionInfo extends Disposable implements azdata.Connect
 
 	public constructor(
 		protected capabilitiesService: ICapabilitiesService,
-		model: string | azdata.IConnectionProfile | undefined
+		model: string | azdata.IConnectionProfile | azdata.connection.ConnectionProfile | undefined
 	) {
 		super();
 		// we can't really do a whole lot if we don't have a provider
-		if (isString(model) || (model && model.providerName)) {
+		if (model) {
-			this.providerName = isString(model) ? model : model.providerName;
+			this.providerName = isString(model) ? model : 'providerName' in model ? model.providerName : model.providerId;
 
 			if (!isString(model)) {
 				if (model.options && this.serverCapabilities) {
@@ -56,7 +56,7 @@ export class ProviderConnectionInfo extends Disposable implements azdata.Connect
 	 *
 	 * This handles the case where someone hasn't passed in a valid property bag, but doesn't cause errors when
 	 */
-	private updateSpecialValueType(typeName: SettableProperty, model: azdata.IConnectionProfile): void {
+	private updateSpecialValueType(typeName: SettableProperty, model: azdata.IConnectionProfile | azdata.connection.ConnectionProfile): void {
 		if (!this[typeName]) {
 			this[typeName] = model[typeName]!;
 		}

src/sql/platform/connection/test/common/connectionProfile.test.ts

@@ -15,7 +15,7 @@ suite('SQL ConnectionProfileInfo tests', () => {
 	let msSQLCapabilities: ConnectionProviderProperties;
 	let capabilitiesService: TestCapabilitiesService;
 
-	let connectionProfile: IConnectionProfile = {
+	let iConnectionProfile: IConnectionProfile = {
 		connectionName: 'new name',
 		serverName: 'new server',
 		databaseName: 'database',
@@ -33,6 +33,22 @@ suite('SQL ConnectionProfileInfo tests', () => {
 		id: undefined!
 	};
 
+	let connectionProfile: azdata.connection.ConnectionProfile = {
+		connectionName: 'new name',
+		serverName: 'new server',
+		databaseName: 'database',
+		userName: 'user',
+		password: 'password',
+		authenticationType: '',
+		savePassword: true,
+		groupFullName: 'g2/g2-2',
+		groupId: 'group id',
+		providerId: mssqlProviderName,
+		options: {},
+		saveProfile: true,
+		connectionId: 'my id'
+	};
+
 	let storedProfile: IConnectionProfileStore = {
 		groupId: 'groupId',
 		id: 'id',
@@ -135,27 +151,42 @@ suite('SQL ConnectionProfileInfo tests', () => {
 	test('set properties should set the values correctly', () => {
 		let conn = new ConnectionProfile(capabilitiesService, undefined!);
 		assert.strictEqual(conn.serverName, undefined);
-		conn.connectionName = connectionProfile.connectionName!;
+		conn.connectionName = iConnectionProfile.connectionName!;
-		conn.serverName = connectionProfile.serverName;
+		conn.serverName = iConnectionProfile.serverName;
-		conn.databaseName = connectionProfile.databaseName!;
+		conn.databaseName = iConnectionProfile.databaseName!;
-		conn.authenticationType = connectionProfile.authenticationType;
+		conn.authenticationType = iConnectionProfile.authenticationType;
-		conn.password = connectionProfile.password;
+		conn.password = iConnectionProfile.password;
-		conn.userName = connectionProfile.userName;
+		conn.userName = iConnectionProfile.userName;
-		conn.groupId = connectionProfile.groupId;
+		conn.groupId = iConnectionProfile.groupId;
-		conn.groupFullName = connectionProfile.groupFullName;
+		conn.groupFullName = iConnectionProfile.groupFullName;
-		conn.savePassword = connectionProfile.savePassword;
+		conn.savePassword = iConnectionProfile.savePassword;
-		assert.strictEqual(conn.connectionName, connectionProfile.connectionName);
+		assert.strictEqual(conn.connectionName, iConnectionProfile.connectionName);
-		assert.strictEqual(conn.serverName, connectionProfile.serverName);
+		assert.strictEqual(conn.serverName, iConnectionProfile.serverName);
-		assert.strictEqual(conn.databaseName, connectionProfile.databaseName);
+		assert.strictEqual(conn.databaseName, iConnectionProfile.databaseName);
-		assert.strictEqual(conn.authenticationType, connectionProfile.authenticationType);
+		assert.strictEqual(conn.authenticationType, iConnectionProfile.authenticationType);
-		assert.strictEqual(conn.password, connectionProfile.password);
+		assert.strictEqual(conn.password, iConnectionProfile.password);
-		assert.strictEqual(conn.userName, connectionProfile.userName);
+		assert.strictEqual(conn.userName, iConnectionProfile.userName);
-		assert.strictEqual(conn.groupId, connectionProfile.groupId);
+		assert.strictEqual(conn.groupId, iConnectionProfile.groupId);
-		assert.strictEqual(conn.groupFullName, connectionProfile.groupFullName);
+		assert.strictEqual(conn.groupFullName, iConnectionProfile.groupFullName);
-		assert.strictEqual(conn.savePassword, connectionProfile.savePassword);
+		assert.strictEqual(conn.savePassword, iConnectionProfile.savePassword);
 	});
 
-	test('constructor should initialize the options given a valid model', () => {
+	test('constructor should initialize the options given a valid IConnectionProfile model', () => {
+		let conn = new ConnectionProfile(capabilitiesService, iConnectionProfile);
+
+		assert.strictEqual(conn.connectionName, iConnectionProfile.connectionName);
+		assert.strictEqual(conn.serverName, iConnectionProfile.serverName);
+		assert.strictEqual(conn.databaseName, iConnectionProfile.databaseName);
+		assert.strictEqual(conn.authenticationType, iConnectionProfile.authenticationType);
+		assert.strictEqual(conn.password, iConnectionProfile.password);
+		assert.strictEqual(conn.userName, iConnectionProfile.userName);
+		assert.strictEqual(conn.groupId, iConnectionProfile.groupId);
+		assert.strictEqual(conn.groupFullName, iConnectionProfile.groupFullName);
+		assert.strictEqual(conn.savePassword, iConnectionProfile.savePassword);
+		assert.strictEqual(conn.providerName, iConnectionProfile.providerName);
+	});
+
+	test('constructor should initialize the options given a valid azdata.connection.ConnectionProfile model', () => {
 		let conn = new ConnectionProfile(capabilitiesService, connectionProfile);
 
 		assert.strictEqual(conn.connectionName, connectionProfile.connectionName);
@@ -167,10 +198,11 @@ suite('SQL ConnectionProfileInfo tests', () => {
 		assert.strictEqual(conn.groupId, connectionProfile.groupId);
 		assert.strictEqual(conn.groupFullName, connectionProfile.groupFullName);
 		assert.strictEqual(conn.savePassword, connectionProfile.savePassword);
+		assert.strictEqual(conn.providerName, connectionProfile.providerId);
 	});
 
 	test('getOptionsKey should create a valid unique id', () => {
-		let conn = new ConnectionProfile(capabilitiesService, connectionProfile);
+		let conn = new ConnectionProfile(capabilitiesService, iConnectionProfile);
 		let expectedId = 'providerName:MSSQL|authenticationType:|databaseName:database|serverName:new server|userName:user|databaseDisplayName:database|group:group id';
 		let id = conn.getOptionsKey();
 		assert.strictEqual(id, expectedId);
@@ -196,20 +228,20 @@ suite('SQL ConnectionProfileInfo tests', () => {
 	});
 
 	test('withoutPassword should create a new instance without password', () => {
-		let conn = new ConnectionProfile(capabilitiesService, connectionProfile);
+		let conn = new ConnectionProfile(capabilitiesService, iConnectionProfile);
 		assert.notStrictEqual(conn.password, '');
 		let withoutPassword = conn.withoutPassword();
 		assert.strictEqual(withoutPassword.password, '');
 	});
 
 	test('unique id should not include password', () => {
-		let conn = new ConnectionProfile(capabilitiesService, connectionProfile);
+		let conn = new ConnectionProfile(capabilitiesService, iConnectionProfile);
 		let withoutPassword = conn.withoutPassword();
 		assert.strictEqual(withoutPassword.getOptionsKey(), conn.getOptionsKey());
 	});
 
 	test('cloneWithDatabase should create new profile with new id', () => {
-		let conn = new ConnectionProfile(capabilitiesService, connectionProfile);
+		let conn = new ConnectionProfile(capabilitiesService, iConnectionProfile);
 		let newProfile = conn.cloneWithDatabase('new db');
 		assert.notStrictEqual(newProfile.id, conn.id);
 		assert.strictEqual(newProfile.databaseName, 'new db');

src/sql/workbench/api/browser/mainThreadQueryEditor.ts

@@ -18,6 +18,7 @@ import { ConnectionProfile } from 'sql/platform/connection/common/connectionProf
 import { ILogService } from 'vs/platform/log/common/log';
 import { URI } from 'vs/base/common/uri';
 import { IQueryEditorService } from 'sql/workbench/services/queryEditor/common/queryEditorService';
+import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 
 @extHostNamedCustomer(SqlMainContext.MainThreadQueryEditor)
 export class MainThreadQueryEditor extends Disposable implements MainThreadQueryEditorShape {
@@ -32,7 +33,8 @@ export class MainThreadQueryEditor extends Disposable implements MainThreadQuery
 		@IEditorService private _editorService: IEditorService,
 		@IQueryManagementService private _queryManagementService: IQueryManagementService,
 		@ILogService private _logService: ILogService,
-		@IQueryEditorService private _queryEditorService: IQueryEditorService
+		@IQueryEditorService private _queryEditorService: IQueryEditorService,
+		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService
 	) {
 		super();
 		if (extHostContext) {
@@ -70,34 +72,25 @@ export class MainThreadQueryEditor extends Disposable implements MainThreadQuery
 		});
 	}
 
-	private static connectionProfileToIConnectionProfile(connection: azdata.connection.ConnectionProfile): IConnectionProfile {
+	public async $connectWithProfile(fileUri: string, connection: azdata.connection.ConnectionProfile): Promise<void> {
-		let profile: ConnectionProfile = new ConnectionProfile(undefined, undefined);
+		let editors = this._editorService.visibleEditorPanes.filter(resource => {
-		profile.options = connection.options;
+			return !!resource && resource.input.resource.toString() === fileUri;
-		profile.providerName = connection.options['providerName'];
-		return profile.toIConnectionProfile();
-	}
-
-	public $connectWithProfile(fileUri: string, connection: azdata.connection.ConnectionProfile): Thenable<void> {
-		return new Promise<void>(async (resolve, reject) => {
-			let editors = this._editorService.visibleEditorPanes.filter(resource => {
-				return !!resource && resource.input.resource.toString() === fileUri;
-			});
-			let editor = editors && editors.length > 0 ? editors[0] : undefined;
-
-			let options: IConnectionCompletionOptions = {
-				params: { connectionType: ConnectionType.editor, runQueryOnCompletion: RunQueryOnConnectionMode.none, input: editor ? editor.input as any : undefined },
-				saveTheConnection: false,
-				showDashboard: false,
-				showConnectionDialogOnError: false,
-				showFirewallRuleOnError: false,
-			};
-
-			let profile: IConnectionProfile = MainThreadQueryEditor.connectionProfileToIConnectionProfile(connection);
-			let connectionResult = await this._connectionManagementService.connect(profile, fileUri, options);
-			if (connectionResult && connectionResult.connected) {
-				this._logService.info(`editor ${fileUri} connected`);
-			}
 		});
+		let editor = editors && editors.length > 0 ? editors[0] : undefined;
+
+		let options: IConnectionCompletionOptions = {
+			params: { connectionType: ConnectionType.editor, runQueryOnCompletion: RunQueryOnConnectionMode.none, input: editor ? editor.input as any : undefined },
+			saveTheConnection: false,
+			showDashboard: false,
+			showConnectionDialogOnError: false,
+			showFirewallRuleOnError: false,
+		};
+
+		let profile: IConnectionProfile = new ConnectionProfile(this._capabilitiesService, connection).toIConnectionProfile();
+		let connectionResult = await this._connectionManagementService.connect(profile, fileUri, options);
+		if (connectionResult && connectionResult.connected) {
+			this._logService.info(`editor ${fileUri} connected`);
+		}
 	}
 
 	public $runQuery(fileUri: string, runCurrentQuery: boolean = true): void {

src/sql/workbench/test/browser/taskUtilities.test.ts

@@ -10,7 +10,8 @@ import { IObjectExplorerService } from 'sql/workbench/services/objectExplorer/br
 import { TestConnectionManagementService } from 'sql/platform/connection/test/common/testConnectionManagementService';
 import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
-import { TestEditorService } from 'vs/workbench/test/browser/workbenchTestServices';
+import { TestEditorInput, TestEditorService } from 'vs/workbench/test/browser/workbenchTestServices';
+import { URI } from 'vs/base/common/uri';
 
 suite('TaskUtilities', function () {
 	test('getCurrentGlobalConnection returns the selected OE server if a server or one of its children is selected', () => {
@@ -71,16 +72,17 @@ suite('TaskUtilities', function () {
 		mockConnectionManagementService.setup(x => x.isProfileConnected(TypeMoq.It.is(profile => profile === oeProfile || profile === tabProfile))).returns(() => true);
 
 		// Mock the workbench service to return the active tab connection
-		let tabConnectionUri = 'file://test_uri';
+		const tabConnectionUri = URI.file('file://test_uri');
-		mockConnectionManagementService.setup(x => x.getConnectionProfile(tabConnectionUri)).returns(() => tabProfile);
+		mockWorkbenchEditorService.setup(x => x.activeEditor).returns(() => new TestEditorInput(tabConnectionUri, 'my_type'));
+		mockConnectionManagementService.setup(x => x.getConnectionProfile(tabConnectionUri.toString(true))).returns(() => tabProfile);
 
 		// If I call getCurrentGlobalConnection, it should return the expected profile from the active tab
 		let actualProfile = TaskUtilities.getCurrentGlobalConnection(mockObjectExplorerService.object, mockConnectionManagementService.object, mockWorkbenchEditorService.object);
-		assert.strictEqual(actualProfile.databaseName, tabProfile.databaseName);
+		assert.strictEqual(actualProfile.databaseName, connectionProfile2.databaseName);
-		assert.strictEqual(actualProfile.authenticationType, tabProfile.authenticationType);
+		assert.strictEqual(actualProfile.authenticationType, connectionProfile2.authenticationType);
-		assert.strictEqual(actualProfile.password, tabProfile.password);
+		assert.strictEqual(actualProfile.password, connectionProfile2.password);
-		assert.strictEqual(actualProfile.serverName, tabProfile.serverName);
+		assert.strictEqual(actualProfile.serverName, connectionProfile2.serverName);
-		assert.strictEqual(actualProfile.userName, tabProfile.userName);
+		assert.strictEqual(actualProfile.userName, connectionProfile2.userName);
 	});
 
 	test('getCurrentGlobalConnection returns the connection from OE if there is no active tab, even if OE is not focused', () => {