use of com.azure.cosmos.GatewayConnectionConfig in project YCSB by brianfrankcooper.
the class AzureCosmosClient method initAzureCosmosClient.
private void initAzureCosmosClient() throws DBException {
// Connection properties
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 uri = this.getStringProperty("azurecosmos.uri", null);
if (primaryKey == null || primaryKey.isEmpty()) {
throw new DBException("Missing uri required to connect to the database.");
}
AzureCosmosClient.userAgent = this.getStringProperty("azurecosmos.userAgent", DEFAULT_USER_AGENT);
AzureCosmosClient.useUpsert = this.getBooleanProperty("azurecosmos.useUpsert", DEFAULT_USE_UPSERT);
AzureCosmosClient.databaseName = this.getStringProperty("azurecosmos.databaseName", DEFAULT_DATABASE_NAME);
AzureCosmosClient.maxDegreeOfParallelism = this.getIntProperty("azurecosmos.maxDegreeOfParallelism", DEFAULT_MAX_DEGREE_OF_PARALLELISM);
AzureCosmosClient.maxBufferedItemCount = this.getIntProperty("azurecosmos.maxBufferedItemCount", DEFAULT_MAX_BUFFERED_ITEM_COUNT);
AzureCosmosClient.preferredPageSize = this.getIntProperty("azurecosmos.preferredPageSize", DEFAULT_PREFERRED_PAGE_SIZE);
AzureCosmosClient.includeExceptionStackInLog = this.getBooleanProperty("azurecosmos.includeExceptionStackInLog", DEFAULT_INCLUDE_EXCEPTION_STACK_IN_LOG);
ConsistencyLevel consistencyLevel = ConsistencyLevel.valueOf(this.getStringProperty("azurecosmos.consistencyLevel", DEFAULT_CONSISTENCY_LEVEL.toString().toUpperCase()));
boolean useGateway = this.getBooleanProperty("azurecosmos.useGateway", DEFAULT_USE_GATEWAY);
ThrottlingRetryOptions retryOptions = new ThrottlingRetryOptions();
int maxRetryAttemptsOnThrottledRequests = this.getIntProperty("azurecosmos.maxRetryAttemptsOnThrottledRequests", -1);
if (maxRetryAttemptsOnThrottledRequests != -1) {
retryOptions.setMaxRetryAttemptsOnThrottledRequests(maxRetryAttemptsOnThrottledRequests);
}
// Direct connection config options.
DirectConnectionConfig directConnectionConfig = new DirectConnectionConfig();
int directMaxConnectionsPerEndpoint = this.getIntProperty("azurecosmos.directMaxConnectionsPerEndpoint", -1);
if (directMaxConnectionsPerEndpoint != -1) {
directConnectionConfig.setMaxConnectionsPerEndpoint(directMaxConnectionsPerEndpoint);
}
int directIdleConnectionTimeoutInSeconds = this.getIntProperty("azurecosmos.directIdleConnectionTimeoutInSeconds", -1);
if (directIdleConnectionTimeoutInSeconds != -1) {
directConnectionConfig.setIdleConnectionTimeout(Duration.ofSeconds(directIdleConnectionTimeoutInSeconds));
}
// Gateway connection config options.
GatewayConnectionConfig gatewayConnectionConfig = new GatewayConnectionConfig();
int gatewayMaxConnectionPoolSize = this.getIntProperty("azurecosmos.gatewayMaxConnectionPoolSize", -1);
if (gatewayMaxConnectionPoolSize != -1) {
gatewayConnectionConfig.setMaxConnectionPoolSize(gatewayMaxConnectionPoolSize);
}
int gatewayIdleConnectionTimeoutInSeconds = this.getIntProperty("azurecosmos.gatewayIdleConnectionTimeoutInSeconds", -1);
if (gatewayIdleConnectionTimeoutInSeconds != -1) {
gatewayConnectionConfig.setIdleConnectionTimeout(Duration.ofSeconds(gatewayIdleConnectionTimeoutInSeconds));
}
try {
LOGGER.info("Creating Cosmos DB client {}, useGateway={}, consistencyLevel={}," + " maxRetryAttemptsOnThrottledRequests={}, maxRetryWaitTimeInSeconds={}" + " useUpsert={}, maxDegreeOfParallelism={}, maxBufferedItemCount={}, preferredPageSize={}", uri, useGateway, consistencyLevel.toString(), retryOptions.getMaxRetryAttemptsOnThrottledRequests(), 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);
if (useGateway) {
builder = builder.gatewayMode(gatewayConnectionConfig);
} else {
builder = builder.directMode(directConnectionConfig);
}
AzureCosmosClient.client = builder.buildClient();
LOGGER.info("Azure Cosmos DB connection created to {}", uri);
} catch (IllegalArgumentException e) {
if (!AzureCosmosClient.includeExceptionStackInLog) {
e = null;
}
throw new DBException("Illegal argument passed in. Check the format of your parameters.", e);
}
AzureCosmosClient.containerCache = new ConcurrentHashMap<>();
// Verify the database exists
try {
AzureCosmosClient.database = AzureCosmosClient.client.getDatabase(databaseName);
AzureCosmosClient.database.read();
} catch (CosmosException e) {
if (!AzureCosmosClient.includeExceptionStackInLog) {
e = null;
}
throw new DBException("Invalid database name (" + AzureCosmosClient.databaseName + ") or failed to read database.", e);
}
}
Aggregations