TODO

Azure · Nov 12, 2024 · a4c59bf · a4c59bf
1 parent 3650f99
commit a4c59bf
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 8 deletions.
diff --git a/data/src/main/java/com/microsoft/azure/kusto/data/auth/AsyncCallbackTokenProvider.java b/data/src/main/java/com/microsoft/azure/kusto/data/auth/AsyncCallbackTokenProvider.java
@@ -1,27 +1,32 @@
 package com.microsoft.azure.kusto.data.auth;
 
-import com.azure.core.http.HttpClient;
 import com.microsoft.azure.kusto.data.exceptions.DataClientException;
+import com.microsoft.azure.kusto.data.exceptions.ExceptionsUtils;
 import org.jetbrains.annotations.NotNull;
 import reactor.core.publisher.Mono;
 
 import java.net.URISyntaxException;
-import java.util.function.Function;
 
-// TODO - Add to KCSB
 public class AsyncCallbackTokenProvider extends TokenProviderBase {
     public static final String CALLBACK_TOKEN_PROVIDER = "CallbackTokenProvider";
-    private final Function<HttpClient, Mono<String>> tokenProvider;
+    private final Mono<String> tokenProvider;
 
-    AsyncCallbackTokenProvider(@NotNull String clusterUrl, @NotNull Function<HttpClient, Mono<String>> tokenProvider) throws URISyntaxException {
+    AsyncCallbackTokenProvider(@NotNull String clusterUrl, @NotNull Mono<String> tokenProvider) throws URISyntaxException {
         super(clusterUrl, null);
         this.tokenProvider = tokenProvider;
     }
 
     @Override
     protected Mono<String> acquireAccessTokenImpl() {
-        return tokenProvider.apply(httpClient)
-                .onErrorMap(e -> new DataClientException(clusterUrl, e.getMessage(), e instanceof Exception ? (Exception) e : null));
+        return tokenProvider
+                .onErrorMap(e -> {
+                    if (e instanceof Exception) {
+                        Exception ex = (Exception) e;
+                        return new DataClientException(clusterUrl, ExceptionsUtils.getMessageEx(ex), ex);
+                    } else {
+                        return new DataClientException(clusterUrl, e.toString(), null);
+                    }
+                });
     }
 
     @Override

diff --git a/data/src/main/java/com/microsoft/azure/kusto/data/auth/CallbackTokenProvider.java b/data/src/main/java/com/microsoft/azure/kusto/data/auth/CallbackTokenProvider.java
@@ -33,7 +33,14 @@ public class CallbackTokenProvider extends TokenProviderBase {
     protected Mono<String> acquireAccessTokenImpl() {
         return Mono.fromCallable(() -> tokenProvider.apply(httpClient))
                 // TODO - is this a better way?
-                .onErrorMap(e -> new DataClientException(clusterUrl, ExceptionsUtils.getMessageEx((Exception) e), (Exception) e));
+                .onErrorMap(e -> {
+                    if (e instanceof Exception) {
+                        Exception ex = (Exception) e;
+                        return new DataClientException(clusterUrl, ExceptionsUtils.getMessageEx(ex), ex);
+                    } else {
+                        return new DataClientException(clusterUrl, e.toString(), null);
+                    }
+                });
     }
 
     @Override

diff --git a/data/src/main/java/com/microsoft/azure/kusto/data/auth/ConnectionStringBuilder.java b/data/src/main/java/com/microsoft/azure/kusto/data/auth/ConnectionStringBuilder.java
@@ -6,6 +6,7 @@
 import com.microsoft.azure.kusto.data.ClientDetails;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.Pair;
+import reactor.core.publisher.Mono;
 import reactor.util.annotation.Nullable;
 
 import java.security.PrivateKey;
@@ -30,6 +31,7 @@ public class ConnectionStringBuilder {
     private String aadAuthorityId;
     private String accessToken;
     private Callable<String> tokenProvider;
+    private Mono<String> asyncTokenProvider;
     private String managedIdentityClientId;
     private boolean useDeviceCodeAuth;
     private boolean useManagedIdentityAuth;
@@ -85,6 +87,7 @@ private ConnectionStringBuilder() {
         this.aadAuthorityId = null;
         this.accessToken = null;
         this.tokenProvider = null;
+        this.asyncTokenProvider = null;
         this.managedIdentityClientId = null;
         this.useDeviceCodeAuth = false;
         this.useManagedIdentityAuth = false;
@@ -147,6 +150,7 @@ public ConnectionStringBuilder(ConnectionStringBuilder other) {
         this.aadAuthorityId = other.aadAuthorityId;
         this.accessToken = other.accessToken;
         this.tokenProvider = other.tokenProvider;
+        this.asyncTokenProvider = other.asyncTokenProvider;
         this.managedIdentityClientId = other.managedIdentityClientId;
         this.useAzureCli = other.useAzureCli;
         this.useDeviceCodeAuth = other.useDeviceCodeAuth;
@@ -230,6 +234,10 @@ public Callable<String> getTokenProvider() {
         return tokenProvider;
     }
 
+    public Mono<String> getAsyncTokenProvider() {
+        return asyncTokenProvider;
+    }
+
     public String getManagedIdentityClientId() {
         return managedIdentityClientId;
     }
@@ -470,6 +478,21 @@ public static ConnectionStringBuilder createWithAadTokenProviderAuthentication(S
         return csb;
     }
 
+    public static ConnectionStringBuilder createWithAadAsyncTokenProviderAuthentication(String clusterUrl, Mono<String> tokenProviderCallable) {
+        if (StringUtils.isEmpty(clusterUrl)) {
+            throw new IllegalArgumentException("clusterUrl cannot be null or empty");
+        }
+
+        if (tokenProviderCallable == null) {
+            throw new IllegalArgumentException("tokenProviderCallback cannot be null");
+        }
+
+        ConnectionStringBuilder csb = new ConnectionStringBuilder();
+        csb.clusterUrl = clusterUrl;
+        csb.asyncTokenProvider = tokenProviderCallable;
+        return csb;
+    }
+
     public static ConnectionStringBuilder createWithAadManagedIdentity(String clusterUrl) {
         return createWithAadManagedIdentity(clusterUrl, null);
     }