如何使用 Spring Data JPA 访问 Azure AD 集成的 Postgres 数据库?

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

我们在 Azure 中有一个 PostgreSQL 单一服务实例 运行,它集成了 Azure AD。所以要通过 psql 连接,我按照以下步骤操作:

  1. 通过az login
  2. 登录
  3. 检索访问令牌$env:PGPASSWORD=$(az account get-access-token --resource-type oss-rdbms --query accessToken --output tsv)
  4. 登录psql "host=single-server-instance.postgres.database.azure.com user=aad_group@single-server-instance dbname=demodb"

到目前为止一切顺利。但是我如何使用 Spring Data JPA 来做到这一点?

这是我当前的 application.properties 文件的样子。我当然不想一遍又一遍地插入访问令牌。

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

这是我来自 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'
}
// ...

问题:

  1. 我是否应该将 spring.datasourcespring.jpa 配置部分移动到我的代码中,并从那里向 spring.datasource.password 提供访问令牌?如果是这样,我该怎么做?它应该去哪里?
  2. 如何检索访问令牌并将其传递给 Spring 数据?
  3. 我应该如何处理访问令牌的缓存?我是否必须处理刷新令牌并处理访问令牌过期?

我们扩展了 HikairDataSource 并覆盖了它的 getPassword() 方法,该方法在创建新连接之前调用。

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");
    }
}

使用的库:'com.azure:azure-identity:1.3.3', 'com.zaxxer:HikariCP:4.0.3'

调整属性,或在配置文件中创建 AzureAdDataSource 实例,而不是将其注释为组件:

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