From 0f64c8df3722012aab48fa2415df04f2f9c0c3e4 Mon Sep 17 00:00:00 2001 From: Daniel Salazar Date: Thu, 20 Jun 2024 14:52:14 -0600 Subject: [PATCH] Adding msi ua client id property --- azurecosmos/README.md | 12 ++- .../java/site/ycsb/db/AzureCosmosClient.java | 81 +++++++++++++++---- 2 files changed, 75 insertions(+), 18 deletions(-) diff --git a/azurecosmos/README.md b/azurecosmos/README.md index 25838f223e..3a7cb66ffe 100644 --- a/azurecosmos/README.md +++ b/azurecosmos/README.md @@ -33,12 +33,14 @@ get. You can override the default database name with the azurecosmos.databaseName configuration value for side-by-side benchmarking. -You must set the uri and the primaryKey in the azurecosmos.properties file in the commands below. +You must set the uri and the primaryKey or userAssignedIdentityClientId in the azurecosmos.properties file in the commands below. $YCSB_HOME/bin/ycsb load azurecosmos -P workloads/workloada -P azurecosmos/conf/azurecosmos.properties $YCSB_HOME/bin/ycsb run azurecosmos -P workloads/workloada -P azurecosmos/conf/azurecosmos.properties -Optionally you can set the uri and primaryKey as follows: +Optionally you can set the uri and primaryKey OR userAssignedIdentityClientId as follows: $YCSB_HOME/bin/ycsb load azurecosmos -P workloads/workloada -p azurecosmos.primaryKey= -p azurecosmos.uri= +Or + $YCSB_HOME/bin/ycsb load azurecosmos -P workloads/workloada -p azurecosmos.userAssignedIdentityClientId= -p azurecosmos.uri= ### 2. Cosmos DB Configuration Parameters @@ -53,6 +55,12 @@ Optionally you can set the uri and primaryKey as follows: doing read only workloads you can substitute the readonly key from the portal. +OR + +- azurecosmos.userAssignedIdentityClientId < key string > : + - Obtained from the UA resource in the portal, named as Client ID. + + #### Options parameters - azurecosmos.databaseName < name string > : diff --git a/azurecosmos/src/main/java/site/ycsb/db/AzureCosmosClient.java b/azurecosmos/src/main/java/site/ycsb/db/AzureCosmosClient.java index 37355f603d..e539c56e9c 100644 --- a/azurecosmos/src/main/java/site/ycsb/db/AzureCosmosClient.java +++ b/azurecosmos/src/main/java/site/ycsb/db/AzureCosmosClient.java @@ -37,6 +37,8 @@ import com.azure.cosmos.models.SqlParameter; import com.azure.cosmos.models.SqlQuerySpec; import com.azure.cosmos.util.CosmosPagedIterable; +import com.azure.identity.DefaultAzureCredentialBuilder; +import com.azure.identity.ManagedIdentityCredentialBuilder; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.node.ObjectNode; @@ -86,6 +88,8 @@ public class AzureCosmosClient extends DB { private static final int DEFAULT_DIAGNOSTICS_LATENCY_THRESHOLD_IN_MS = -1; private static final boolean DEFAULT_INCLUDE_EXCEPTION_STACK_IN_LOG = false; private static final String DEFAULT_USER_AGENT = "azurecosmos-ycsb"; + private static final String AAD_AUTH_METHOD = "AAD"; + private static final String MASTER_KEY_AUTH_METHOD = "MASTER_KEY"; private static final Logger LOGGER = LoggerFactory.getLogger(AzureCosmosClient.class); private static final Marker CREATE_DIAGNOSTIC = MarkerFactory.getMarker("CREATE_DIAGNOSTIC"); @@ -152,10 +156,11 @@ public void init() throws DBException { private void initAzureCosmosClient() throws DBException { - // Connection properties + // Connection properties. Requires at least the primary key OR the user assigned identity id. String primaryKey = this.getStringProperty("azurecosmos.primaryKey", null); - if (primaryKey == null || primaryKey.isEmpty()) { - throw new DBException("Missing primary key required to connect to the database."); + String userAssignedIdentityClientId = this.getStringProperty("azurecosmos.userAssignedIdentityClientId", null); + if ( (primaryKey == null || primaryKey.isEmpty()) && (userAssignedIdentityClientId == null || userAssignedIdentityClientId.isEmpty()) ) { + throw new DBException("Missing primary key or user assigned identity id required to connect to the database."); } String uri = this.getStringProperty("azurecosmos.uri", null); @@ -247,19 +252,22 @@ private void initAzureCosmosClient() throws DBException { retryOptions.getMaxRetryWaitTime().toMillis() / 1000, AzureCosmosClient.useUpsert, AzureCosmosClient.maxDegreeOfParallelism, AzureCosmosClient.maxBufferedItemCount, AzureCosmosClient.preferredPageSize); - - CosmosClientBuilder builder = new CosmosClientBuilder() - .endpoint(uri) - .key(primaryKey) - .throttlingRetryOptions(retryOptions) - .consistencyLevel(consistencyLevel) - .userAgentSuffix(userAgent) - .clientTelemetryConfig(new CosmosClientTelemetryConfig() - .diagnosticsThresholds( - new CosmosDiagnosticsThresholds() - .setPointOperationLatencyThreshold(Duration.ofMillis(pointOperationLatencyThresholdInMS)) - .setNonPointOperationLatencyThreshold(Duration.ofMillis(nonPointOperationLatencyThresholdInMS)) - .setRequestChargeThreshold(requestChargeThreshold))); + + String authMethod = ""; + String token = ""; + if (primaryKey != null && !primaryKey.isEmpty()) + { + authMethod = AAD_AUTH_METHOD; + token = primaryKey; + } + else if (userAssignedIdentityClientId != null && !userAssignedIdentityClientId.isEmpty()) + { + authMethod = MASTER_KEY_AUTH_METHOD; + token = userAssignedIdentityClientId; + } + + // Initialize the CosmosClientBuilder based on the authentication method. + CosmosClientBuilder builder = getCosmosClientBuilder(authMethod, uri, token); if (useGateway) { builder = builder.gatewayMode(gatewayConnectionConfig); @@ -300,6 +308,47 @@ private void initAzureCosmosClient() throws DBException { } } + // Initialize the CosmosClientBuilder based on the authentication method. + private static CosmosClientBuilder getCosmosClientBuilder(String authMethod, String uri, String token) { + if (authMethod.equals(AAD_AUTH_METHOD)) + { + // Gets token from MSI UA + ManagedIdentityCredentialBuilder credentialBuilder = new ManagedIdentityCredentialBuilder(); + credentialBuilder.clientId(token); + + return new CosmosClientBuilder() + .endpoint(uri) + .credential(credentialBuilder.build()) + .throttlingRetryOptions(retryOptions) + .consistencyLevel(consistencyLevel) + .userAgentSuffix(userAgent) + .clientTelemetryConfig(new CosmosClientTelemetryConfig() + .diagnosticsThresholds( + new CosmosDiagnosticsThresholds() + .setPointOperationLatencyThreshold(Duration.ofMillis(pointOperationLatencyThresholdInMS)) + .setNonPointOperationLatencyThreshold(Duration.ofMillis(nonPointOperationLatencyThresholdInMS)) + .setRequestChargeThreshold(requestChargeThreshold))); + } + else if (authMethod.equals(MASTER_KEY_AUTH_METHOD)) + { + return new CosmosClientBuilder() + .endpoint(uri) + .key(token) + .throttlingRetryOptions(retryOptions) + .consistencyLevel(consistencyLevel) + .userAgentSuffix(userAgent) + .clientTelemetryConfig(new CosmosClientTelemetryConfig() + .diagnosticsThresholds( + new CosmosDiagnosticsThresholds() + .setPointOperationLatencyThreshold(Duration.ofMillis(pointOperationLatencyThresholdInMS)) + .setNonPointOperationLatencyThreshold(Duration.ofMillis(nonPointOperationLatencyThresholdInMS)) + .setRequestChargeThreshold(requestChargeThreshold))); + } + else { + throw new DBException("Invalid authentication method. Supported methods are AAD and MASTER_KEY."); + } + } + private String getStringProperty(String propertyName, String defaultValue) { return getProperties().getProperty(propertyName, defaultValue); }