Adding Old Constructor back to AKV Provider

src/main/java/com/microsoft/sqlserver/jdbc/KeyVaultCredential.java

@@ -24,20 +24,30 @@
  */
 class KeyVaultCredential extends KeyVaultCredentials {
 
+	SQLServerKeyVaultAuthenticationCallback authenticationCallback = null;
     String clientId = null;
     String clientKey = null;
+    String accessToken = null;
 
     KeyVaultCredential(String clientId,
             String clientKey) {
         this.clientId = clientId;
         this.clientKey = clientKey;
     }
-
+
+    KeyVaultCredential(SQLServerKeyVaultAuthenticationCallback authenticationCallback) {
+        this.authenticationCallback = authenticationCallback;
+    }
+
     public String doAuthenticate(String authorization,
             String resource,
             String scope) {
-        AuthenticationResult token = getAccessTokenFromClientCredentials(authorization, resource, clientId, clientKey);
-        return token.getAccessToken();
+    	if(authenticationCallback==null) {
+    		AuthenticationResult token = getAccessTokenFromClientCredentials(authorization, resource, clientId, clientKey);
+    		return token.getAccessToken();
+    	}else {
+    		return authenticationCallback.getAccessToken(authorization, resource, scope);
+    	}
     }
 
     private static AuthenticationResult getAccessTokenFromClientCredentials(String authorization,

src/main/java/com/microsoft/sqlserver/jdbc/SQLServerColumnEncryptionAzureKeyVaultProvider.java

@@ -18,13 +18,20 @@
 import java.security.NoSuchAlgorithmException;
 import java.text.MessageFormat;
 import java.util.Locale;
+import java.util.concurrent.ExecutorService;
 
+import com.microsoft.azure.AzureResponseBuilder;
 import com.microsoft.azure.keyvault.KeyVaultClient;
 import com.microsoft.azure.keyvault.models.KeyBundle;
 import com.microsoft.azure.keyvault.models.KeyOperationResult;
 import com.microsoft.azure.keyvault.models.KeyVerifyResult;
 import com.microsoft.azure.keyvault.webkey.JsonWebKeyEncryptionAlgorithm;
 import com.microsoft.azure.keyvault.webkey.JsonWebKeySignatureAlgorithm;
+import com.microsoft.azure.serializer.AzureJacksonAdapter;
+import com.microsoft.rest.RestClient;
+
+import okhttp3.OkHttpClient;
+import retrofit2.Retrofit;
 
 /**
  * Provides implementation similar to certificate store provider. A CEK encrypted with certificate store provider should be decryptable by this
@@ -41,7 +48,9 @@ public class SQLServerColumnEncryptionAzureKeyVaultProvider extends SQLServerCol
      * Column Encryption Key Store Provider string
      */
     String name = "AZURE_KEY_VAULT";
-
+
+    private final String baseUrl = "https://{vaultBaseUrl}";
+
     private final String azureKeyVaultDomainName = "vault.azure.net";
 
     private final String rsaEncryptionAlgorithmWithOAEPForAKV = "RSA-OAEP";
@@ -53,7 +62,7 @@ public class SQLServerColumnEncryptionAzureKeyVaultProvider extends SQLServerCol
 
     private KeyVaultClient keyVaultClient;
 
-    private KeyVaultCredential credential;
+    private KeyVaultCredential credentials;
 
     public void setName(String name) {
         this.name = name;
@@ -62,7 +71,35 @@ public void setName(String name) {
     public String getName() {
         return this.name;
     }
-
+
+    /**
+     * Constructor that takes a callback function to authenticate to AAD. This is used by KeyVaultClient at runtime to authenticate to Azure Key
+     * Vault.
+     * 
+     * @param authenticationCallback
+     *            - Callback function used for authenticating to AAD.
+     * @param executorService
+     *            - The ExecutorService used to create the keyVaultClient
+     * @throws SQLServerException
+     *             when an error occurs
+     */
+    public SQLServerColumnEncryptionAzureKeyVaultProvider(SQLServerKeyVaultAuthenticationCallback authenticationCallback,
+            ExecutorService executorService) throws SQLServerException {
+        if (null == authenticationCallback) {
+            MessageFormat form = new MessageFormat(SQLServerException.getErrString("R_NullValue"));
+            Object[] msgArgs1 = {"SQLServerKeyVaultAuthenticationCallback"};
+            throw new SQLServerException(form.format(msgArgs1), null);
+        }
+        credentials = new KeyVaultCredential(authenticationCallback);        
+        RestClient restClient = new RestClient.Builder(new OkHttpClient.Builder(), new Retrofit.Builder())
+                .withBaseUrl(baseUrl)
+                .withCredentials(credentials)
+                .withSerializerAdapter(new AzureJacksonAdapter())
+                .withResponseBuilderFactory(new AzureResponseBuilder.Factory())
+                .build();
+        keyVaultClient = new KeyVaultClient(restClient);
+    }
+
     /**
      * Constructor that authenticates to AAD. This is used by KeyVaultClient at runtime to authenticate to Azure Key
      * Vault.
@@ -76,8 +113,8 @@ public String getName() {
      */
     public SQLServerColumnEncryptionAzureKeyVaultProvider(String clientId,
             String clientKey) throws SQLServerException {
-        credential = new KeyVaultCredential(clientId, clientKey);
-        keyVaultClient = new KeyVaultClient(credential);
+        credentials = new KeyVaultCredential(clientId, clientKey);
+        keyVaultClient = new KeyVaultClient(credentials);
     }
 
     /**

src/main/java/com/microsoft/sqlserver/jdbc/SQLServerKeyVaultAuthenticationCallback.java

@@ -0,0 +1,27 @@
+/*
+ * Microsoft JDBC Driver for SQL Server
+ * 
+ * Copyright(c) Microsoft Corporation All rights reserved.
+ * 
+ * This program is made available under the terms of the MIT License. See the LICENSE file in the project root for more information.
+ */
+
+package com.microsoft.sqlserver.jdbc;
+
+public interface SQLServerKeyVaultAuthenticationCallback {
+
+    /**
+     * The authentication callback delegate which is to be implemented by the client code
+     * 
+     * @param authority
+     *            - Identifier of the authority, a URL.
+     * @param resource
+     *            - Identifier of the target resource that is the recipient of the requested token, a URL.
+     * @param scope
+     *            - The scope of the authentication request.
+     * @return access token
+     */
+    public String getAccessToken(String authority,
+            String resource,
+            String scope);
+}