How to authenticate an Azure identity against a Postgres instance using Spring Boot

Introduction

This post will demonstrate how you can authenticate your Spring Boot featured application against an Azure AD integrated Postgres instance.

To accomplish that, we are going to use the Azure Identity Library and create our own DataSource type by extending the HikariDataSource.

😴 TLDR? You only came here for the code? Fair enough, you'll find it further down at "Putting everything together".

It completes previous posts I have published that can be found below. If you are new to this topic, I recommend reading my article to gain some background knowledge.

Before we get to see some code, please note that only Azure Database for PostgreSQL single server provides Azure AD integration. So you won't be able to follow this post by using a Postgres flexible server.

☝🏻 Only PostgreSQL single server provides Azure AD integration. According to Microsoft, flexible server will follow in the future.

I assume you know how to set up and enable Azure AD integration on an Azure Postgres instance. So I am not going to cover that here.

Let's dive in 🤿

Azure Identity Library

As stated in my previous article (you read them right?), the access token is getting sent in the password field. So you first need a way to retrieve an access token.

We are going to use the com.azure.azure-identity library for that purpose, which is part of the Azure SDK for Java. Assuming you are using maven, adjust your pom.xml file to match mine as shown below.

</project>
  <parent> 
    <groupId>org.springframework.boot</groupId>  
    <artifactId>spring-boot-starter-parent</artifactId>  
    <version>2.6.4</version>  
    <relativePath/>  
  </parent>  
  
  ...
  
  <properties> 
    <java.version>11</java.version>
    <azure.version>3.14.0</azure.version> 
  </properties>  
  
  <dependencies> 

    <dependency> 
      <groupId>com.azure</groupId>  
      <artifactId>azure-identity</artifactId> 
    </dependency>  
      
	...
      
  <dependencyManagement> 
    <dependencies> 
      <dependency> 
        <groupId>com.azure.spring</groupId>  
        <artifactId>azure-spring-boot-bom</artifactId>  
        <version>${azure.version}</version>  
        <type>pom</type>  
        <scope>import</scope> 
      </dependency> 
    </dependencies> 
  </dependencyManagement> 
  
  ...
  
</project>

☝🏼 As of writing the latest version was 3.14.0. You might want to check search.maven.org first to grab the latest and greatest version!

After downloading the dependencies we can move to the next step!

System & user-assigned managed identities

Let's assume you'd like to deploy your Spring Boot application to Azure App Service. You are going to have two choices:

• System-assigned managed identities
• User-assigned managed identities

☝🏼 System-assigned and user-assigned managed identities only differ in their lifetime. A system assigned identities lifetime is bound to the resource it got enabled for. When you delete the resource, you are deleting the identity. A user-assigned managed identity on the other hand has its own lifetime. However, both managed identities provide the benefit of not having to care about the credentials itself

Both can be enabled from the Identity blade. For the system assigned just flip the switch to on and save! Needless to say, you'd choose one or the other.

For user-assigned, you'd have to create a managed ID first and at it later on.

Don't forget to add the managed identity to the Azure AD group, which was set as the Active Directory admin.

With that out of the way, we finally get to see some code! 🤓 To request an access token for a system managed identity you'd create a ManagedIdentityCredential first and then invoke the getToken() method, that takes a TokenRequestContext that sets the scope.

var credential = new ManagedIdentityCredentialBuilder().build();

var request = new TokenRequestContext()
		.addScopes("https://ossrdbms-aad.database.windows.net/.default");

var accessToken = credential
		.getToken(request)
		.retry(3L)
		.blockOptional()
		.orElseThrow(() -> new RuntimeException("Couldn't retrieve JWT"));

To request an access token for a user-assigned managed identity simply invoke clientId() and pass the appropriate GUID as shown in the Azure portal.

var credential = new ManagedIdentityCredentialBuilder()
		.clientId("f63574e7-c67d-4d6b-a4ea-78c55e4081c7")
		.build();

var request = new TokenRequestContext()
		.addScopes("https://ossrdbms-aad.database.windows.net/.default");

var accessToken = credential
		.getToken(request)
		.retry(3L)
		.blockOptional()
		.orElseThrow(() -> new RuntimeException("Failed to retrieve JWT"));

I might be stating the obvious here, but you won't be able to request an access token for a managed identity from your local developer machine. The endpoint required is only available from within Azure.

This particular REST endpoint is called IMDS (Azure Instance Metadata Service), that's available at a well-known, non-routable IP address 169.254.169.254. You can only access it from resources running in Azure. Like virtual machines, your app service instance, Azure Spring Cloud, and so on.

Using Azure CLI credentials

However, there are other ways to retrieve an access token for your own personal account, which is specifically important when developing locally. This can be done e.g. by using an AzureCliCredentialBuilder.

var credential = new AzureCliCredentialBuilder().build();

var request = new TokenRequestContext()
		.addScopes("https://ossrdbms-aad.database.windows.net/.default");

var accessToken = credential
		.getToken(request)
		.retry(3L)
		.blockOptional()
		.orElseThrow(() -> new RuntimeException("Failed to retrieve JWT"));

The AzureCliCredentialBuilder will pick up your existing Azure CLI session and request a token with that identity.

There are similar credential types, like VisualStudioCodeCredential and IntelliJCredential, whereas the latter is supposed to pick up the identity from the Azure Toolkit installed in IntelliJ IDEA. Unfortunately, the IntelliJCredential never worked for me.

Credential Chaining

Another interesting option is the ChainedTokenCredentialBuilder, which lets you chain several credential types together. All of the credential types are implementing the same interface TokenCredential. Chaining can be useful for fallback scenarios.

var azureCliCredential = new AzureCliCredentialBuilder()
		.build();

var intelliJCredential = new IntelliJCredentialBuilder()
		.keePassDatabasePath("C:\\Users\\MatthiasGuentert\\AppData\\Roaming\\JetBrains\\IntelliJIdea2021.3\\c.kdbx")
		.build();

var managedIdentityCredential = new ManagedIdentityCredentialBuilder()
		.clientId("f63574e7-c67d-4d6b-a4ea-78c55e4081c7")
		.build();

var credentialChain = new ChainedTokenCredentialBuilder()
		.addLast(azureCliCredential)
		.addLast(intelliJCredential)
		.addLast(managedIdentityCredential)
		.build();
		
var request = new TokenRequestContext()
		.addScopes("https://ossrdbms-aad.database.windows.net/.default");

var accessToken = credentialChain
		.getToken(request)
		.retry(3L)
		.blockOptional()
		.orElseThrow(() -> new RuntimeException("Failed to retrieve JWT"));		

JDBC Connection Pooling & HikariCP

The earlier 1.x versions of Spring Boot were using the Tomcat JDBC Connection Pooling library. Since Spring Boot version 2.x, HikariCP is the default, which provides improved performance and comes with the Spring-Boot-Starter-Data-JPA > Spring-Boot-Starter-JDBC > HikariCP dependency chain.

As mentioned in the introduction, the access token must be passed in the password field. This can easily be achieved by extending the HikariDataSource class, which itself implements the DataSource interface. From there we can override the getPassword() method and inject our logic.

Since an access token stays valid for some period (a couple of minutes), we don't want to ask for a new token each and every time a connection is required from the pool. Also, we don't want to think about refreshing it.

Instead, we are going to use a caching mechanism that Microsoft provides in the form of the SimpleTokenCache class. This class also makes sure our access token gets refreshed when required.

public class SimpleTokenCache {
    ...
    public SimpleTokenCache(Supplier<Mono<AccessToken>> tokenSupplier) {
       ...
    }
    ...
}

Putting everything together

When putting everything together you should end up with something like bellow.

import com.azure.core.credential.SimpleTokenCache;
import com.azure.core.credential.TokenCredential;
import com.azure.core.credential.TokenRequestContext;
import com.zaxxer.hikari.HikariDataSource;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.stereotype.Component;

@Component
@ConfigurationProperties(prefix = "spring.datasource.hikari")
public class AzureAdDataSource extends HikariDataSource {

    private final SimpleTokenCache cache;

    public AzureAdDataSource(TokenCredential credential) {
        this.cache = new SimpleTokenCache(() -> credential.getToken(createRequestContext()));
    }

    @Override
    public String getPassword() {
        var accessToken = cache
                .getToken()
                .retry(1L)
                .blockOptional()
                .orElseThrow(() -> new RuntimeException("Attempt to retrieve AAD token failed"));

        return accessToken.getToken();
    }

    private static TokenRequestContext createRequestContext() {
        return new TokenRequestContext().addScopes("https://ossrdbms-aad.database.windows.net/.default");
    }
}

The @ConfigurationProperties annotation is required so that we can configure our AzureAdDataSource as we are used to.

spring:
    liquibase:
        enabled: true
    datasource:
        driver-class-name: org.postgresql.Driver
        hikari:
            jdbc-url: jdbc:postgresql://psql-demo.postgres.database.azure.com:5432/demo
            username: psql-administrators@psql-demo

This is how I configured the beans. As mentioned before, ChainedTokenCredential implements the TokenCredential interface, that's getting injected into the AzureAdDataSource class.

import com.azure.identity.*;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class AppConfig {
    @Bean
    public ChainedTokenCredential chainedTokenCredential() {
        var azureCliCredential = new AzureCliCredentialBuilder()
                .build();

        var managedIdentityCredential = new ManagedIdentityCredentialBuilder()
                .clientId("f63574e7-c67d-4d6b-a4ea-78c55e4081c7")
                .build();

        var credentialChain = new ChainedTokenCredentialBuilder()
                .addLast(azureCliCredential)
                .addLast(intelliJCredential)
                .addLast(managedIdentityCredential)
                .build();

        return credentialChain;
    }
}

A more pragmatic configuration approach that increases the application's startup performance is using the Profile annotation. Also, it won't clutter your log with errors when running in production.

@Configuration
public class AppConfig {

    @Bean
    @Profile("local")
    public AzureCliCredential azureCliCredential() {
        return new AzureCliCredentialBuilder().build();
    }

    @Bean
    @Profile("!local")
    public ManagedIdentityCredential managedIdentityCredential() {
        return new ManagedIdentityCredentialBuilder()
                .clientId("f63574e7-c67d-4d6b-a4ea-78c55e4081c7")
                .build();
    }
}

Conclusion

In this post, I have demonstrated how the Azure Identity library can be used together with a custom Spring Boot DataSource to authenticate Azure AD identities against a Postgres Single Server instance.

I hope this article was informative to my readers and as always I am looking forward to feedback. Happy coding, Matthias 😎👨🏻‍💻

Further reading