ckaczor/azuredatastudio
1
0
Fork 0
mirror of https://github.com/ckaczor/azuredatastudio.git synced 2026-01-19 01:25:36 -05:00
Code Issues Packages Projects Releases Wiki Activity

Enable SQL Auth Provider support (#21903)

Browse Source
This commit is contained in:
Cheena Malhotra
2023-03-03 12:49:01 -08:00
committed by GitHub
parent 0ac6f40559
commit aa350f7e49
25 changed files with 198 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/sql/workbench/api/browser/mainThreadAccountManagement.ts
View File

@@ -72,7 +72,6 @@ export class MainThreadAccountManagement extends Disposable implements MainThrea
clear(accountKey: azdata.AccountKey): Thenable<void> {
return self._proxy.$clear(handle, accountKey);
},
getSecurityToken(account: azdata.Account, resource: azdata.AzureResource): Thenable<{}> {
return self._proxy.$getSecurityToken(account, resource);
},

1
src/sql/workbench/api/common/extHostAccountManagement.ts
View File

@@ -114,7 +114,6 @@ export class ExtHostAccountManagement extends ExtHostAccountManagementShape {
});
}
public get onDidChangeAccounts(): Event<azdata.DidChangeAccountsParams> {
return this._onDidChangeAccounts.event;
}

3
src/sql/workbench/api/common/sqlExtHostTypes.ts
View File

@@ -474,7 +474,8 @@ export enum AzureResource {
AzureLogAnalytics = 8,
AzureStorage = 9,
AzureKusto = 10,
PowerBi = 11
PowerBi = 11,
Custom = 12 // Handles custom resource URIs as received from server endpoint.
}
export class TreeItem extends vsExtTypes.TreeItem {

11
src/sql/workbench/contrib/backup/browser/backup.contribution.ts
View File

@@ -20,6 +20,7 @@ import { ConnectionContextKey } from 'sql/workbench/services/connection/common/c
import { ServerInfoContextKey } from 'sql/workbench/services/connection/common/serverInfoContextKey';
import { ServicesAccessor, IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
import { DatabaseEngineEdition } from 'sql/workbench/api/common/sqlExtHostTypes';
import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
new BackupAction().registerTask();
@@ -30,7 +31,9 @@ CommandsRegistry.registerCommand({
handler: async (accessor, args: TreeViewItemHandleArg) => {
if (args.$treeItem?.payload) {
const commandService = accessor.get(ICommandService);
return commandService.executeCommand(BackupAction.ID, args.$treeItem.payload);
const connectionService = accessor.get(IConnectionManagementService);
let payload = await connectionService.fixProfile(args.$treeItem.payload);
return commandService.executeCommand(BackupAction.ID, payload);
}
}
});
@@ -69,9 +72,11 @@ MenuRegistry.appendMenuItem(MenuId.ObjectExplorerItemContext, {
// dashboard explorer
const ExplorerBackUpActionID = 'explorer.backup';
CommandsRegistry.registerCommand(ExplorerBackUpActionID, (accessor, context: ManageActionContext) => {
CommandsRegistry.registerCommand(ExplorerBackUpActionID, async (accessor, context: ManageActionContext) => {
const commandService = accessor.get(ICommandService);
return commandService.executeCommand(BackupAction.ID, context.profile);
const connectionService = accessor.get(IConnectionManagementService);
let profile = await connectionService.fixProfile(context.profile);
return commandService.executeCommand(BackupAction.ID, profile);
});
MenuRegistry.appendMenuItem(MenuId.ExplorerWidgetContext, {

6
src/sql/workbench/contrib/dashboard/browser/dashboardActions.ts
View File

@@ -42,7 +42,8 @@ CommandsRegistry.registerCommand({
showConnectionDialogOnError: true,
showFirewallRuleOnError: true
};
let profile = new ConnectionProfile(capabilitiesService, args.$treeItem.payload);
let payload = await connectionService.fixProfile(args.$treeItem.payload);
let profile = new ConnectionProfile(capabilitiesService, payload);
let uri = generateUri(profile, 'dashboard');
return connectionService.connect(new ConnectionProfile(capabilitiesService, args.$treeItem.payload), uri, options);
}
@@ -96,7 +97,8 @@ export class OEManageConnectionAction extends Action {
if (actionContext instanceof ObjectExplorerActionsContext) {
// Must use a real connection profile for this action due to lookup
connectionProfile = ConnectionProfile.fromIConnectionProfile(this._capabilitiesService, actionContext.connectionProfile);
let updatedIConnProfile = await this._connectionManagementService.fixProfile(actionContext.connectionProfile);
connectionProfile = ConnectionProfile.fromIConnectionProfile(this._capabilitiesService, updatedIConnProfile);
if (!actionContext.isConnectionNode) {
treeNode = await getTreeNode(actionContext, this._objectExplorerService);
if (TreeUpdateUtils.isDatabaseNode(treeNode)) {

7
src/sql/workbench/contrib/notebook/browser/notebook.contribution.ts
View File

@@ -67,6 +67,7 @@ import { IViewDescriptorService, ViewContainerLocation } from 'vs/workbench/comm
import { ToggleTabFocusModeAction } from 'vs/editor/contrib/toggleTabFocusMode/browser/toggleTabFocusMode';
import { ActiveEditorContext } from 'vs/workbench/common/contextkeys';
import { ILanguageService } from 'vs/editor/common/languages/language';
import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
Registry.as<IEditorFactoryRegistry>(EditorExtensions.EditorFactory)
.registerEditorSerializer(FileNotebookInput.ID, FileNoteBookEditorSerializer);
@@ -101,9 +102,11 @@ const DE_NEW_NOTEBOOK_COMMAND_ID = 'dataExplorer.newNotebook';
// New Notebook
CommandsRegistry.registerCommand({
id: DE_NEW_NOTEBOOK_COMMAND_ID,
handler: (accessor, args: TreeViewItemHandleArg) => {
handler: async (accessor, args: TreeViewItemHandleArg) => {
const instantiationService = accessor.get(IInstantiationService);
const connectedContext: ConnectedContext = { connectionProfile: args.$treeItem.payload };
const connectionService = accessor.get(IConnectionManagementService);
let payload = await connectionService.fixProfile(args.$treeItem.payload);
const connectedContext: ConnectedContext = { connectionProfile: payload };
return instantiationService.createInstance(NewNotebookAction, NewNotebookAction.ID, NewNotebookAction.LABEL).run({ connectionProfile: connectedContext.connectionProfile, isConnectionNode: false, nodeInfo: undefined });
}
});

8
src/sql/workbench/contrib/objectExplorer/test/browser/connectionTreeActions.test.ts
View File

@@ -43,6 +43,7 @@ import { TestEditorService } from 'vs/workbench/test/browser/workbenchTestServic
import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
import { TestDialogService } from 'vs/platform/dialogs/test/common/testDialogService';
import { IDialogService } from 'vs/platform/dialogs/common/dialogs';
import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
suite('SQL Connection Tree Action tests', () => {
let errorMessageService: TypeMoq.Mock<TestErrorMessageService>;
@@ -61,7 +62,7 @@ suite('SQL Connection Tree Action tests', () => {
errorMessageService.setup(x => x.showDialog(Severity.Error, TypeMoq.It.isAnyString(), TypeMoq.It.isAnyString())).returns(() => nothing);
});
function createConnectionManagementService(isConnectedReturnValue: boolean, profileToReturn: ConnectionProfile): TypeMoq.Mock<TestConnectionManagementService> {
function createConnectionManagementService(isConnectedReturnValue: boolean, profileToReturn: IConnectionProfile): TypeMoq.Mock<TestConnectionManagementService> {
let connectionManagementService = TypeMoq.Mock.ofType(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
connectionManagementService.callBase = true;
connectionManagementService.setup(x => x.isConnected(undefined, TypeMoq.It.isAny())).returns(() => isConnectedReturnValue);
@@ -77,6 +78,7 @@ suite('SQL Connection Tree Action tests', () => {
connectionManagementService.setup(x => x.deleteConnectionGroup(TypeMoq.It.isAny())).returns(() => Promise.resolve(true));
connectionManagementService.setup(x => x.deleteConnection(TypeMoq.It.isAny())).returns(() => Promise.resolve(true));
connectionManagementService.setup(x => x.getConnectionProfile(TypeMoq.It.isAny())).returns(() => profileToReturn);
connectionManagementService.setup(x => x.fixProfile(TypeMoq.It.isAny())).returns(() => new Promise<IConnectionProfile>((resolve, reject) => resolve(profileToReturn)));
connectionManagementService.setup(x => x.showEditConnectionDialog(TypeMoq.It.isAny())).returns(() => new Promise<void>((resolve, reject) => resolve()));
return connectionManagementService;
}
@@ -117,7 +119,7 @@ suite('SQL Connection Tree Action tests', () => {
test('ManageConnectionAction - test if connect is called for manage action if not already connected', () => {
let isConnectedReturnValue: boolean = false;
let connection: ConnectionProfile = new ConnectionProfile(capabilitiesService, {
let connection: IConnectionProfile = new ConnectionProfile(capabilitiesService, {
connectionName: 'Test',
savePassword: false,
groupFullName: 'testGroup',
@@ -197,7 +199,7 @@ suite('SQL Connection Tree Action tests', () => {
viewsService);
let actionContext = new ObjectExplorerActionsContext();
actionContext.connectionProfile = connection.toIConnectionProfile();
actionContext.connectionProfile = connection;
actionContext.isConnectionNode = true;
return manageConnectionAction.run(actionContext).then(() => {
connectionManagementService.verify(x => x.connect(TypeMoq.It.isAny(), undefined, TypeMoq.It.isAny(), undefined), TypeMoq.Times.once());

3
src/sql/workbench/contrib/query/browser/queryActions.ts
View File

@@ -182,7 +182,8 @@ CommandsRegistry.registerCommand({
showConnectionDialogOnError: true,
showFirewallRuleOnError: true
};
return connectionService.connect(new ConnectionProfile(capabilitiesService, args.$treeItem.payload), owner.uri, options);
let payload = await connectionService.fixProfile(args.$treeItem.payload);
return connectionService.connect(new ConnectionProfile(capabilitiesService, payload), owner.uri, options);
}
return true;
}

17
src/sql/workbench/contrib/restore/browser/restore.contribution.ts
View File

@@ -19,6 +19,7 @@ import { ManageActionContext } from 'sql/workbench/browser/actions';
import { ItemContextKey } from 'sql/workbench/contrib/dashboard/browser/widgets/explorer/explorerContext';
import { ServerInfoContextKey } from 'sql/workbench/services/connection/common/serverInfoContextKey';
import { DatabaseEngineEdition } from 'sql/workbench/api/common/sqlExtHostTypes';
import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
new RestoreAction().registerTask();
@@ -29,7 +30,9 @@ CommandsRegistry.registerCommand({
handler: async (accessor, args: TreeViewItemHandleArg) => {
if (args.$treeItem?.payload) {
const commandService = accessor.get(ICommandService);
return commandService.executeCommand(RestoreAction.ID, args.$treeItem.payload);
const connectionService = accessor.get(IConnectionManagementService);
let payload = await connectionService.fixProfile(args.$treeItem.payload);
return commandService.executeCommand(RestoreAction.ID, payload);
}
}
});
@@ -51,9 +54,11 @@ MenuRegistry.appendMenuItem(MenuId.DataExplorerContext, {
const OE_RESTORE_COMMAND_ID = 'objectExplorer.restore';
CommandsRegistry.registerCommand({
id: OE_RESTORE_COMMAND_ID,
handler: (accessor, args: ObjectExplorerActionsContext) => {
handler: async (accessor, args: ObjectExplorerActionsContext) => {
const commandService = accessor.get(ICommandService);
return commandService.executeCommand(RestoreAction.ID, args.connectionProfile);
const connectionService = accessor.get(IConnectionManagementService);
let profile = await connectionService.fixProfile(args.connectionProfile);
return commandService.executeCommand(RestoreAction.ID, profile);
}
});
@@ -69,9 +74,11 @@ MenuRegistry.appendMenuItem(MenuId.ObjectExplorerItemContext, {
});
const ExplorerRestoreActionID = 'explorer.restore';
CommandsRegistry.registerCommand(ExplorerRestoreActionID, (accessor, context: ManageActionContext) => {
CommandsRegistry.registerCommand(ExplorerRestoreActionID, async (accessor, context: ManageActionContext) => {
const commandService = accessor.get(ICommandService);
return commandService.executeCommand(RestoreAction.ID, context.profile);
const connectionService = accessor.get(IConnectionManagementService);
let profile = await connectionService.fixProfile(context.profile);
return commandService.executeCommand(RestoreAction.ID, profile);
});
MenuRegistry.appendMenuItem(MenuId.ExplorerWidgetContext, {

24
src/sql/workbench/contrib/scripting/browser/scriptingActions.ts
View File

@@ -49,7 +49,9 @@ CommandsRegistry.registerCommand({
const scriptingService = accessor.get(IScriptingService);
const errorMessageService = accessor.get(IErrorMessageService);
const progressService = accessor.get(IProgressService);
const profile = new ConnectionProfile(capabilitiesService, args.$treeItem.payload);
const connectionService = accessor.get(IConnectionManagementService);
let payload = await connectionService.fixProfile(args.$treeItem.payload);
const profile = new ConnectionProfile(capabilitiesService, payload);
const baseContext: BaseActionContext = {
profile: profile,
object: oeShimService.getNodeInfoForTreeItem(args.$treeItem)!.metadata
@@ -73,7 +75,9 @@ CommandsRegistry.registerCommand({
const scriptingService = accessor.get(IScriptingService);
const errorMessageService = accessor.get(IErrorMessageService);
const progressService = accessor.get(IProgressService);
const profile = new ConnectionProfile(capabilitiesService, args.$treeItem.payload);
const connectionService = accessor.get(IConnectionManagementService);
let payload = await connectionService.fixProfile(args.$treeItem.payload);
const profile = new ConnectionProfile(capabilitiesService, payload);
const baseContext: BaseActionContext = {
profile: profile,
object: oeShimService.getNodeInfoForTreeItem(args.$treeItem)!.metadata
@@ -97,7 +101,9 @@ CommandsRegistry.registerCommand({
const scriptingService = accessor.get(IScriptingService);
const progressService = accessor.get(IProgressService);
const errorMessageService = accessor.get(IErrorMessageService);
const profile = new ConnectionProfile(capabilitiesService, args.$treeItem.payload);
const connectionService = accessor.get(IConnectionManagementService);
let payload = await connectionService.fixProfile(args.$treeItem.payload);
const profile = new ConnectionProfile(capabilitiesService, payload);
const baseContext: BaseActionContext = {
profile: profile,
object: oeShimService.getNodeInfoForTreeItem(args.$treeItem)!.metadata
@@ -121,7 +127,9 @@ CommandsRegistry.registerCommand({
const scriptingService = accessor.get(IScriptingService);
const progressService = accessor.get(IProgressService);
const errorMessageService = accessor.get(IErrorMessageService);
const profile = new ConnectionProfile(capabilitiesService, args.$treeItem.payload);
const connectionService = accessor.get(IConnectionManagementService);
let payload = await connectionService.fixProfile(args.$treeItem.payload);
const profile = new ConnectionProfile(capabilitiesService, payload);
const baseContext: BaseActionContext = {
profile: profile,
object: oeShimService.getNodeInfoForTreeItem(args.$treeItem)!.metadata
@@ -145,7 +153,9 @@ CommandsRegistry.registerCommand({
const scriptingService = accessor.get(IScriptingService);
const progressService = accessor.get(IProgressService);
const errorMessageService = accessor.get(IErrorMessageService);
const profile = new ConnectionProfile(capabilitiesService, args.$treeItem.payload);
const connectionService = accessor.get(IConnectionManagementService);
let payload = await connectionService.fixProfile(args.$treeItem.payload);
const profile = new ConnectionProfile(capabilitiesService, payload);
const baseContext: BaseActionContext = {
profile: profile,
object: oeShimService.getNodeInfoForTreeItem(args.$treeItem)!.metadata
@@ -169,7 +179,9 @@ CommandsRegistry.registerCommand({
const scriptingService = accessor.get(IScriptingService);
const progressService = accessor.get(IProgressService);
const errorMessageService = accessor.get(IErrorMessageService);
const profile = new ConnectionProfile(capabilitiesService, args.$treeItem.payload);
const connectionService = accessor.get(IConnectionManagementService);
let payload = await connectionService.fixProfile(args.$treeItem.payload);
const profile = new ConnectionProfile(capabilitiesService, payload);
const baseContext: BaseActionContext = {
profile: profile,
object: oeShimService.getNodeInfoForTreeItem(args.$treeItem)!.metadata

18
src/sql/workbench/services/connection/browser/connectionManagementService.ts
View File

@@ -403,6 +403,24 @@ export class ConnectionManagementService extends Disposable implements IConnecti
return this.tryConnect(connection, input, options);
}
public async fixProfile(profile?: interfaces.IConnectionProfile): Promise<interfaces.IConnectionProfile> {
if (profile) {
if (profile.authenticationType !== undefined && profile.authenticationType === '') {
// we need to set auth type here, because it's value is part of the session key
profile.authenticationType = this.getDefaultAuthenticationTypeId(profile.providerName);
}
// If this is Azure MFA Authentication, fix username to azure Account user. Falls back to current user name.
// This is required, as by default, server login / administrator is the username.
if (profile.authenticationType === 'AzureMFA') {
let accounts = await this._accountManagementService?.getAccounts();
profile.userName = accounts?.find(a => a.key.accountId === profile.azureAccount)?.displayInfo.displayName
?? profile.userName;
}
}
return profile;
}
/**
* If there's already a connection for given profile and purpose, returns the ownerUri for the connection
* otherwise tries to make a connection and returns the owner uri when connection is complete

8
src/sql/workbench/services/objectExplorer/browser/objectExplorerViewTreeShim.ts
View File

@@ -43,13 +43,13 @@ export class OEShimService extends Disposable implements IOEShimService {
@IObjectExplorerService private oe: IObjectExplorerService,
@IConnectionManagementService private cm: IConnectionManagementService,
@ICapabilitiesService private capabilities: ICapabilitiesService,
@IConfigurationService private configurationService: IConfigurationService
) {
super();
}
private async createSession(viewId: string, providerId: string, node: ITreeItem): Promise<string> {
let connProfile = new ConnectionProfile(this.capabilities, node.payload);
let payload = await this.cm.fixProfile(node.payload);
let connProfile = new ConnectionProfile(this.capabilities, payload);
connProfile.saveProfile = false;
if (this.cm.providerRegistered(providerId)) {
connProfile = await this.connectOrPrompt(connProfile);
@@ -119,9 +119,7 @@ export class OEShimService extends Disposable implements IOEShimService {
public async getChildren(node: ITreeItem, viewId: string): Promise<ITreeItem[]> {
if (node.payload) {
if (node.payload.authenticationType !== undefined && node.payload.authenticationType === '') {
node.payload.authenticationType = this.getDefaultAuthenticationType(this.configurationService); // we need to set auth type here, because it's value is part of the session key
}
node.payload = await this.cm.fixProfile(node.payload);
if (node.sessionId === undefined) {
node.sessionId = await this.createSession(viewId, node.childProvider!, node);