com.datastax.data.prepare.spark.dataset.database.CassandraHandler Maven / Gradle / Ivy

Go to download
package com.datastax.data.prepare.spark.dataset.database;

import com.datastax.insight.core.driver.SparkContextBuilder;
import com.datastax.insight.spec.Operator;
import com.datastax.insight.annonation.InsightComponent;
import com.datastax.insight.annonation.InsightComponentArg;
import com.datastax.driver.core.Cluster;
import com.datastax.driver.core.Session;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SaveMode;

import java.util.HashMap;

public class CassandraHandler implements Operator{

    @InsightComponent( name = "loadCassandra", description = "loadCassandra")
    public static Dataset load(
            @InsightComponentArg(name = "host", description = "host", request = true) String host,
            @InsightComponentArg(name = "keySpace", description = "keySpace", request = true) String keySpace,
            @InsightComponentArg(name = "table", description = "table", request = true) String table) {
        Dataset ds = SparkContextBuilder.getSession().read()
                .format("org.apache.spark.sql.cassandra")
                .options(new HashMap(){
                    {
                        put("keyspace", keySpace);
                        put("table", table);
                        put("spark.cassandra.connection.host",host);
                    }
                }).load();
        return ds;
    }

    @InsightComponent( name = "saveCassandra", description = "saveCassandra")
    public static void save(
            @InsightComponentArg(externalInput = true, name = "数据集", description = "数据集",request = true,defaultValue = "${output}") Dataset dataset,
            @InsightComponentArg(name = "host", description = "host", request = true) String host,
            @InsightComponentArg(name = "keySpace", description = "keySpace", request = true) String keySpace,
            @InsightComponentArg(name = "table", description = "table", request = true) String table,
            @InsightComponentArg(name = "saveMode", description = "saveMode", request = true, defaultValue = "ignore",items = "append;overwrite;errorIfExists;ignore") String saveMode) {
        dataset.write()
                .format("org.apache.spark.sql.cassandra")
                .options(new HashMap() {
                    {
                        put("keyspace", keySpace);
                        put("table", table);
                        put("spark.cassandra.connection.host",host);

                    }
                })
                .mode(getSaveMode(saveMode))
                .save();
    }

    @InsightComponent( name = "excuteCassandra", description = "excuteCassandra")
    public static com.datastax.driver.core.ResultSet excute(
            @InsightComponentArg(name = "keySpace", description = "keySpace", request = true) String keySpace,
            @InsightComponentArg(name = "host", description = "host", request = true) String host,
            @InsightComponentArg(name = "query", description = "query", request = true) String query) {
        Cluster cluster = Cluster.builder().addContactPoint(host).build();
        Session session = cluster.connect(keySpace);
        return session.execute(query);
    }

    private static SaveMode getSaveMode(String saveMode) {
        switch (saveMode) {
            case "append":
                return SaveMode.Append;
            case "overwrite":
                return SaveMode.Overwrite;
            case "errorIfExists":
                return SaveMode.ErrorIfExists;
            case "ignore":
                return SaveMode.Ignore;
            default:
                return SaveMode.Ignore;
        }
    }
}