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 org.apache.apex.malhar.contrib.cassandra;
  
  import java.io.Serializable;
  import java.net.InetAddress;
  import java.net.UnknownHostException;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.HashSet;
  import java.util.Set;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.hadoop.classification.InterfaceStability;
  
  import com.datastax.driver.core.Cluster;
  import com.datastax.driver.core.ConsistencyLevel;
  import com.datastax.driver.core.Host;
  import com.datastax.driver.core.Metadata;
  import com.datastax.driver.core.ProtocolVersion;
  import com.datastax.driver.core.QueryOptions;
  import com.datastax.driver.core.Session;
  import com.datastax.driver.core.policies.DCAwareRoundRobinPolicy;
  import com.datastax.driver.core.policies.DowngradingConsistencyRetryPolicy;
  import com.datastax.driver.core.policies.ExponentialReconnectionPolicy;
  import com.datastax.driver.core.policies.LoadBalancingPolicy;
  import com.datastax.driver.core.policies.ReconnectionPolicy;
  import com.datastax.driver.core.policies.RetryPolicy;
  import com.datastax.driver.core.policies.TokenAwarePolicy;
  
  import static com.google.common.base.Preconditions.checkNotNull;
  
  /**
   * <p>This is used to specify the connection parameters that is used to process mutations for the
   * {@link AbstractUpsertOutputOperator}.
   * The Connection can only be set for one table in a given keyspace and as such is used to define the following
   * <ul>
   * <li>Connection Parameters</li>
   * <li>Defaults to be used in case the execution context tuple {@link UpsertExecutionContext} does not want to
   * set one explicitly.</li>
   * </ul></p>
   * <p>
   * Note that the {@link ConnectionBuilder} is used to build an instance of the ConnectionStateManager.
   * An instance of this class is typically instantiated by implementing the following inside
   * {@link AbstractUpsertOutputOperator} withConnectionBuilder() method.
   * </p>
   * <p> A typical implementation of the ConnectionBuilder would like this:
   * <pre>
   *
   *     public ConnectionStateManager.ConnectionBuilder withConnectionBuilder()
   *     {
   *       return ConnectionStateManager.withNewBuilder()
   *       .withSeedNodes("localhost") // format of host1:port;host2:port;host3:port
   *       .withClusterNameAs("Test Cluster") // mandatory param
   *       .withDCNameAs("datacenter1") // mandatory param
   *      .withTableNameAs("users") // mandatory param
   *       .withKeySpaceNameAs("unittests") // mandatory param
   *       .withdefaultConsistencyLevel(ConsistencyLevel.LOCAL_ONE); // Set if required. Default of LOCAL_QUORUM
   *       // Rest of the configs are initialized to sane defaults
   *     }
   * </pre>
   * </p>
   * <p>Please refer {@link ConnectionBuilder} for details about parameters that can be used to define the connection
   * and its default behaviour </p>
   *
   * @since 3.6.0
   */
  @InterfaceStability.Evolving
  public class ConnectionStateManager implements AutoCloseable, Serializable
  {
  
    private static final long serialVersionUID = -6024334738016015213L;
    private static final long DEFAULT_MAX_DELAY_MS = 30000L;
    private static final long DEFAULT_BASE_DELAY_MS = 10000L;
    // Cassandra cluster name
    private final String clusterName;
    // Cassandra DC Name
    private final String dcName;
    // Seeds nodes. The format for specifying the host names are host1:port;host2:port
    private final String seedNodesStr;
    private final long baseDelayMs;
    private final long maxDelayMs;
    // Other operational constraints
    private final transient LoadBalancingPolicy loadBalancingPolicy;
    private final transient RetryPolicy retryPolicy;
    private final transient QueryOptions queryOptions;
    private final transient ReconnectionPolicy reconnectionPolicy;
    private final transient ProtocolVersion protocolVersion;
    private transient Logger LOG = LoggerFactory.getLogger(ConnectionStateManager.class);
    // Connection specific config elements
    private transient Cluster cluster;
    // Session specific metadata
    private transient Session session;
    private String keyspaceName;
    private String tableName;
    private int defaultTtlInSecs;
    private ConsistencyLevel defaultConsistencyLevel;
    // Standard defaults
    private boolean isTTLSet = false;
  
    private ConnectionStateManager(final ConnectionBuilder connectionBuilder)
    {
      checkNotNull(connectionBuilder, "Connection Builder passed in as Null");
      checkNotNull(connectionBuilder.clusterName, "Cluster Name not set for Cassandra");
      checkNotNull(connectionBuilder.dcName, "DataCenter Name not set for Cassandra");
      checkNotNull(connectionBuilder.seedNodesStr,
          "Seed nodes not set for Cassandra. Pattern is host1:port;host2:port");
      checkNotNull(connectionBuilder.keyspaceName, "Keyspace Name not set for Cassandra");
      checkNotNull(connectionBuilder.tableName, "Table Name not set for Cassandra");
      //Required params
      this.clusterName = connectionBuilder.clusterName;
      this.dcName = connectionBuilder.dcName;
      this.seedNodesStr = connectionBuilder.seedNodesStr;
      this.keyspaceName = connectionBuilder.keyspaceName;
      this.tableName = connectionBuilder.tableName;
  
      // optional params
      if (connectionBuilder.maxDelayMs != null) {
        maxDelayMs = connectionBuilder.maxDelayMs;
      } else {
        maxDelayMs = DEFAULT_MAX_DELAY_MS; // 30 seconds
      }
      if (connectionBuilder.baseDelayMs != null) {
        baseDelayMs = connectionBuilder.baseDelayMs;
      } else {
        baseDelayMs = DEFAULT_BASE_DELAY_MS; // 10 seconds
      }
      if (connectionBuilder.defaultTtlInSecs != null) {
        defaultTtlInSecs = connectionBuilder.defaultTtlInSecs;
        isTTLSet = true;
      }
      if (connectionBuilder.defaultConsistencyLevel != null) {
        defaultConsistencyLevel = connectionBuilder.defaultConsistencyLevel;
      } else {
        defaultConsistencyLevel = ConsistencyLevel.LOCAL_QUORUM;
      }
      if (connectionBuilder.loadBalancingPolicy != null) {
        loadBalancingPolicy = connectionBuilder.loadBalancingPolicy;
      } else {
        loadBalancingPolicy = new TokenAwarePolicy(
          DCAwareRoundRobinPolicy.builder()
            .withLocalDc(dcName)
            .build());
      }
      if (connectionBuilder.retryPolicy != null) {
        this.retryPolicy = connectionBuilder.retryPolicy;
      } else {
        retryPolicy = DowngradingConsistencyRetryPolicy.INSTANCE;
      }
      if (connectionBuilder.queryOptions != null) {
        this.queryOptions = connectionBuilder.queryOptions;
      } else {
        this.queryOptions = new QueryOptions().setConsistencyLevel(defaultConsistencyLevel);
      }
      if (connectionBuilder.reconnectionPolicy != null) {
        this.reconnectionPolicy = connectionBuilder.reconnectionPolicy;
      } else {
        this.reconnectionPolicy = new ExponentialReconnectionPolicy(baseDelayMs, maxDelayMs);
      }
      if (connectionBuilder.protocolVersion != null) {
        this.protocolVersion = connectionBuilder.protocolVersion;
      } else {
        this.protocolVersion = ProtocolVersion.NEWEST_SUPPORTED;
      }
      establishSessionWithPolicies();
    }
  
    public static ConnectionBuilder withNewBuilder()
    {
      return new ConnectionBuilder();
    }
  
    private void establishSessionWithPolicies()
    {
      Cluster.Builder clusterBuilder = Cluster.builder();
      String[] seedNodesSplit = seedNodesStr.split(";");
      Collection<InetAddress> allSeeds = new ArrayList<>();
      Set<String> allKnownPorts = new HashSet<>();
      for (String seedNode : seedNodesSplit) {
        String[] nodeAndPort = seedNode.split(":");
        if (nodeAndPort.length > 1) {
          allKnownPorts.add(nodeAndPort[1]);
        }
        try {
          allSeeds.add(InetAddress.getByName(nodeAndPort[0]));
        } catch (UnknownHostException e) {
          LOG.error(" Error while trying to initialize the seed brokers for the cassandra cluster " + e.getMessage(), e);
        }
      }
      clusterBuilder = clusterBuilder.addContactPoints(allSeeds);
      clusterBuilder
        .withClusterName(clusterName)
        // We can fail if all of the nodes die in the local DC
        .withLoadBalancingPolicy(loadBalancingPolicy)
        // Want Strong consistency
        .withRetryPolicy(retryPolicy)
        // Tolerate some nodes down
        .withQueryOptions(queryOptions)
        // Keep establishing connections after detecting a node down
        .withReconnectionPolicy(reconnectionPolicy);
      // Finally initialize the cluster info
      if (allKnownPorts.size() > 0) {
        int shortlistedPort = Integer.parseInt(allKnownPorts.iterator().next());
        clusterBuilder = clusterBuilder.withPort(shortlistedPort);
      }
      cluster = clusterBuilder.build();
      Metadata metadata = cluster.getMetadata();
      LOG.info("Connected to cluster: \n" + metadata.getClusterName());
      for (Host host : metadata.getAllHosts()) {
        LOG.info(String.format("Datacenter: %s; Host: %s; Rack: %s\n",
            host.getDatacenter(), host.getAddress(), host.getRack()));
      }
      session = cluster.connect(keyspaceName);
    }
  
    @Override
    public void close()
    {
      if (session != null) {
        session.close();
      }
      if (cluster != null) {
        cluster.close();
      }
    }
  
    public Cluster getCluster()
    {
      return cluster;
    }
  
    public void setCluster(Cluster cluster)
    {
      this.cluster = cluster;
    }
  
    public Session getSession()
    {
      return session;
    }
  
    public void setSession(Session session)
    {
      this.session = session;
    }
  
    public String getKeyspaceName()
    {
      return keyspaceName;
    }
  
    public void setKeyspaceName(String keyspaceName)
    {
      this.keyspaceName = keyspaceName;
    }
  
    public String getTableName()
    {
      return tableName;
    }
  
    public void setTableName(String tableName)
    {
      this.tableName = tableName;
    }
  
    public int getDefaultTtlInSecs()
    {
      return defaultTtlInSecs;
    }
  
    public void setDefaultTtlInSecs(int defaultTtlInSecs)
    {
      this.defaultTtlInSecs = defaultTtlInSecs;
    }
  
    public boolean isTTLSet()
    {
      return isTTLSet;
    }
  
    public void setTTLSet(boolean TTLSet)
    {
      isTTLSet = TTLSet;
    }
  
    public static class ConnectionBuilder
    {
  
      private String clusterName;
      private String dcName;
      private String seedNodesStr;
      private Long baseDelayMs; // Class to enable check for nulls and is optional
      private Long maxDelayMs; // Class to enable check for nulls and is optional
  
      private String keyspaceName;
      private String tableName;
      private Integer defaultTtlInSecs; // Class to enable checks for nulls and is optional
      private ConsistencyLevel defaultConsistencyLevel;
  
      private LoadBalancingPolicy loadBalancingPolicy;
      private RetryPolicy retryPolicy;
      private QueryOptions queryOptions;
      private ReconnectionPolicy reconnectionPolicy;
      private ProtocolVersion protocolVersion;
  
      public static final String CLUSTER_NAME_IN_PROPS_FILE = "cluster.name";
      public static final String DC_NAME_IN_PROPS_FILE = "dc.name";
      public static final String KEYSPACE_NAME_IN_PROPS_FILE = "keyspace.name";
      public static final String TABLE_NAME_IN_PROPS_FILE = "table.name";
      public static final String SEEDNODES_IN_PROPS_FILE = "seednodes";
  
      public ConnectionBuilder withClusterNameAs(String clusterName)
      {
        this.clusterName = clusterName;
        return this;
      }
  
      public ConnectionBuilder withDCNameAs(String dcName)
      {
        this.dcName = dcName;
        return this;
      }
  
      /**
       * Used to specify the seed nodes of the target cassandra cluster. Format is
       * host1:port;host2:port;host3:port
       * @param seedNodesStr
       * @return The builder instance as initially created updated with this value
         */
      public ConnectionBuilder withSeedNodes(String seedNodesStr)
      {
        this.seedNodesStr = seedNodesStr;
        return this;
      }
  
      /**
       * Used to specify the base delay while trying to set a Connection attempt policy
       * @param baseDelayMillis
       * @return The builder instance as initially created updated with this value
         */
      public ConnectionBuilder withBaseDelayMillis(long baseDelayMillis)
      {
        this.baseDelayMs = baseDelayMillis;
        return this;
      }
  
      /**
       * Used to specify the maximum time that can elapse before which a connection is given up as a failure attmept
       * @param maxDelayMillis
       * @return The builder instance as initially created updated with this value
         */
      public ConnectionBuilder withMaxDelayMillis(long maxDelayMillis)
      {
        this.maxDelayMs = maxDelayMillis;
        return this;
      }
  
      public ConnectionBuilder withKeySpaceNameAs(String keyspaceName)
      {
        this.keyspaceName = keyspaceName;
        return this;
      }
  
      public ConnectionBuilder withTableNameAs(String tableName)
      {
        this.tableName = tableName;
        return this;
      }
  
      public ConnectionBuilder withdefaultTTL(Integer defaultTtlInSecs)
      {
        this.defaultTtlInSecs = defaultTtlInSecs;
        return this;
      }
  
      /**
       * Used to specify the default consistency level when executing the mutations on the cluster.
       * Default if not set is LOCAL_QUORUM. Can be overriden at the tuple level using {@link UpsertExecutionContext}
       * @param consistencyLevel
       * @return The builder instance as initially created updated with this value
         */
      public ConnectionBuilder withdefaultConsistencyLevel(ConsistencyLevel consistencyLevel)
      {
        this.defaultConsistencyLevel = consistencyLevel;
        return this;
      }
  
      /**
       * Used to define how the nodes in the cluster will be contacted for executing a mutation.
       * The following is the default behaviour if not set.
       * 1. Use a TokenAware approach i.e. the row key is used to decide the right node to execute the mutation
       *    on the target cassandra node. i.e. One of the R-1 replicas is used as the coordinator node.
       *    This effectively balances the traffic onto all nodes of the cassandra cluster for the given
       *    Operator instance. Of course this assumes the keys are evenly distributed in the cluster
       *    which is normally the case
       * 2. Overlay TokenAware with DC aware approach - The above token aware approach is further segmented to use only
       *    the local DC for the mutation executions. Cassandras multi-DC execution will take care of the cross DC
       *    replication thus achieving the lowest possible latencies for the given mutation of writes.
       *
       * Using this effectively removes the need for an extra implementation of the Partitioning logic of the Operator
       * Nor would we need any extra logic ( for most use cases ) for dynamic partitioning implementations as the
       * underlying driver normalizes the traffic pattern anyways.
       * @param loadBalancingPolicy
       * @return The builder instance as initially created updated with this value
         */
      public ConnectionBuilder withLoadBalancingPolicy(LoadBalancingPolicy loadBalancingPolicy)
      {
        this.loadBalancingPolicy = loadBalancingPolicy;
        return this;
      }
  
      /**
       * Used to specify how queries will need to be retried in case the current in progress one fails.
       * The default is to use a DowngradingConsistency Policy i.e. first LOCAL_QUORUM is attempted and if
       * there is a failure because less than RF/2-1 nodes are alive, it automatically switches to the Consistency Level
       * of LOCAL_ONE and so on. ( and hope that hint windows take care of the rest when the nodes come back up )
       * @param retryPolicy
       * @return The builder instance as initially created updated with this value
         */
      public ConnectionBuilder withRetryPolicy(RetryPolicy retryPolicy)
      {
        this.retryPolicy = retryPolicy;
        return this;
      }
  
      /**
       * Used to set various aspects for executing a given query / mutation.
       * The default is to use  LOCAL_QUORUM consistency for all mutation queries
       * @param queryOptions
       * @return The builder instance as initially created updated with this value
         */
      public ConnectionBuilder withQueryOptions(QueryOptions queryOptions)
      {
        this.queryOptions = queryOptions;
        return this;
      }
  
      /**
       * Used to decide how to establish a connection to the cluster in case the current one fails.
       * The default if not set is to use an ExponentialRetry Policy.
       * The baseDelay and maxDelay are the two time windows that are used to specify the retry attempts
       * in an exponential manner
       * @param reconnectionPolicy
       * @return The builder instance as initially created updated with this value
       */
      public ConnectionBuilder withReconnectionPolicy(ReconnectionPolicy reconnectionPolicy)
      {
        this.reconnectionPolicy = reconnectionPolicy;
        return this;
      }
  
      public ConnectionBuilder withProtocolVersion(ProtocolVersion protocolVersion)
      {
        this.protocolVersion = protocolVersion;
        return this;
      }
  
      protected ConnectionStateManager initialize()
      {
        ConnectionStateManager operatorConnection = new ConnectionStateManager(this);
        return operatorConnection;
      }
  
    }
  
  }
  