update

microsoft · Mar 4, 2025 · ed2e189 · ed2e189
1 parent 15bf760
commit ed2e189
Showing 1 changed file with 16 additions and 15 deletions.
diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/FedauthCommon.java b/src/test/java/com/microsoft/sqlserver/jdbc/fedauth/FedauthCommon.java
@@ -7,6 +7,10 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
+import com.azure.core.credential.AccessToken;
+import com.azure.core.credential.TokenRequestContext;
+import com.azure.identity.ManagedIdentityCredential;
+import com.azure.identity.ManagedIdentityCredentialBuilder;
 import com.microsoft.aad.msal4j.ClientCredentialFactory;
 import com.microsoft.aad.msal4j.ClientCredentialParameters;
 import com.microsoft.aad.msal4j.ConfidentialClientApplication;
@@ -21,6 +25,7 @@
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.Collections;
 import java.util.Date;
 import java.util.HashSet;
 import java.util.Locale;
@@ -216,25 +221,21 @@ public static void getConfigs() throws Exception {
     static void getFedauthInfo() {
         int retry = 0;
         long interval = THROTTLE_RETRY_INTERVAL;
+        ManagedIdentityCredential credential = new ManagedIdentityCredentialBuilder()
+                .clientId(akvProviderManagedClientId).build();
+
         while (retry <= THROTTLE_RETRY_COUNT) {
             try {
-                Set<String> scopes = new HashSet<>();
-                scopes.add(spn + "/.default");
-                if (null == fedauthClientApp) {
-                    IClientCredential credential = ClientCredentialFactory.createFromSecret(applicationKey);
-                    fedauthClientApp = ConfidentialClientApplication.builder(applicationClientID, credential)
-                            .executorService(Executors.newFixedThreadPool(1))
-                            .setTokenCacheAccessAspect(FedauthTokenCache.getInstance()).authority(stsurl).build();
-                }
+                TokenRequestContext requestContext = new TokenRequestContext()
+                        .setScopes(Collections.singletonList(spn + "/.default"));
 
-                final CompletableFuture<IAuthenticationResult> future = fedauthClientApp
-                        .acquireToken(ClientCredentialParameters.builder(scopes).build());
+                AccessToken token = credential.getToken(requestContext).block();
 
-                final IAuthenticationResult authenticationResult = future.get();
-
-                secondsBeforeExpiration = TimeUnit.MILLISECONDS
-                        .toSeconds(authenticationResult.expiresOnDate().getTime() - new Date().getTime());
-                accessToken = authenticationResult.accessToken();
+                if (token != null) {
+                    secondsBeforeExpiration = TimeUnit.MILLISECONDS
+                            .toSeconds(token.getExpiresAt().toInstant().toEpochMilli() - new Date().getTime());
+                    accessToken = token.getToken();
+                }
 
                 retry = THROTTLE_RETRY_COUNT + 1;
             } catch (MsalThrottlingException te) {