demos.BatchInsert.java Source code

Introduction

Here is the source code for demos.BatchInsert.java

Source

/*
 * Copyright 2015 Red Hat, Inc. and/or its affiliates
 * and other contributors as indicated by the @author tags.
 *
 * Licensed 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 demos;

import java.util.Random;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

import com.datastax.driver.core.BatchStatement;
import com.datastax.driver.core.Cluster;
import com.datastax.driver.core.PreparedStatement;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.ResultSetFuture;
import com.datastax.driver.core.Session;
import com.google.common.base.Stopwatch;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;

import org.joda.time.DateTime;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Demonstrates inserting data using batches asynchronously. You must first
 * run <code>src/main/resources/demo.cql</code> prior to running this example.
 *
 * @author jsanda
 */
public class BatchInsert implements Runnable {

    private static final Logger logger = LoggerFactory.getLogger(BatchInsert.class);

    private final int NUM_INSERTS = 100000;

    private final int BATCH_SIZE = 100;

    private final int NUM_METRICS = 25;

    public void run() {
        try {
            logger.info("Preparing to insert metric data points");

            Cluster cluster = Cluster.builder().addContactPoint("127.0.0.1").build();
            Session session = cluster.connect("demo");
            PreparedStatement insert = session
                    .prepare("insert into metric_data (metric_id, time, value) values (?, ?, ?)");
            Random random = new Random();
            DateTime time = DateTime.now().minusYears(1);
            final CountDownLatch latch = new CountDownLatch(NUM_INSERTS / BATCH_SIZE);

            FutureCallback<ResultSet> callback = new FutureCallback<ResultSet>() {
                @Override
                public void onSuccess(ResultSet result) {
                    latch.countDown();
                }

                @Override
                public void onFailure(Throwable t) {
                    logger.warn("There was an error inserting data", t);
                    latch.countDown();
                }
            };

            Stopwatch stopwatch = new Stopwatch().start();
            BatchStatement batch = new BatchStatement();
            for (int i = 0; i < NUM_INSERTS; ++i) {
                String metricId = "metric-" + Math.abs(random.nextInt() % NUM_METRICS);
                double value = random.nextDouble();
                batch.add(insert.bind(metricId, time.toDate(), value));
                time = time.plusSeconds(10);
                if (batch.size() == BATCH_SIZE) {
                    ResultSetFuture future = session.executeAsync(batch);
                    Futures.addCallback(future, callback);
                    batch = new BatchStatement();
                }
            }
            latch.await();
            stopwatch.stop();

            logger.info("Finished inserting {} data points in {} ms", NUM_INSERTS,
                    stopwatch.elapsed(TimeUnit.MILLISECONDS));
        } catch (InterruptedException e) {
            logger.info("There was an interrupt while waiting for inserts to complete");
        }
    }
}