private void initAzureCosmosClient()

in azurecosmos/src/main/java/site/ycsb/db/AzureCosmosClient.java [156:287]
102 lines of code
19 McCabe index (conditional complexity)

  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 (uri == null || uri.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.diagnosticsLatencyThresholdInMS = this.getIntProperty(
        "azurecosmos.diagnosticsLatencyThresholdInMS",
        DEFAULT_DIAGNOSTICS_LATENCY_THRESHOLD_IN_MS);

    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));
    }

    String preferredRegions = this.getStringProperty("azurecosmos.preferredRegionList", null);
    List<String> preferredRegionList = null;
    if (StringUtils.isNotEmpty(preferredRegions)) {
      preferredRegions = preferredRegions.trim();
      preferredRegionList = new ArrayList<>(Arrays.asList(preferredRegions.split(",")));
    }
    
    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);
      }

      if (preferredRegionList != null && preferredRegionList.size() > 0) {
        builder.preferredRegions(preferredRegionList);
      }

      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);
    }

    String appInsightConnectionString = this.getStringProperty("azurecosmos.appInsightConnectionString", null);
    if (appInsightConnectionString != null) {
      this.azureMonitorMeterRegistry = this.azureMonitorMeterRegistry(appInsightConnectionString);
      registerMeter();
    }
  }