content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package com.datastax.oss.driver.internal.core.context;
  
  import com.datastax.dse.driver.api.core.config.DseDriverOption;
  import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
  import com.datastax.oss.driver.api.core.loadbalancing.LoadBalancingPolicy;
  import com.datastax.oss.driver.api.core.session.Session;
  import com.datastax.oss.driver.api.core.uuid.Uuids;
  import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
  import com.datastax.oss.protocol.internal.request.Startup;
  import com.datastax.oss.protocol.internal.util.collection.NullAllowingImmutableMap;
  import com.fasterxml.jackson.databind.ObjectMapper;
  import edu.umd.cs.findbugs.annotations.Nullable;
  import java.util.Map;
  import java.util.Optional;
  import java.util.UUID;
  import net.jcip.annotations.Immutable;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  @Immutable
  public class StartupOptionsBuilder {
  
    public static final String DRIVER_NAME_KEY = "DRIVER_NAME";
    public static final String DRIVER_VERSION_KEY = "DRIVER_VERSION";
    public static final String DRIVER_BAGGAGE = "DRIVER_BAGGAGE";
    public static final String APPLICATION_NAME_KEY = "APPLICATION_NAME";
    public static final String APPLICATION_VERSION_KEY = "APPLICATION_VERSION";
    public static final String CLIENT_ID_KEY = "CLIENT_ID";
  
    private static final Logger LOG = LoggerFactory.getLogger(StartupOptionsBuilder.class);
    private static final ObjectMapper mapper = new ObjectMapper();
  
    protected final InternalDriverContext context;
    private UUID clientId;
    private String applicationName;
    private String applicationVersion;
  
    public StartupOptionsBuilder(InternalDriverContext context) {
      this.context = context;
    }
  
    /**
     * Sets the client ID to be sent in the Startup message options.
     *
     * <p>If this method is not invoked, or the id passed in is null, a random {@link UUID} will be
     * generated and used by default.
     */
    public StartupOptionsBuilder withClientId(@Nullable UUID clientId) {
      this.clientId = clientId;
      return this;
    }
  
    /**
     * Sets the client application name to be sent in the Startup message options.
     *
     * <p>If this method is not invoked, or the name passed in is null, no application name option
     * will be sent in the startup message options.
     */
    public StartupOptionsBuilder withApplicationName(@Nullable String applicationName) {
      this.applicationName = applicationName;
      return this;
    }
  
    /**
     * Sets the client application version to be sent in the Startup message options.
     *
     * <p>If this method is not invoked, or the name passed in is null, no application version option
     * will be sent in the startup message options.
     */
    public StartupOptionsBuilder withApplicationVersion(@Nullable String applicationVersion) {
      this.applicationVersion = applicationVersion;
      return this;
    }
  
    /**
     * Builds a map of options to send in a Startup message.
     *
     * <p>The default set of options are built here and include {@link
     * com.datastax.oss.protocol.internal.request.Startup#COMPRESSION_KEY} (if the context passed in
     * has a compressor/algorithm set), and the driver's {@link #DRIVER_NAME_KEY} and {@link
     * #DRIVER_VERSION_KEY}. The {@link com.datastax.oss.protocol.internal.request.Startup}
     * constructor will add {@link
     * com.datastax.oss.protocol.internal.request.Startup#CQL_VERSION_KEY}.
     *
     * @return Map of Startup Options.
     */
    public Map<String, String> build() {
      DriverExecutionProfile config = context.getConfig().getDefaultProfile();
  
      NullAllowingImmutableMap.Builder<String, String> builder = NullAllowingImmutableMap.builder(3);
      // add compression (if configured) and driver name and version
      String compressionAlgorithm = context.getCompressor().algorithm();
      if (compressionAlgorithm != null && !compressionAlgorithm.trim().isEmpty()) {
        builder.put(Startup.COMPRESSION_KEY, compressionAlgorithm.trim());
      }
      builder.put(DRIVER_NAME_KEY, getDriverName()).put(DRIVER_VERSION_KEY, getDriverVersion());
  
      // Add Insights entries, falling back to generation / config if no programmatic values provided:
      if (clientId == null) {
        clientId = Uuids.random();
      }
      builder.put(CLIENT_ID_KEY, clientId.toString());
      if (applicationName == null) {
        applicationName = config.getString(DseDriverOption.APPLICATION_NAME, null);
      }
      if (applicationName != null) {
        builder.put(APPLICATION_NAME_KEY, applicationName);
      }
      if (applicationVersion == null) {
        applicationVersion = config.getString(DseDriverOption.APPLICATION_VERSION, null);
      }
      if (applicationVersion != null) {
        builder.put(APPLICATION_VERSION_KEY, applicationVersion);
      }
      driverBaggage().ifPresent(s -> builder.put(DRIVER_BAGGAGE, s));
  
      return builder.build();
    }
  
    /**
     * Returns this driver's name.
     *
     * <p>By default, this method will pull from the bundled Driver.properties file. Subclasses should
     * override this method if they need to report a different Driver name on Startup.
     */
    protected String getDriverName() {
      return Session.OSS_DRIVER_COORDINATES.getName();
    }
  
    /**
     * Returns this driver's version.
     *
     * <p>By default, this method will pull from the bundled Driver.properties file. Subclasses should
     * override this method if they need to report a different Driver version on Startup.
     */
    protected String getDriverVersion() {
      return Session.OSS_DRIVER_COORDINATES.getVersion().toString();
    }
  
    private Optional<String> driverBaggage() {
      ImmutableMap.Builder<String, Object> builder = new ImmutableMap.Builder<>();
      for (Map.Entry<String, LoadBalancingPolicy> entry :
          context.getLoadBalancingPolicies().entrySet()) {
        Map<String, ?> config = entry.getValue().getStartupConfiguration();
        if (!config.isEmpty()) {
          builder.put(entry.getKey(), config);
        }
      }
      try {
        return Optional.of(mapper.writeValueAsString(builder.build()));
      } catch (Exception e) {
        LOG.warn("Failed to construct startup driver baggage", e);
        return Optional.empty();
      }
    }
  }
  