Retrieve TokenCredential from ConfigurableBootstrapContext for env post processor

sdk/spring/CHANGELOG.md

@@ -1,7 +1,16 @@
 # Release History
-# (Unreleased)
+
+## 5.16.0-beta.1 (unreleased)
+
+#### Dependency Updates
 Upgrade Spring Boot dependencies version to 3.3.3 and Spring Cloud dependencies version to 2023.0.3
 
+### Spring Cloud Azure Autoconfigure
+This section includes changes in `spring-cloud-azure-autoconfigure` module.
+
+#### Features Added
+- Provide extension point to configure token credential for Key Vault property source  [#41580](https://github.com/Azure/azure-sdk-for-java/pull/41580).
+
 ## 5.15.0 (2024-08-07)
 - This release is compatible with Spring Boot 3.0.0-3.0.13, 3.1.0-3.1.12, 3.2.0-3.2.7, 3.3.0-3.3.2. (Note: 3.0.x (x>13), 3.1.y (y>12), 3.2.z (z>7) and 3.3.m (m>2) should be supported, but they aren't tested with this release.)
 - This release is compatible with Spring Cloud 2022.0.0-2022.0.5, 2023.0.0-2023.0.3. (Note: 2022.0.x (x>5) and 2023.0.y (y>3) should be supported, but they aren't tested with this release.)

sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/implementation/context/AzureTokenCredentialAutoConfiguration.java

@@ -22,6 +22,8 @@
 import com.azure.spring.cloud.core.implementation.factory.credential.ManagedIdentityCredentialBuilderFactory;
 import com.azure.spring.cloud.core.implementation.factory.credential.UsernamePasswordCredentialBuilderFactory;
 import com.azure.spring.cloud.core.provider.authentication.TokenCredentialOptionsProvider;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.BeanFactoryAware;
@@ -51,6 +53,7 @@
 @Configuration(proxyBeanMethods = false)
 @AutoConfigureAfter(TaskExecutionAutoConfiguration.class)
 public class AzureTokenCredentialAutoConfiguration extends AzureServiceConfigurationBase {
+    private static final Logger LOGGER = LoggerFactory.getLogger(AzureTokenCredentialAutoConfiguration.class);
 
     private final IdentityClientProperties identityClientProperties;
 
@@ -68,6 +71,7 @@ TokenCredential tokenCredential(DefaultAzureCredentialBuilderFactory factory,
         if (globalTokenCredential != null) {
             return globalTokenCredential;
         } else {
+            LOGGER.debug("No global token credential found, constructing default credential.");
             return factory.build().build();
         }
     }

sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/implementation/keyvault/environment/KeyVaultEnvironmentPostProcessor.java

@@ -3,20 +3,23 @@
 
 package com.azure.spring.cloud.autoconfigure.implementation.keyvault.environment;
 
+import com.azure.core.credential.TokenCredential;
 import com.azure.security.keyvault.secrets.SecretClient;
 import com.azure.spring.cloud.autoconfigure.implementation.context.properties.AzureGlobalProperties;
 import com.azure.spring.cloud.autoconfigure.implementation.keyvault.secrets.properties.AzureKeyVaultPropertySourceProperties;
 import com.azure.spring.cloud.autoconfigure.implementation.keyvault.secrets.properties.AzureKeyVaultSecretProperties;
+import com.azure.spring.cloud.core.implementation.credential.resolver.AzureTokenCredentialResolver;
 import com.azure.spring.cloud.core.implementation.util.AzurePropertiesUtils;
 import com.azure.spring.cloud.core.implementation.util.AzureSpringIdentifier;
 import com.azure.spring.cloud.service.implementation.keyvault.secrets.SecretClientBuilderFactory;
 import org.apache.commons.logging.Log;
+import org.springframework.boot.ConfigurableBootstrapContext;
 import org.springframework.boot.SpringApplication;
 import org.springframework.boot.context.config.ConfigDataEnvironmentPostProcessor;
 import org.springframework.boot.context.properties.bind.Bindable;
 import org.springframework.boot.context.properties.bind.Binder;
 import org.springframework.boot.env.EnvironmentPostProcessor;
-import org.springframework.boot.logging.DeferredLog;
+import org.springframework.boot.logging.DeferredLogFactory;
 import org.springframework.core.Ordered;
 import org.springframework.core.env.ConfigurableEnvironment;
 import org.springframework.core.env.MutablePropertySources;
@@ -44,21 +47,17 @@ public class KeyVaultEnvironmentPostProcessor implements EnvironmentPostProcesso
     private static final String SKIP_CONFIGURE_REASON_FORMAT = "Skip configuring Key Vault PropertySource because %s.";
 
     private final Log logger;
+    private final ConfigurableBootstrapContext bootstrapContext;
 
 
     /**
      * Creates a new instance of {@link KeyVaultEnvironmentPostProcessor}.
-     * @param logger The logger used in this class.
+     * @param loggerFactory The logger factory to get the logger.
+     * @param bootstrapContext The bootstrap context.
      */
-    public KeyVaultEnvironmentPostProcessor(Log logger) {
-        this.logger = logger;
-    }
-
-    /**
-     * Construct a {@link KeyVaultEnvironmentPostProcessor} instance with a new {@link DeferredLog}.
-     */
-    public KeyVaultEnvironmentPostProcessor() {
-        this.logger = new DeferredLog();
+    public KeyVaultEnvironmentPostProcessor(DeferredLogFactory loggerFactory, ConfigurableBootstrapContext bootstrapContext) {
+        this.logger = loggerFactory.getLog(getClass());
+        this.bootstrapContext = bootstrapContext;
     }
 
     /**
@@ -155,6 +154,17 @@ private AzureKeyVaultSecretProperties toAzureKeyVaultSecretProperties(
     SecretClient buildSecretClient(AzureKeyVaultSecretProperties secretProperties) {
         SecretClientBuilderFactory factory = new SecretClientBuilderFactory(secretProperties);
         factory.setSpringIdentifier(AzureSpringIdentifier.AZURE_SPRING_KEY_VAULT_SECRETS);
+
+        if (bootstrapContext != null && bootstrapContext.isRegistered(TokenCredential.class)) {
+            // If TokenCredential is registered in bootstrap context, use it to build SecretClient.
+            // This will ignore the credential properties configured
+            TokenCredential registerCredential = bootstrapContext.get(TokenCredential.class);
+            logger.debug(registerCredential.getClass().getSimpleName() + " is registered in bootstrap context, use it to build SecretClient.");
+            factory.setTokenCredentialResolver(
+                new AzureTokenCredentialResolver(ignored -> registerCredential)
+            );
+        }
+
         return factory.build().buildClient();
     }
 

sdk/spring/spring-cloud-azure-autoconfigure/src/test/java/com/azure/spring/cloud/autoconfigure/implementation/keyvault/environment/KeyVaultEnvironmentPostProcessorTests.java

@@ -3,15 +3,17 @@
 
 package com.azure.spring.cloud.autoconfigure.implementation.keyvault.environment;
 
+import com.azure.core.credential.TokenCredential;
 import com.azure.security.keyvault.secrets.SecretClient;
 import com.azure.spring.cloud.autoconfigure.implementation.keyvault.secrets.properties.AzureKeyVaultPropertySourceProperties;
 import com.azure.spring.cloud.autoconfigure.implementation.keyvault.secrets.properties.AzureKeyVaultSecretProperties;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.mockito.MockedStatic;
+import org.springframework.boot.ConfigurableBootstrapContext;
 import org.springframework.boot.SpringApplication;
-import org.springframework.boot.logging.DeferredLog;
+import org.springframework.boot.logging.DeferredLogs;
 import org.springframework.core.env.MutablePropertySources;
 import org.springframework.core.env.PropertySource;
 import org.springframework.core.env.SystemEnvironmentPropertySource;
@@ -29,7 +31,11 @@
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 import static org.springframework.core.env.StandardEnvironment.SYSTEM_ENVIRONMENT_PROPERTY_SOURCE_NAME;
 
 class KeyVaultEnvironmentPostProcessorTests {
@@ -43,19 +49,50 @@ class KeyVaultEnvironmentPostProcessorTests {
     private KeyVaultEnvironmentPostProcessor processor;
     private MockEnvironment environment;
     private MutablePropertySources propertySources;
+    private ConfigurableBootstrapContext context;
 
     @BeforeEach
     void beforeEach() {
-        processor = spy(new KeyVaultEnvironmentPostProcessor(new DeferredLog()));
+        processor = spy(new KeyVaultEnvironmentPostProcessor(new DeferredLogs(), null));
         environment = new MockEnvironment();
         propertySources = environment.getPropertySources();
         SecretClient secretClient = mock(SecretClient.class);
         doReturn(secretClient).when(processor).buildSecretClient(any(AzureKeyVaultSecretProperties.class));
     }
 
+    @Test
+    void testContextRegisterWithTokenCredentialRegistered() {
+        context = mock(ConfigurableBootstrapContext.class);
+        TokenCredential tokenCredential = mock(TokenCredential.class);
+        when(context.get(TokenCredential.class)).thenReturn(tokenCredential);
+        when(context.isRegistered(TokenCredential.class)).thenReturn(true);
+        processor = spy(new KeyVaultEnvironmentPostProcessor(new DeferredLogs(), context));
+        AzureKeyVaultSecretProperties secretProperties = new AzureKeyVaultSecretProperties();
+        secretProperties.setEndpoint(ENDPOINT_0);
+
+        processor.buildSecretClient(secretProperties);
+
+        verify(context, times(1)).get(TokenCredential.class);
+    }
+
+    @Test
+    void testContextRegisterWithoutTokenCredentialRegistered() {
+        context = mock(ConfigurableBootstrapContext.class);
+        TokenCredential tokenCredential = mock(TokenCredential.class);
+        when(context.get(TokenCredential.class)).thenReturn(tokenCredential);
+        when(context.isRegistered(TokenCredential.class)).thenReturn(false);
+        processor = spy(new KeyVaultEnvironmentPostProcessor(new DeferredLogs(), context));
+        AzureKeyVaultSecretProperties secretProperties = new AzureKeyVaultSecretProperties();
+        secretProperties.setEndpoint(ENDPOINT_0);
+
+        processor.buildSecretClient(secretProperties);
+
+        verify(context, never()).get(TokenCredential.class);
+    }
+
     @Test
     void postProcessorHasConfiguredOrder() {
-        final KeyVaultEnvironmentPostProcessor processor = new KeyVaultEnvironmentPostProcessor();
+        final KeyVaultEnvironmentPostProcessor processor = new KeyVaultEnvironmentPostProcessor(new DeferredLogs(), null);
         assertEquals(processor.getOrder(), KeyVaultEnvironmentPostProcessor.ORDER);
     }
 
@@ -308,7 +345,7 @@ void buildKeyVaultPropertySourceWithExceptionTest() {
         environment.setProperty("spring.cloud.azure.keyvault.secret.property-sources[0].name", NAME_0);
         environment.setProperty("spring.cloud.azure.keyvault.secret.property-sources[0].endpoint", ENDPOINT_0);
         assertThrows(IllegalStateException.class,
-                () -> new KeyVaultEnvironmentPostProcessor().postProcessEnvironment(environment, application));
+                () -> new KeyVaultEnvironmentPostProcessor(new DeferredLogs(), null).postProcessEnvironment(environment, application));
     }
 }
 

sdk/spring/spring-cloud-azure-autoconfigure/src/test/java/com/azure/spring/cloud/autoconfigure/implementation/keyvault/environment/KeyVaultSecretPropertySourceUserAgentTests.java

@@ -11,6 +11,7 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.api.parallel.Isolated;
+import org.springframework.boot.logging.DeferredLogs;
 import org.springframework.boot.test.system.CapturedOutput;
 import org.springframework.boot.test.system.OutputCaptureExtension;
 
@@ -32,7 +33,7 @@ public void userAgentTest(CapturedOutput output) {
         properties.getRetry().getFixed().setDelay(Duration.ofSeconds(1));
         properties.getRetry().getFixed().setMaxRetries(0);
 
-        KeyVaultEnvironmentPostProcessor environmentPostProcessor = new KeyVaultEnvironmentPostProcessor();
+        KeyVaultEnvironmentPostProcessor environmentPostProcessor = new KeyVaultEnvironmentPostProcessor(new DeferredLogs(), null);
         SecretClient secretClient = environmentPostProcessor.buildSecretClient(properties);
         try {
             secretClient.getSecret("property-source-name1");

sdk/spring/spring-cloud-azure-core/src/main/java/com/azure/spring/cloud/core/implementation/factory/AbstractAzureServiceClientBuilderFactory.java

@@ -200,6 +200,8 @@ protected void configureCredential(T builder) {
                                                  () -> new IllegalArgumentException("Consumer should not be null"));
 
 
+        LOGGER.debug("Will configure the credential of type {} for {}.", azureCredential.getClass().getSimpleName(),
+            builder.getClass().getSimpleName());
         consumer.accept(azureCredential);
         credentialConfigured = true;
     }
@@ -244,8 +246,8 @@ protected void configureConnectionString(T builder) {
      */
     protected void configureDefaultCredential(T builder) {
         if (!credentialConfigured) {
-            LOGGER.info("Will configure the default credential of type {} for {}.",
-                this.defaultTokenCredential.getClass().getSimpleName(), builder.getClass());
+            LOGGER.debug("Will configure the default credential of type {} for {}.",
+                this.defaultTokenCredential.getClass().getSimpleName(), builder.getClass().getSimpleName());
             consumeDefaultTokenCredential().accept(builder, this.defaultTokenCredential);
         }
     }