org.apache.streams.cassandra.CassandraClient.java Source code

Introduction

Here is the source code for org.apache.streams.cassandra.CassandraClient.java

Source

/*
 * 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
 *
 *   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.streams.cassandra;

import com.datastax.driver.core.Cluster;
import com.datastax.driver.core.Host;
import com.datastax.driver.core.JdkSSLOptions;
import com.datastax.driver.core.Metadata;
import com.datastax.driver.core.SSLOptions;
import com.datastax.driver.core.Session;
import com.datastax.driver.core.SocketOptions;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.FileInputStream;
import java.io.InputStream;
import java.net.InetSocketAddress;
import java.security.KeyStore;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Objects;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManagerFactory;

import static com.datastax.driver.core.SocketOptions.DEFAULT_CONNECT_TIMEOUT_MILLIS;
import static com.datastax.driver.core.SocketOptions.DEFAULT_READ_TIMEOUT_MILLIS;

public class CassandraClient {

    private static final Logger LOGGER = LoggerFactory.getLogger(CassandraClient.class);

    private Cluster cluster;
    private Session session;

    public CassandraConfiguration config;

    public CassandraClient(CassandraConfiguration config) throws Exception {
        this.config = config;
        org.apache.cassandra.config.Config.setClientMode(true);
    }

    public void start() throws Exception {

        Objects.nonNull(config);

        LOGGER.info("CassandraClient.start {}", config);

        Cluster.Builder builder = Cluster.builder().withPort(config.getPort().intValue()).withoutJMXReporting()
                .withoutMetrics()
                .withSocketOptions(new SocketOptions().setConnectTimeoutMillis(DEFAULT_CONNECT_TIMEOUT_MILLIS * 10)
                        .setReadTimeoutMillis(DEFAULT_READ_TIMEOUT_MILLIS * 10));

        if (config.getSsl() != null && config.getSsl().getEnabled() == true) {

            Ssl ssl = config.getSsl();

            KeyStore ks = KeyStore.getInstance("JKS");

            InputStream trustStore = new FileInputStream(ssl.getTrustStore());
            ks.load(trustStore, ssl.getTrustStorePassword().toCharArray());
            InputStream keyStore = new FileInputStream(ssl.getKeyStore());
            ks.load(keyStore, ssl.getKeyStorePassword().toCharArray());

            TrustManagerFactory tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
            tmf.init(ks);

            KeyManagerFactory kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
            kmf.init(ks, ssl.getKeyStorePassword().toCharArray());

            SSLContext sslContext = SSLContext.getInstance("SSLv3");
            sslContext.init(kmf.getKeyManagers(), tmf.getTrustManagers(), null);

            SSLOptions sslOptions = JdkSSLOptions.builder().withSSLContext(sslContext).build();

            builder = builder.withSSL(sslOptions);
        }

        Collection<InetSocketAddress> addresses = new ArrayList<>();
        for (String h : config.getHosts()) {
            LOGGER.info("Adding Host: {}", h);
            InetSocketAddress socketAddress = new InetSocketAddress(h, config.getPort().intValue());
            addresses.add(socketAddress);
        }
        builder.addContactPointsWithPorts(addresses);

        if (StringUtils.isNotBlank(config.getUser()) && StringUtils.isNotBlank(config.getPassword())) {
            builder.withCredentials(config.getUser(), config.getPassword());
        }
        cluster = builder.build();

        Objects.nonNull(cluster);

        try {
            Metadata metadata = cluster.getMetadata();
            LOGGER.info("Connected to cluster: {}\n", metadata.getClusterName());
            for (Host host : metadata.getAllHosts()) {
                LOGGER.info("Datacenter: {}; Host: {}; Rack: {}\n", host.getDatacenter(), host.getAddress(),
                        host.getRack());
            }
        } catch (Exception e) {
            LOGGER.error("Exception: {}", e);
            throw e;
        }

        try {
            session = cluster.connect();
        } catch (Exception e) {
            LOGGER.error("Exception: {}", e);
            throw e;
        }

        Objects.nonNull(session);

    }

    public void stop() throws Exception {
        session.close();
        cluster.close();
    }

    public CassandraConfiguration config() {
        return config;
    }

    public Session client() {
        return session;
    }

    public Cluster cluster() {
        return cluster;
    }
}