Using Managed Identities in a Java Azure Function

I have been experimenting with Java Azure Functions and was having trouble figuring out how to use managed identities in my code. My goal was to use the BlobServiceClient and create a function that would list the containers in an Azure Storage Account. I wanted to create an User-Assigned Managed Identity, and manage the function's access permissions at the storage account level by assigning Storage Blob Data Reader to the identity.

What I found was the BlobServiceClient could be configured with a TokenCredential object. If you are using a System-Assigned Managed Identity, or you are running locally and have used az login, then you just need to instantiate the DefaultAzureCredentialBuilder and that will pull in the local identity. If you are using a User-Assigned Managed Identity, then you need to instantiate the ManagedIdentityCredentialBuilder and set the Client ID of your identity which will allow the API to load the correct identity.

I created this block of code to allow me to test with local, System-Assigned or User-Assigned Managed Identity. If I want to configure an User-Assigned Managed Identity, then I specific the Client ID in the environment variable USER_ASSIGNED_MANAGED_ID, and if the environment variable is not set code will pull the local or System-Assigned Managed Identity.

TokenCredential credential = null;
final String userAssignedManagedIdentity = System.getenv("USER_ASSIGNED_MANAGED_ID");
if (userAssignedManagedIdentity == null) {
    context.getLogger().info("Using default credentials");
    final AzureProfile profile = new AzureProfile(AzureEnvironment.AZURE);
    credential = new DefaultAzureCredentialBuilder()
        .authorityHost(profile.getEnvironment().getActiveDirectoryEndpoint())
        .build();
} else {
    context.getLogger().info(String.format("Using User Assigned Managed Identity - %s", userAssignedManagedIdentity));
    credential = new ManagedIdentityCredentialBuilder()
        .clientId(userAssignedManagedIdentity)
        .build();
}

You can find the full Azure Function project on my GitHub repo.