Examples with CosmosClientBuilder com.azure.cosmos.CosmosClientBuilder used on opensource projects

Example 1 with CosmosClientBuilder

use of com.azure.cosmos.CosmosClientBuilder 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);
    }
}