javaspringspring-data-jpapostgresql-11azure-postgresql

How to access Azure AD integrated Postgres DB using Spring Data JPA?


We are having a PostgreSQL single service instance running in Azure, which is Azure AD integrated. So to connect via psql I follow this steps:

  1. Log in via az login
  2. Retrieve access token $env:PGPASSWORD=$(az account get-access-token --resource-type oss-rdbms --query accessToken --output tsv)
  3. Login psql "host=single-server-instance.postgres.database.azure.com user=aad_group@single-server-instance dbname=demodb"

So far so good. But how would I do that with Spring Data JPA?

This is what my current application.properties file looks like. Of course I don't want to insert the access token over and over again.

logging.level.org.hibernate.SQL=DEBUG

spring.datasource.url=jdbc:postgresql://single-server-instance.postgres.database.azure.com:5432/demodb
spring.datasource.username=aad_group@single-server-instance
spring.datasource.password=eyJ...there goes the access token

spring.jpa.database-platform=org.hibernate.dialect.PostgreSQLDialect    
spring.jpa.show-sql=true
spring.jpa.hibernate.ddl-auto=create-drop

This are my dependencies from build.gradle.

// ...
dependencies {
    implementation 'org.springframework.boot:spring-boot-starter-data-jpa'
    implementation 'org.springframework.boot:spring-boot-starter-web'
    runtimeOnly 'org.postgresql:postgresql'
    testImplementation 'org.springframework.boot:spring-boot-starter-test'
}
// ...

Questions:

  1. Should I move the spring.datasource and spring.jpa configuration part into my code and supply the access token to the spring.datasource.password from there? If so, how can I do this? Where should it go?
  2. How can I retrieve an access token and pass it to Spring Data?
  3. How should I deal with the caching of the access token? Do I have to take care of the refresh token and deal with access token expiry?

Solution

  • We've extended a HikairDataSource and overrided its getPassword() method, which is called before creation of new connection.

    import com.azure.core.credential.AccessToken;
    import com.azure.core.credential.SimpleTokenCache;
    import com.azure.core.credential.TokenCredential;
    import com.azure.core.credential.TokenRequestContext;
    import com.zaxxer.hikari.HikariDataSource;
    
    @Component
    @ConfigurationProperties(prefix = "spring.datasource.hikari")
    public class AzureAdDataSource extends HikariDataSource {
    
        private final SimpleTokenCache cache;
    
        public AzureAdDataSource(TokenCredential aadTokenCredential) {
            this.cache = new SimpleTokenCache(() -> aadTokenCredential.getToken(createRequestContext()));
        }
    
        @Override
        public String getPassword() {
            final AccessToken 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");
        }
    }
    

    used libs: 'com.azure:azure-identity:1.3.3', 'com.zaxxer:HikariCP:4.0.3'

    adjust properties, or create the instance of AzureAdDataSource in configuration file instead annotating it as Component:

    spring.datasource.hikari.jdbc-url=jdbc:postgresql://single-server-instance.postgres.database.azure.com:5432/demodb
    spring.datasource.hikari.username=aad_group@single-server-instance