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 index 0f25a7e7..f428982b 100644 --- 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> tokenProvider; + private final Mono tokenProvider; - AsyncCallbackTokenProvider(@NotNull String clusterUrl, @NotNull Function> tokenProvider) throws URISyntaxException { + AsyncCallbackTokenProvider(@NotNull String clusterUrl, @NotNull Mono tokenProvider) throws URISyntaxException { super(clusterUrl, null); this.tokenProvider = tokenProvider; } @Override protected Mono 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 index 65101966..f0809ca6 100644 --- 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 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 index 6d6e9bec..f6a537c3 100644 --- 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 tokenProvider; + private Mono 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 getTokenProvider() { return tokenProvider; } + public Mono 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 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); }