Cosmos Key Credential (Azure#4885)

Cosmos Key Credential holds the key credentials, and supports key rotations. User can update the key in CosmosKeyCredential object, and that will be reflected in SDK on the fly.

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

@@ -12,19 +12,19 @@
  * DO NOT USE. For internal use only by the SDK. These methods might break at any time. No support will be provided.
  */
 public class CosmosBridgeInternal {
-    
+
     public static DocumentCollection toDocumentCollection(CosmosContainerProperties cosmosContainerProperties) {
         return new DocumentCollection(cosmosContainerProperties.toJson());
     }
 
     public static AsyncDocumentClient getAsyncDocumentClient(CosmosClient client) {
         return client.getDocClientWrapper();
     }
-    
+
     public static CosmosDatabase getCosmosDatabaseWithNewClient(CosmosDatabase cosmosDatabase, CosmosClient client) {
         return new CosmosDatabase(cosmosDatabase.id(), client);
     }
-    
+
     public static CosmosContainer getCosmosContainerWithNewClient(CosmosContainer cosmosContainer, CosmosDatabase cosmosDatabase, CosmosClient client) {
         return new CosmosContainer(cosmosContainer.id(), CosmosBridgeInternal.getCosmosDatabaseWithNewClient(cosmosDatabase, client));
     }

sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/CosmosClient.java

@@ -29,24 +29,27 @@ public class CosmosClient implements AutoCloseable {
     private final ConsistencyLevel desiredConsistencyLevel;
     private final List<Permission> permissions;
     private final TokenResolver tokenResolver;
+    private final CosmosKeyCredential cosmosKeyCredential;
 
 
      CosmosClient(CosmosClientBuilder builder) {
          this.configs = builder.configs();
          this.serviceEndpoint = builder.endpoint();
-        this.keyOrResourceToken = builder.key();
-        this.connectionPolicy = builder.connectionPolicy();
-        this.desiredConsistencyLevel = builder.consistencyLevel();
-        this.permissions = builder.permissions();
-        this.tokenResolver = builder.tokenResolver();
-        this.asyncDocumentClient = new AsyncDocumentClient.Builder()
-                .withServiceEndpoint(this.serviceEndpoint)
-                .withMasterKeyOrResourceToken(this.keyOrResourceToken)
-                .withConnectionPolicy(this.connectionPolicy)
-                .withConsistencyLevel(this.desiredConsistencyLevel)
-                .withConfigs(this.configs)
-                .withTokenResolver(this.tokenResolver)
-                .build();
+         this.keyOrResourceToken = builder.key();
+         this.connectionPolicy = builder.connectionPolicy();
+         this.desiredConsistencyLevel = builder.consistencyLevel();
+         this.permissions = builder.permissions();
+         this.tokenResolver = builder.tokenResolver();
+         this.cosmosKeyCredential = builder.cosmosKeyCredential();
+         this.asyncDocumentClient = new AsyncDocumentClient.Builder()
+             .withServiceEndpoint(this.serviceEndpoint)
+             .withMasterKeyOrResourceToken(this.keyOrResourceToken)
+             .withConnectionPolicy(this.connectionPolicy)
+             .withConsistencyLevel(this.desiredConsistencyLevel)
+             .withConfigs(this.configs)
+             .withTokenResolver(this.tokenResolver)
+             .withCosmosKeyCredential(this.cosmosKeyCredential)
+             .build();
     }
 
     AsyncDocumentClient getContextClient() {
@@ -121,10 +124,18 @@ TokenResolver getTokenResolver() {
         return tokenResolver;
     }
 
+    /**
+     * Gets the cosmos key credential
+     * @return cosmos key credential
+     */
+    CosmosKeyCredential cosmosKeyCredential() {
+        return cosmosKeyCredential;
+    }
+
     /**
      * CREATE a Database if it does not already exist on the service
-     * 
-     * The {@link Mono} upon successful completion will contain a single cosmos database response with the 
+     *
+     * The {@link Mono} upon successful completion will contain a single cosmos database response with the
      * created or existing database.
      * @param databaseSettings CosmosDatabaseProperties
      * @return a {@link Mono} containing the cosmos database response with the created or existing database or
@@ -136,7 +147,7 @@ public Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(CosmosDatabaseProp
 
     /**
      * CREATE a Database if it does not already exist on the service
-     * The {@link Mono} upon successful completion will contain a single cosmos database response with the 
+     * The {@link Mono} upon successful completion will contain a single cosmos database response with the
      * created or existing database.
      * @param id the id of the database
      * @return a {@link Mono} containing the cosmos database response with the created or existing database or
@@ -145,7 +156,7 @@ public Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(CosmosDatabaseProp
     public Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(String id) {
         return createDatabaseIfNotExistsInternal(getDatabase(id));
     }
-    
+
     private Mono<CosmosDatabaseResponse> createDatabaseIfNotExistsInternal(CosmosDatabase database){
         return database.read().onErrorResume(exception -> {
             if (exception instanceof CosmosClientException) {
@@ -160,12 +171,12 @@ private Mono<CosmosDatabaseResponse> createDatabaseIfNotExistsInternal(CosmosDat
 
     /**
      * Creates a database.
-     * 
-     * After subscription the operation will be performed. 
-     * The {@link Mono} upon successful completion will contain a single resource response with the 
+     *
+     * After subscription the operation will be performed.
+     * The {@link Mono} upon successful completion will contain a single resource response with the
      *      created database.
      * In case of failure the {@link Mono} will error.
-     * 
+     *
      * @param databaseSettings {@link CosmosDatabaseProperties}
      * @param options {@link CosmosDatabaseRequestOptions}
      * @return an {@link Mono} containing the single cosmos database response with the created database or an error.
@@ -183,12 +194,12 @@ public Mono<CosmosDatabaseResponse> createDatabase(CosmosDatabaseProperties data
 
     /**
      * Creates a database.
-     * 
-     * After subscription the operation will be performed. 
-     * The {@link Mono} upon successful completion will contain a single resource response with the 
+     *
+     * After subscription the operation will be performed.
+     * The {@link Mono} upon successful completion will contain a single resource response with the
      *      created database.
      * In case of failure the {@link Mono} will error.
-     * 
+     *
      * @param databaseSettings {@link CosmosDatabaseProperties}
      * @return an {@link Mono} containing the single cosmos database response with the created database or an error.
      */
@@ -198,12 +209,12 @@ public Mono<CosmosDatabaseResponse> createDatabase(CosmosDatabaseProperties data
 
     /**
      * Creates a database.
-     * 
-     * After subscription the operation will be performed. 
-     * The {@link Mono} upon successful completion will contain a single resource response with the 
+     *
+     * After subscription the operation will be performed.
+     * The {@link Mono} upon successful completion will contain a single resource response with the
      *      created database.
      * In case of failure the {@link Mono} will error.
-     * 
+     *
      * @param id id of the database
      * @return a {@link Mono} containing the single cosmos database response with the created database or an error.
      */
@@ -275,27 +286,27 @@ public Mono<CosmosDatabaseResponse> createDatabase(String id, int throughput) {
 
     /**
      * Reads all databases.
-     * 
-     * After subscription the operation will be performed. 
+     *
+     * After subscription the operation will be performed.
      * The {@link Flux} will contain one or several feed response of the read databases.
      * In case of failure the {@link Flux} will error.
-     * 
+     *
      * @param options {@link FeedOptions}
      * @return a {@link Flux} containing one or several feed response pages of read databases or an error.
      */
     public Flux<FeedResponse<CosmosDatabaseProperties>> readAllDatabases(FeedOptions options) {
         return getDocClientWrapper().readDatabases(options)
-                .map(response-> BridgeInternal.createFeedResponse(CosmosDatabaseProperties.getFromV2Results(response.results()),
+                                    .map(response-> BridgeInternal.createFeedResponse(CosmosDatabaseProperties.getFromV2Results(response.results()),
                         response.responseHeaders()));
     }
 
     /**
      * Reads all databases.
-     * 
-     * After subscription the operation will be performed. 
+     *
+     * After subscription the operation will be performed.
      * The {@link Flux} will contain one or several feed response of the read databases.
      * In case of failure the {@link Flux} will error.
-     * 
+     *
      * @return a {@link Flux} containing one or several feed response pages of read databases or an error.
      */
     public Flux<FeedResponse<CosmosDatabaseProperties>> readAllDatabases() {
@@ -331,7 +342,7 @@ public Flux<FeedResponse<CosmosDatabaseProperties>> queryDatabases(String query,
      */
     public Flux<FeedResponse<CosmosDatabaseProperties>> queryDatabases(SqlQuerySpec querySpec, FeedOptions options){
         return getDocClientWrapper().queryDatabases(querySpec, options)
-                .map(response-> BridgeInternal.createFeedResponse(
+                                    .map(response-> BridgeInternal.createFeedResponse(
                         CosmosDatabaseProperties.getFromV2Results(response.results()),
                         response.responseHeaders()));
     }

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

@@ -4,6 +4,7 @@
 
 import com.azure.data.cosmos.internal.Configs;
 import com.azure.data.cosmos.internal.Permission;
+import org.apache.commons.lang3.StringUtils;
 
 import java.util.List;
 
@@ -33,6 +34,7 @@ public class CosmosClientBuilder {
     private ConsistencyLevel desiredConsistencyLevel;
     private List<Permission> permissions;
     private TokenResolver tokenResolver;
+    private CosmosKeyCredential cosmosKeyCredential;
 
     CosmosClientBuilder() {
     }
@@ -137,15 +139,15 @@ public CosmosClientBuilder permissions(List<Permission> permissions) {
     }
 
     /**
-     * Gets the (@link ConsistencyLevel) to be used
+     * Gets the {@link ConsistencyLevel} to be used
      * @return the consistency level
      */
     public ConsistencyLevel consistencyLevel() {
         return this.desiredConsistencyLevel;
     }
 
     /**
-     * Sets the (@link ConsistencyLevel) to be used
+     * Sets the {@link ConsistencyLevel} to be used
      * @param desiredConsistencyLevel {@link ConsistencyLevel}
      * @return current Builder
      */
@@ -163,7 +165,7 @@ public ConnectionPolicy connectionPolicy() {
     }
 
     /**
-     * Sets the (@link ConnectionPolicy) to be used
+     * Sets the {@link ConnectionPolicy} to be used
      * @param connectionPolicy {@link ConnectionPolicy}
      * @return current Builder
      */
@@ -172,6 +174,24 @@ public CosmosClientBuilder connectionPolicy(ConnectionPolicy connectionPolicy) {
         return this;
     }
 
+    /**
+     * Gets the {@link CosmosKeyCredential} to be used
+     * @return cosmosKeyCredential
+     */
+    public CosmosKeyCredential cosmosKeyCredential() {
+        return cosmosKeyCredential;
+    }
+
+    /**
+     * Sets the {@link CosmosKeyCredential} to be used
+     * @param cosmosKeyCredential {@link CosmosKeyCredential}
+     * @return current builder
+     */
+    public CosmosClientBuilder cosmosKeyCredential(CosmosKeyCredential cosmosKeyCredential) {
+        this.cosmosKeyCredential = cosmosKeyCredential;
+        return this;
+    }
+
     /**
      * Builds a cosmos configuration object with the provided properties
      * @return CosmosClient
@@ -180,8 +200,11 @@ public CosmosClient build() {
 
         ifThrowIllegalArgException(this.serviceEndpoint == null, "cannot build client without service endpoint");
         ifThrowIllegalArgException(
-            this.keyOrResourceToken == null && (permissions == null || permissions.isEmpty()) && this.tokenResolver == null,
-            "cannot build client without any one of key, resource token, permissions, and token resolver");
+            this.keyOrResourceToken == null && (permissions == null || permissions.isEmpty())
+                && this.tokenResolver == null && this.cosmosKeyCredential == null,
+            "cannot build client without any one of key, resource token, permissions, token resolver, and cosmos key credential");
+        ifThrowIllegalArgException(cosmosKeyCredential != null && StringUtils.isEmpty(cosmosKeyCredential.key()),
+            "cannot build client without key credential");
 
         return new CosmosClient(this);
     }

sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/CosmosDatabase.java

@@ -26,7 +26,7 @@ public class CosmosDatabase {
 
     /**
      * Get the id of the CosmosDatabase
-     * 
+     *
      * @return the id of the CosmosDatabase
      */
     public String id() {
@@ -35,7 +35,7 @@ public String id() {
 
     /**
      * Set the id of the CosmosDatabase
-     * 
+     *
      * @param id the id of the CosmosDatabase
      * @return the same CosmosConflict that had the id set
      */
@@ -415,7 +415,7 @@ public Flux<FeedResponse<CosmosContainerProperties>> queryContainers(SqlQuerySpe
 
     /**
      * Gets a CosmosContainer object without making a service call
-     * 
+     *
      * @param id id of the container
      * @return Cosmos Container
      */
@@ -622,4 +622,4 @@ String getLink() {
         builder.append(id());
         return builder.toString();
     }
-}
+}

sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/CosmosKeyCredential.java

@@ -0,0 +1,49 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.data.cosmos;
+
+/**
+ * Cosmos Key Credential is used to store key credentials, in order to support dynamic key rotation.
+ * Singleton instance should be used to support multiple keys.
+ * Azure client library for Cosmos ensures to use the updated key provided in the same singleton instance
+ * which was used when building {@link CosmosClient}
+ */
+public class CosmosKeyCredential {
+
+    private String key;
+
+    //  Stores key's hashcode for performance improvements
+    private int keyHashCode;
+
+    public CosmosKeyCredential(String key) {
+        this.key = key;
+        this.keyHashCode = key.hashCode();
+    }
+
+    /**
+     * Returns the key stored in Cosmos Key Credential
+     * @return key
+     */
+    public String key() {
+        return key;
+    }
+
+    /**
+     * Sets the key to be used in CosmosKeyCredential
+     * @param key key to be used in CosmosKeyCredential
+     * @return current CosmosKeyCredential
+     */
+    public CosmosKeyCredential key(String key) {
+        this.key = key;
+        this.keyHashCode = key.hashCode();
+        return this;
+    }
+
+    /**
+     * CosmosKeyCredential stores the computed hashcode of the key for performance improvements.
+     * @return hashcode of the key
+     */
+    public int keyHashCode() {
+        return this.keyHashCode;
+    }
+}