Improve support for assigning UAID. Supporting adding newly created UAID.

MicroFish91 · MicroFish91 · commit df78db482c57 · 2025-08-04T12:56:07.000-07:00
diff --git a/src/commands/appSettings/connectionSettings/durableTaskScheduler/IDTSConnectionWizardContext.ts b/src/commands/appSettings/connectionSettings/durableTaskScheduler/IDTSConnectionWizardContext.ts
@@ -8,7 +8,7 @@ import { type AzureSubscription } from "@microsoft/vscode-azureresources-api";
 import { type ConnectionType } from "../../../../constants";
 import { type DurableTaskHubResource, type DurableTaskSchedulerResource } from "../../../../tree/durableTaskScheduler/DurableTaskSchedulerClient";
 import { type DurableTaskSchedulerEmulator } from "../../../../tree/durableTaskScheduler/DurableTaskSchedulerEmulatorClient";
-import { type IFunctionAppUserAssignedIdentitiesContext } from "../../../identity/listUserAssignedIdentities/IFunctionAppUserAssignedIdentitiesContext";
+import { type ManagedIdentityAssignContext } from "../../../identity/ManagedIdentityAssignContext";
 import { type StorageConnectionType } from "../IConnectionTypesContext";
 import { type ISetConnectionSettingContext } from "../ISetConnectionSettingContext";
 
@@ -24,7 +24,7 @@ export interface IDTSConnectionWizardContext extends IActionContext, ISetConnect
     // All properties from `IDTSConnectionSetSettingsContext` apply
 }
 
-export interface IDTSAzureConnectionWizardContext extends IFunctionAppUserAssignedIdentitiesContext, IDTSConnectionWizardContext, Partial<ExecuteActivityContext> {
+export interface IDTSAzureConnectionWizardContext extends ManagedIdentityAssignContext, IDTSConnectionWizardContext, ExecuteActivityContext {
     subscription?: AzureSubscription;
 
     newDTSName?: string;
diff --git a/src/commands/appSettings/connectionSettings/durableTaskScheduler/azure/DurableTaskSchedulerListStep.ts b/src/commands/appSettings/connectionSettings/durableTaskScheduler/azure/DurableTaskSchedulerListStep.ts
@@ -3,12 +3,12 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { CommonRoleDefinitions, createRoleId, LocationListStep, parseAzureResourceId, RoleAssignmentExecuteStep, type ILocationWizardContext, type Role } from '@microsoft/vscode-azext-azureutils';
+import { CommonRoleDefinitions, createAuthorizationManagementClient, createRoleId, LocationListStep, parseAzureResourceId, RoleAssignmentExecuteStep, uiUtils, type ILocationWizardContext, type Role } from '@microsoft/vscode-azext-azureutils';
 import { AzureWizardPromptStep, nonNullProp, type AzureWizardExecuteStep, type IAzureQuickPickItem, type IWizardOptions } from '@microsoft/vscode-azext-utils';
 import { localSettingsDescription } from '../../../../../constants-nls';
 import { localize } from '../../../../../localize';
 import { HttpDurableTaskSchedulerClient, type DurableTaskSchedulerClient, type DurableTaskSchedulerResource } from '../../../../../tree/durableTaskScheduler/DurableTaskSchedulerClient';
-import { FunctionAppUserAssignedIdentitiesListStep } from '../../../../identity/listUserAssignedIdentities/FunctionAppUserAssignedIdentitiesListStep';
+import { FunctionAppUserAssignedIdentitiesListStep } from '../../../../identity/FunctionAppUserAssignedIdentitiesListStep';
 import { type IDTSAzureConnectionWizardContext } from '../IDTSConnectionWizardContext';
 import { DurableTaskHubListStep } from './DurableTaskHubListStep';
 import { DurableTaskSchedulerCreateStep } from './DurableTaskSchedulerCreateStep';
@@ -58,22 +58,35 @@ export class DurableTaskSchedulerListStep<T extends IDTSAzureConnectionWizardCon
             roleDefinitionName: CommonRoleDefinitions.durableTaskDataContributor.roleName,
         };
 
-        const identitiesListStep = new FunctionAppUserAssignedIdentitiesListStep(dtsContributorRole /** targetRole */);
-        promptSteps.push(identitiesListStep);
-        executeSteps.push(new RoleAssignmentExecuteStep(getDTSRoleAssignmentCallback(context, identitiesListStep, dtsContributorRole)));
+        promptSteps.push(new FunctionAppUserAssignedIdentitiesListStep(dtsContributorRole /** targetRole */, { identityAssignStepPriority: 180 }));
+        executeSteps.push(new RoleAssignmentExecuteStep(getDTSRoleAssignmentCallback(context, dtsContributorRole), { priority: 190 }));
 
         return { promptSteps, executeSteps };
 
-        function getDTSRoleAssignmentCallback(context: T, functionAppIdentitiesListStep: FunctionAppUserAssignedIdentitiesListStep<T>, role: Role): () => Role[] {
-            return () => {
+        function getDTSRoleAssignmentCallback(context: T, role: Role): () => Promise<Role[]> {
+            return async () => {
                 const roleAssignment: Role = {
                     ...role,
                     // This id may be missing when the role is initially passed in,
                     // but by the time we run the step, we should have the populated id ready.
                     scopeId: context.dts?.id,
                 };
 
-                return functionAppIdentitiesListStep.hasIdentityWithTargetRole ? [] : [roleAssignment];
+                if (!roleAssignment.scopeId) {
+                    return [];
+                }
+
+                const amClient = await createAuthorizationManagementClient(context);
+                const roleAssignments = await uiUtils.listAllIterator(amClient.roleAssignments.listForScope(
+                    roleAssignment.scopeId,
+                    {
+                        // $filter=principalId eq {id}
+                        filter: `principalId eq '{${context.managedIdentity?.principalId}}'`,
+                    }
+                ));
+
+                const hasRoleAssignment = roleAssignments.some(r => !!r.roleDefinitionId?.endsWith(role.roleDefinitionId));
+                return hasRoleAssignment ? [] : [roleAssignment];
             };
         }
     }
diff --git a/src/commands/identity/FunctionAppUserAssignedIdentitiesListStep.ts b/src/commands/identity/FunctionAppUserAssignedIdentitiesListStep.ts
@@ -5,11 +5,12 @@
 
 import { type ManagedServiceIdentityClient } from '@azure/arm-msi';
 import { type ParsedSite } from '@microsoft/vscode-azext-azureappservice';
-import { createAuthorizationManagementClient, createManagedServiceIdentityClient, parseAzureResourceId, uiUtils, type ParsedAzureResourceId, type Role } from '@microsoft/vscode-azext-azureutils';
-import { ActivityChildItem, ActivityChildType, activityInfoContext, activityInfoIcon, AzureWizardPromptStep, createContextValue, nonNullProp, prependOrInsertAfterLastInfoChild, type ActivityInfoChild, type IAzureQuickPickItem } from '@microsoft/vscode-azext-utils';
-import { ext } from '../../../extensionVariables';
-import { localize } from '../../../localize';
-import { type IFunctionAppUserAssignedIdentitiesContext } from './IFunctionAppUserAssignedIdentitiesContext';
+import { createAuthorizationManagementClient, createManagedServiceIdentityClient, parseAzureResourceId, uiUtils, UserAssignedIdentityListStep, type ParsedAzureResourceId, type Role } from '@microsoft/vscode-azext-azureutils';
+import { ActivityChildItem, ActivityChildType, activityInfoContext, activityInfoIcon, AzureWizardPromptStep, createContextValue, nonNullProp, prependOrInsertAfterLastInfoChild, type ActivityInfoChild, type IAzureQuickPickItem, type IWizardOptions } from '@microsoft/vscode-azext-utils';
+import { ext } from '../../extensionVariables';
+import { localize } from '../../localize';
+import { type ManagedIdentityAssignContext } from './ManagedIdentityAssignContext';
+import { ManagedIdentityAssignStep } from './ManagedIdentityAssignStep';
 
 /**
  * Wizard step to select a user-assigned managed identity from the parsed site of a function app.
@@ -20,30 +21,22 @@ import { type IFunctionAppUserAssignedIdentitiesContext } from './IFunctionAppUs
  *
  * @populates `context.managedIdentity`
  */
-export class FunctionAppUserAssignedIdentitiesListStep<T extends IFunctionAppUserAssignedIdentitiesContext> extends AzureWizardPromptStep<T> {
+export class FunctionAppUserAssignedIdentitiesListStep<T extends ManagedIdentityAssignContext> extends AzureWizardPromptStep<T> {
     private _msiClient: ManagedServiceIdentityClient;
-    private _hasTargetRole?: boolean;
 
-    constructor(readonly targetRole?: Role) {
+    constructor(
+        readonly targetRole?: Role,
+        readonly options?: { identityAssignStepPriority?: number },
+    ) {
         super();
     }
 
-    /**
-     * Indicates whether there is at least one user-assigned identity on the function app with the provided role.
-     * If no role is provided, or if the step has not yet been run, this will return `undefined`.
-     */
-    get hasIdentityWithTargetRole(): boolean | undefined {
-        return this._hasTargetRole;
-    }
-
     // Verify if any of the existing user assigned identities for the function app have the required role already
     public async configureBeforePrompt(context: T): Promise<void> {
-        if (!this.targetRole || !this.targetRole?.scopeId) {
-            this._hasTargetRole = undefined;
+        if (!this.targetRole?.scopeId) {
             return;
         }
 
-        this._hasTargetRole = false;
         this._msiClient ??= await createManagedServiceIdentityClient(context);
         const amClient = await createAuthorizationManagementClient(context);
 
@@ -52,6 +45,7 @@ export class FunctionAppUserAssignedIdentitiesListStep<T extends IFunctionAppUse
         const identityIds: string[] = Object.keys(site.identity?.userAssignedIdentities ?? {}) ?? [];
         context.telemetry.properties.functionAppUserAssignedIdentityCount = String(identityIds.length);
 
+        let hasTargetRole: boolean = false;
         for (const identityId of identityIds) {
             const uaid = site.identity?.userAssignedIdentities?.[identityId];
             const roleAssignments = await uiUtils.listAllIterator(amClient.roleAssignments.listForScope(
@@ -65,20 +59,21 @@ export class FunctionAppUserAssignedIdentitiesListStep<T extends IFunctionAppUse
             if (roleAssignments.some(r => !!r.roleDefinitionId?.endsWith(role.roleDefinitionId))) {
                 const parsedIdentity = parseAzureResourceId(identityId);
                 context.managedIdentity = await this._msiClient.userAssignedIdentities.get(parsedIdentity.resourceGroup, parsedIdentity.resourceName);
-                this._hasTargetRole = true;
+                hasTargetRole = true;
                 break;
             }
         }
 
-        context.telemetry.properties.functionAppHasIdentityWithTargetRole = String(this.hasIdentityWithTargetRole);
+        context.telemetry.properties.functionAppHasIdentityWithTargetRole = String(hasTargetRole);
 
-        if (this.hasIdentityWithTargetRole) {
+        if (hasTargetRole) {
             prependOrInsertAfterLastInfoChild(context,
                 new ActivityChildItem({
+                    stepId: this.id,
                     label: localize('useIdentityWithRole', 'Use identity "{0}" with role "{1}"', context.managedIdentity?.name, this.targetRole.roleDefinitionName),
                     contextValue: createContextValue(['functionAppUserAssignedIdentitiesListStepItem', activityInfoContext]),
                     activityType: ActivityChildType.Info,
-                    iconPath: activityInfoIcon
+                    iconPath: activityInfoIcon,
                 }) as ActivityInfoChild,
             );
             ext.outputChannel.appendLog(localize('foundIdentity', 'Located existing user assigned identity "{0}" with role "{1}".', context.managedIdentity?.name, this.targetRole.roleDefinitionName));
@@ -89,12 +84,14 @@ export class FunctionAppUserAssignedIdentitiesListStep<T extends IFunctionAppUse
 
     public async prompt(context: T): Promise<void> {
         const site: ParsedSite = nonNullProp(context, 'site');
-        const identityId: string = (await context.ui.showQuickPick(await this.getPicks(site), {
+        const identityId: string | undefined = (await context.ui.showQuickPick(await this.getPicks(site), {
             placeHolder: localize('selectFunctionAppIdentity', 'Select a function app identity for new role assignments'),
-            // Todo: Remove when create + assign is implemented
-            noPicksMessage: localize('noUserAssignedIdentities', 'No identities found. Add a user assigned identity to the function app before proceeding.'),
         })).data;
 
+        if (!identityId) {
+            return;
+        }
+
         const parsedIdentity: ParsedAzureResourceId = parseAzureResourceId(identityId);
         this._msiClient ??= await createManagedServiceIdentityClient(context);
 
@@ -106,14 +103,34 @@ export class FunctionAppUserAssignedIdentitiesListStep<T extends IFunctionAppUse
         return !context.managedIdentity;
     }
 
-    private async getPicks(site: ParsedSite): Promise<IAzureQuickPickItem<string>[]> {
-        return Object.keys(site.identity?.userAssignedIdentities ?? {}).map((id) => {
-            const parsedResource: ParsedAzureResourceId = parseAzureResourceId(id);
-            return {
-                label: parsedResource.resourceName,
-                description: parsedResource.resourceGroup,
-                data: id,
-            };
-        });
+    public async getSubWizard(context: T): Promise<IWizardOptions<T> | undefined> {
+        if (context.managedIdentity) {
+            return undefined;
+        }
+
+        return {
+            promptSteps: [new UserAssignedIdentityListStep()],
+            executeSteps: [new ManagedIdentityAssignStep({ priority: this.options?.identityAssignStepPriority })],
+        };
+    }
+
+    private async getPicks(site: ParsedSite): Promise<IAzureQuickPickItem<string | undefined>[]> {
+        const picks: IAzureQuickPickItem<string | undefined>[] = [{
+            label: localize('assignIdentity', '$(plus) Assign new user-assigned identity'),
+            data: undefined,
+        }];
+
+        return picks.concat(
+            Object
+                .keys(site.identity?.userAssignedIdentities ?? {})
+                .map((id) => {
+                    const parsedResource: ParsedAzureResourceId = parseAzureResourceId(id);
+                    return {
+                        label: parsedResource.resourceName,
+                        description: parsedResource.resourceGroup,
+                        data: id,
+                    };
+                }),
+        );
     }
 }
diff --git a/src/commands/identity/ManagedIdentityAssignStep.ts b/src/commands/identity/ManagedIdentityAssignStep.ts
@@ -12,7 +12,12 @@ import { localize } from "../../localize";
 import { type ManagedIdentityAssignContext } from "./ManagedIdentityAssignContext";
 
 export class ManagedIdentityAssignStep extends AzureWizardExecuteStep<ManagedIdentityAssignContext> {
-    public priority: number = 500;
+    public priority: number;
+
+    constructor(options?: { priority?: number }) {
+        super();
+        this.priority = options?.priority ?? 500;
+    }
 
     public async execute(context: ManagedIdentityAssignContext, _progress: Progress<{ message?: string | undefined; increment?: number | undefined; }>): Promise<void> {
         const site: ParsedSite = nonNullProp(context, 'site');
diff --git a/src/commands/identity/listUserAssignedIdentities/IFunctionAppUserAssignedIdentitiesContext.ts b/src/commands/identity/listUserAssignedIdentities/IFunctionAppUserAssignedIdentitiesContext.ts
