Change from CosmosKeyCredential to AzureKeyCredential

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java

@@ -5,6 +5,7 @@
 
 import com.azure.cosmos.implementation.Configs;
 import com.azure.cosmos.implementation.Constants;
+import com.azure.cosmos.implementation.CosmosError;
 import com.azure.cosmos.implementation.CosmosItemProperties;
 import com.azure.cosmos.implementation.DatabaseAccount;
 import com.azure.cosmos.implementation.Document;
@@ -27,7 +28,6 @@
 import com.azure.cosmos.implementation.query.metrics.ClientSideMetrics;
 import com.azure.cosmos.implementation.routing.PartitionKeyInternal;
 import com.azure.cosmos.models.CosmosAsyncItemResponse;
-import com.azure.cosmos.implementation.CosmosError;
 import com.azure.cosmos.models.CosmosItemResponse;
 import com.azure.cosmos.models.CosmosStoredProcedureProperties;
 import com.azure.cosmos.models.FeedResponse;
@@ -442,11 +442,6 @@ public static <T> CosmosItemProperties getProperties(CosmosItemResponse<T> cosmo
         return ModelBridgeInternal.getCosmosItemProperties(cosmosItemResponse);
     }
 
-    @Warning(value = INTERNAL_USE_ONLY_WARNING)
-    public static int getHashCode(CosmosKeyCredential keyCredential) {
-        return keyCredential.getKeyHashCode();
-    }
-
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
     public static String getLink(CosmosAsyncContainer cosmosAsyncContainer) {
         return cosmosAsyncContainer.getLink();

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java

@@ -3,6 +3,7 @@
 package com.azure.cosmos;
 
 import com.azure.core.annotation.ServiceClient;
+import com.azure.core.credential.AzureKeyCredential;
 import com.azure.cosmos.implementation.AsyncDocumentClient;
 import com.azure.cosmos.implementation.Configs;
 import com.azure.cosmos.implementation.ConnectionPolicy;
@@ -48,7 +49,7 @@ public final class CosmosAsyncClient implements Closeable {
     private final ConsistencyLevel desiredConsistencyLevel;
     private final List<CosmosPermissionProperties> permissions;
     private final CosmosAuthorizationTokenResolver cosmosAuthorizationTokenResolver;
-    private final CosmosKeyCredential cosmosKeyCredential;
+    private final AzureKeyCredential credential;
     private final boolean sessionCapturingOverride;
     private final boolean enableTransportClientSharing;
     private final boolean contentResponseOnWriteEnabled;
@@ -61,7 +62,7 @@ public final class CosmosAsyncClient implements Closeable {
         this.desiredConsistencyLevel = builder.getConsistencyLevel();
         this.permissions = builder.getPermissions();
         this.cosmosAuthorizationTokenResolver = builder.getAuthorizationTokenResolver();
-        this.cosmosKeyCredential = builder.getKeyCredential();
+        this.credential = builder.getCredential();
         this.sessionCapturingOverride = builder.isSessionCapturingOverrideEnabled();
         this.enableTransportClientSharing = builder.isConnectionSharingAcrossClientsEnabled();
         this.contentResponseOnWriteEnabled = builder.isContentResponseOnWriteEnabled();
@@ -73,7 +74,7 @@ public final class CosmosAsyncClient implements Closeable {
                                        .withSessionCapturingOverride(this.sessionCapturingOverride)
                                        .withConfigs(this.configs)
                                        .withTokenResolver(this.cosmosAuthorizationTokenResolver)
-                                       .withCosmosKeyCredential(this.cosmosKeyCredential)
+                                       .withCredential(this.credential)
                                        .withTransportClientSharing(this.enableTransportClientSharing)
                                        .withContentResponseOnWriteEnabled(this.contentResponseOnWriteEnabled)
                                        .build();
@@ -160,12 +161,12 @@ CosmosAuthorizationTokenResolver getCosmosAuthorizationTokenResolver() {
     }
 
     /**
-     * Gets the cosmos key credential
+     * Gets the azure key credential
      *
-     * @return cosmos key credential
+     * @return azure key credential
      */
-    CosmosKeyCredential cosmosKeyCredential() {
-        return cosmosKeyCredential;
+    AzureKeyCredential credential() {
+        return credential;
     }
 
     /**

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosBridgeInternal.java

@@ -80,7 +80,7 @@ public static CosmosClientBuilder cloneCosmosClientBuilder(CosmosClientBuilder b
             .directMode(builder.getDirectConnectionConfig())
             .gatewayMode(builder.getGatewayConnectionConfig())
             .consistencyLevel(builder.getConsistencyLevel())
-            .keyCredential(builder.getKeyCredential())
+            .credential(builder.getCredential())
             .permissions(builder.getPermissions())
             .authorizationTokenResolver(builder.getAuthorizationTokenResolver())
             .resourceToken(builder.getResourceToken())

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java

@@ -3,6 +3,7 @@
 package com.azure.cosmos;
 
 import com.azure.core.annotation.ServiceClientBuilder;
+import com.azure.core.credential.AzureKeyCredential;
 import com.azure.cosmos.implementation.Configs;
 import com.azure.cosmos.implementation.ConnectionPolicy;
 import com.azure.cosmos.implementation.CosmosAuthorizationTokenResolver;
@@ -38,7 +39,7 @@ public class CosmosClientBuilder {
     private ConsistencyLevel desiredConsistencyLevel;
     private List<CosmosPermissionProperties> permissions;
     private CosmosAuthorizationTokenResolver cosmosAuthorizationTokenResolver;
-    private CosmosKeyCredential cosmosKeyCredential;
+    private AzureKeyCredential credential;
     private boolean sessionCapturingOverrideEnabled;
     private boolean connectionSharingAcrossClientsEnabled;
     private boolean contentResponseOnWriteEnabled;
@@ -273,22 +274,22 @@ ConnectionPolicy getConnectionPolicy() {
     }
 
     /**
-     * Gets the {@link CosmosKeyCredential} to be used
+     * Gets the {@link AzureKeyCredential} to be used
      *
-     * @return cosmosKeyCredential
+     * @return {@link AzureKeyCredential}
      */
-    CosmosKeyCredential getKeyCredential() {
-        return cosmosKeyCredential;
+    AzureKeyCredential getCredential() {
+        return credential;
     }
 
     /**
-     * Sets the {@link CosmosKeyCredential} to be used
+     * Sets the {@link AzureKeyCredential} to be used
      *
-     * @param cosmosKeyCredential {@link CosmosKeyCredential}
+     * @param credential {@link AzureKeyCredential}
      * @return current cosmosClientBuilder
      */
-    public CosmosClientBuilder keyCredential(CosmosKeyCredential cosmosKeyCredential) {
-        this.cosmosKeyCredential = cosmosKeyCredential;
+    public CosmosClientBuilder credential(AzureKeyCredential credential) {
+        this.credential = credential;
         return this;
     }
 
@@ -634,10 +635,10 @@ private void validateConfig() {
         ifThrowIllegalArgException(this.serviceEndpoint == null,
             "cannot buildAsyncClient client without service endpoint");
         ifThrowIllegalArgException(
-            this.keyOrResourceToken == null && (permissions == null || permissions.isEmpty()) && this.cosmosKeyCredential == null,
+            this.keyOrResourceToken == null && (permissions == null || permissions.isEmpty()) && this.credential == null,
             "cannot buildAsyncClient client without any one of key, resource token, permissions, and "
-                + "cosmos key credential");
-        ifThrowIllegalArgException(cosmosKeyCredential != null && StringUtils.isEmpty(cosmosKeyCredential.getKey()),
+                + "azure key credential");
+        ifThrowIllegalArgException(credential != null && StringUtils.isEmpty(credential.getKey()),
             "cannot buildAsyncClient client without key credential");
         ifThrowIllegalArgException(directConnectionConfig == null && gatewayConnectionConfig == null,
             "cannot buildAsyncClient client without connection config");

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java

@@ -2,8 +2,8 @@
 // Licensed under the MIT License.
 package com.azure.cosmos.implementation;
 
+import com.azure.core.credential.AzureKeyCredential;
 import com.azure.cosmos.ConsistencyLevel;
-import com.azure.cosmos.CosmosKeyCredential;
 import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
 import com.azure.cosmos.models.FeedOptions;
 import com.azure.cosmos.models.FeedResponse;
@@ -71,7 +71,7 @@ class Builder {
         String masterKeyOrResourceToken;
         URI serviceEndpoint;
         CosmosAuthorizationTokenResolver cosmosAuthorizationTokenResolver;
-        CosmosKeyCredential cosmosKeyCredential;
+        AzureKeyCredential credential;
         boolean sessionCapturingOverride;
         boolean transportClientSharing;
         boolean contentResponseOnWriteEnabled;
@@ -148,11 +148,11 @@ public Builder withTransportClientSharing(boolean transportClientSharing) {
             return this;
         }
 
-        public Builder withCosmosKeyCredential(CosmosKeyCredential cosmosKeyCredential) {
-            if (cosmosKeyCredential != null && StringUtils.isEmpty(cosmosKeyCredential.getKey())) {
+        public Builder withCredential(AzureKeyCredential credential) {
+            if (credential != null && StringUtils.isEmpty(credential.getKey())) {
                 throw new IllegalArgumentException("Cannot buildAsyncClient client with empty key credential");
             }
-            this.cosmosKeyCredential = cosmosKeyCredential;
+            this.credential = credential;
             return this;
         }
 
@@ -183,10 +183,10 @@ public AsyncDocumentClient build() {
             ifThrowIllegalArgException(this.serviceEndpoint == null, "cannot buildAsyncClient client without service endpoint");
             ifThrowIllegalArgException(
                     this.masterKeyOrResourceToken == null && (permissionFeed == null || permissionFeed.isEmpty())
-                        && this.cosmosKeyCredential == null,
+                        && this.credential == null,
                     "cannot buildAsyncClient client without any one of masterKey, " +
-                        "resource token, permissionFeed and cosmos key credential");
-            ifThrowIllegalArgException(cosmosKeyCredential != null && StringUtils.isEmpty(cosmosKeyCredential.getKey()),
+                        "resource token, permissionFeed and azure key credential");
+            ifThrowIllegalArgException(credential != null && StringUtils.isEmpty(credential.getKey()),
                 "cannot buildAsyncClient client without key credential");
 
             RxDocumentClientImpl client = new RxDocumentClientImpl(serviceEndpoint,
@@ -196,7 +196,7 @@ public AsyncDocumentClient build() {
                 desiredConsistencyLevel,
                 configs,
                 cosmosAuthorizationTokenResolver,
-                cosmosKeyCredential,
+                credential,
                 sessionCapturingOverride,
                 transportClientSharing,
                 contentResponseOnWriteEnabled);
@@ -260,8 +260,8 @@ public void setCosmosAuthorizationTokenResolver(CosmosAuthorizationTokenResolver
             this.cosmosAuthorizationTokenResolver = cosmosAuthorizationTokenResolver;
         }
 
-        public CosmosKeyCredential getCosmosKeyCredential() {
-            return cosmosKeyCredential;
+        public AzureKeyCredential getCredential() {
+            return credential;
         }
     }
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/BaseAuthorizationTokenProvider.java

@@ -3,11 +3,10 @@
 
 package com.azure.cosmos.implementation;
 
-import com.azure.cosmos.BridgeInternal;
-import com.azure.cosmos.CosmosKeyCredential;
+import com.azure.core.credential.AzureKeyCredential;
 import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
-import com.azure.cosmos.models.ModelBridgeInternal;
 import com.azure.cosmos.implementation.directconnectivity.HttpUtils;
+import com.azure.cosmos.models.ModelBridgeInternal;
 
 import javax.crypto.Mac;
 import javax.crypto.SecretKey;
@@ -28,14 +27,14 @@
 public class BaseAuthorizationTokenProvider implements AuthorizationTokenProvider {
 
     private static final String AUTH_PREFIX = "type=master&ver=1.0&sig=";
-    private final CosmosKeyCredential cosmosKeyCredential;
+    private final AzureKeyCredential credential;
     private final Mac macInstance;
 
     //  stores current master key's hashcode for performance reasons.
     private int masterKeyHashCode;
 
-    public BaseAuthorizationTokenProvider(CosmosKeyCredential cosmosKeyCredential) {
-        this.cosmosKeyCredential = cosmosKeyCredential;
+    public BaseAuthorizationTokenProvider(AzureKeyCredential credential) {
+        this.credential = credential;
         this.macInstance = getMacInstance();
     }
 
@@ -120,7 +119,7 @@ public String generateKeyAuthorizationSignature(RequestVerb verb,
             throw new IllegalArgumentException("headers");
         }
 
-        if (StringUtils.isEmpty(this.cosmosKeyCredential.getKey())) {
+        if (StringUtils.isEmpty(this.credential.getKey())) {
             throw new IllegalArgumentException("key credentials cannot be empty");
         }
 
@@ -246,11 +245,11 @@ private String generateKeyAuthorizationSignatureNew(RequestVerb verb, String res
     }
 
     private Mac getMacInstance() {
-        int masterKeyLatestHashCode = BridgeInternal.getHashCode(this.cosmosKeyCredential);
+        int masterKeyLatestHashCode = this.credential.getKey().hashCode();
 
         //  Master key has changed, or this is the first time we are getting mac instance
         if (masterKeyLatestHashCode != this.masterKeyHashCode) {
-            byte[] masterKeyBytes = this.cosmosKeyCredential.getKey().getBytes(StandardCharsets.UTF_8);
+            byte[] masterKeyBytes = this.credential.getKey().getBytes(StandardCharsets.UTF_8);
             byte[] masterKeyDecodedBytes = Utils.Base64Decoder.decode(masterKeyBytes);
             SecretKey signingKey = new SecretKeySpec(masterKeyDecodedBytes, "HMACSHA256");
             try {