com.klarna.hiverunner.data.TableDataInserter Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of hiverunner Show documentation

HiveRunner is a unit test framework based on JUnit (4 or 5) that enables TDD development of Hive SQL without the need of any installed dependencies.

There is a newer version: 6.0.1

Show newest version

package com.klarna.hiverunner.data;

import java.util.Iterator;
import java.util.Map;

import org.apache.hadoop.hive.conf.HiveConf;
import org.apache.hive.hcatalog.common.HCatException;
import org.apache.hive.hcatalog.data.HCatRecord;
import org.apache.hive.hcatalog.data.transfer.DataTransferFactory;
import org.apache.hive.hcatalog.data.transfer.HCatWriter;
import org.apache.hive.hcatalog.data.transfer.WriteEntity;
import org.apache.hive.hcatalog.data.transfer.WriterContext;

import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;

class TableDataInserter {

  private final String databaseName;
  private final String tableName;
  private final Map config;

  TableDataInserter(String databaseName, String tableName, HiveConf conf) {
    this.databaseName = databaseName;
    this.tableName = tableName;
    config = Maps.fromProperties(conf.getAllProperties());
  }

  void insert(Multimap, HCatRecord> data) {
    Iterator> iterator = data.keySet().iterator();
    while (iterator.hasNext()) {
      Map partitionSpec = iterator.next();
      insert(partitionSpec, data.get(partitionSpec));
    }
  }

  private void insert(Map partitionSpec, Iterable rows) {
    WriteEntity entity = new WriteEntity.Builder()
        .withDatabase(databaseName)
        .withTable(tableName)
        .withPartition(partitionSpec)
        .build();

    try {
      HCatWriter master = DataTransferFactory.getHCatWriter(entity, config);
      WriterContext context = master.prepareWrite();
      HCatWriter writer = DataTransferFactory.getHCatWriter(context);
      writer.write(rows.iterator());
      master.commit(context);
    } catch (HCatException e) {
      throw new RuntimeException("An error occurred while inserting data to " + databaseName + "." + tableName, e);
    }
  }

}