D-Net

package eu.dnetlib.data.hadoop;

import java.io.IOException;

import java.net.URI;

import java.util.*;

import java.util.Map.Entry;

import java.util.stream.Collectors;

import com.google.common.collect.Lists;

import com.google.common.collect.Maps;

import com.google.common.collect.Sets;

import eu.dnetlib.data.hadoop.config.ClusterName;

import eu.dnetlib.data.hadoop.config.ConfigurationEnumerator;

import eu.dnetlib.data.hadoop.rmi.HadoopServiceException;

import eu.dnetlib.data.hadoop.rmi.hbase.Column;

import eu.dnetlib.data.hadoop.rmi.hbase.HBaseRowDescriptor;

import eu.dnetlib.data.hadoop.rmi.hbase.schema.HBaseTableDescriptor;

import eu.dnetlib.data.hadoop.rmi.hbase.schema.HBaseTableRegionInfo;

import org.apache.commons.lang.StringUtils;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.apache.hadoop.conf.Configuration;

import org.apache.hadoop.fs.FileSystem;

import org.apache.hadoop.fs.Path;

import org.apache.hadoop.hbase.HColumnDescriptor;

import org.apache.hadoop.hbase.HRegionInfo;

import org.apache.hadoop.hbase.HTableDescriptor;

import org.apache.hadoop.hbase.client.*;

import org.apache.hadoop.hbase.util.Bytes;

import org.springframework.beans.factory.annotation.Autowired;

import org.springframework.beans.factory.annotation.Required;

public class HadoopServiceCore {

	private static final Log log = LogFactory.getLog(HadoopServiceCore.class); // NOPMD by marko on 11/24/08 5:02 PM

	@Autowired

	protected ConfigurationEnumerator configurationEnumerator;

	@Autowired

	private HadoopClientMap clients;

	private int maxVersions;

	public List<String> listTables(final ClusterName clusterName) throws IOException, HadoopServiceException {

		try(final HBaseAdmin admin = clients.getHbaseAdmin(clusterName)) {

			return Arrays.asList(admin.listTables())

					.stream()

					.map(HTableDescriptor::getNameAsString)

					.collect(Collectors.toList());

		}

	}

	public String getHBaseTableDescriptor(final ClusterName clusterName, final String tableName) throws HadoopServiceException, IOException {

		try(final HBaseAdmin admin = clients.getHbaseAdmin(clusterName)) {

			if (StringUtils.isBlank(tableName)) throw new HadoopServiceException("Table name cannot be empty or null");

			if (admin == null) throw new HadoopServiceException(String.format("HBase admin not available for cluster: '%s'", clusterName.toString()));

			final List<HRegionInfo> tableRegions = admin.getTableRegions(tableName.getBytes());

			final HTableDescriptor desc = admin.getTableDescriptor(tableName.getBytes());

			final Set<String> columns = Sets.newHashSet();

			for (HColumnDescriptor hColDesc : Arrays.asList(desc.getColumnFamilies())) {

				columns.add(hColDesc.getNameAsString());

			}

			HBaseTableDescriptor htDescriptor = new HBaseTableDescriptor();

			htDescriptor.setColumns(columns);

			List<HBaseTableRegionInfo> regions = Lists.newArrayList();

			for (HRegionInfo info : tableRegions) {

				regions.add(new HBaseTableRegionInfo(new String(info.getStartKey()), new String(info.getEndKey())));

			}

			htDescriptor.setRegions(regions);

			if (log.isDebugEnabled()) {

				log.info("got configuration for table '" + tableName + "': " + htDescriptor.toString());

			}

			return htDescriptor.toString();

		}

	}

	public List<String> describeTable(final ClusterName clusterName, final String table) throws IOException, HadoopServiceException {

		try(final HBaseAdmin admin = clients.getHbaseAdmin(clusterName)) {

			final HTableDescriptor desc = admin.getTableDescriptor(table.getBytes());

			return desc.getFamilies().stream()

					.map(d -> d.getNameAsString())

					.collect(Collectors.toList());

		}

	}

	public void truncateTable(final ClusterName clusterName, final String table) throws IOException, HadoopServiceException {

		try(final HBaseAdmin admin = clients.getHbaseAdmin(clusterName)) {

			if (!admin.tableExists(table)) throw new IllegalStateException("cannot truncate unexisting table");

			final HTableDescriptor desc = admin.getTableDescriptor(table.getBytes());

			log.info("disabling table: " + table);

			admin.disableTable(table);

			log.info("deleting table: " + table);

			admin.deleteTable(table);

			log.info("creating table: " + table);

			admin.createTable(desc);

		}

	}

	public boolean existTable(final ClusterName clusterName, final String table) throws IOException, HadoopServiceException {

		try(final HBaseAdmin admin = clients.getHbaseAdmin(clusterName)) {

			return admin.tableExists(table);

		}

	}

	public void dropTable(final ClusterName clusterName, final String table) throws IOException, HadoopServiceException {

		try(final HBaseAdmin admin = clients.getHbaseAdmin(clusterName)) {

			if (!admin.tableExists(table)) throw new IllegalStateException("cannot drop unexisting table: '" + table + "'");

			log.info("disabling table: " + table);

			admin.disableTable(table);

			log.info("deleting table: " + table);

			admin.deleteTable(table);

		}

	}

	public void createTable(final ClusterName clusterName, final String table, final String tableConfiguration) throws IOException, HadoopServiceException {

		try(final HBaseAdmin admin = clients.getHbaseAdmin(clusterName)) {

			if (admin.tableExists(table)) throw new IllegalStateException("table already exists");

			if (StringUtils.isBlank(tableConfiguration)) throw new HadoopServiceException("empty table configuration");

			final HBaseTableDescriptor tableConf = HBaseTableDescriptor.fromJSON(tableConfiguration);

			doCreateTable(clusterName, table, tableConf.getColumns(), tableConf.getRegions());

		}

	}

	public void createTable(final ClusterName clusterName, final String table, final Set<String> columns) throws IOException, HadoopServiceException {

		try(final HBaseAdmin admin = clients.getHbaseAdmin(clusterName)) {

			if (admin.tableExists(table)) throw new IllegalStateException("table already exists");

			doCreateTable(clusterName, table, columns, null);

		}

	}

	public void doCreateTable(final ClusterName clusterName, final String table, final Set<String> columns, final List<HBaseTableRegionInfo> regions)

			throws IOException, HadoopServiceException {

		try(final HBaseAdmin admin = clients.getHbaseAdmin(clusterName)) {

			if (admin.tableExists(table)) throw new IllegalStateException("table already exists");

			final HTableDescriptor desc = new HTableDescriptor(table);

			for (final String column : columns) {

				final HColumnDescriptor hds = new HColumnDescriptor(column);

				hds.setMaxVersions(getMaxVersions());

				desc.addFamily(hds);

			}

			log.info("creating hbase table: " + table);

			if (regions != null && !regions.isEmpty()) {

				log.debug(String.format("create using %s regions: %s", regions.size(), regions));

				admin.createTable(desc, getSplitKeys(regions));

			} else {

				admin.createTable(desc);

			}

			log.info("created hbase table: [" + table + "]");

			log.debug("descriptor: [" + desc.toString() + "]");

		}

	}

	private byte[][] getSplitKeys(final List<HBaseTableRegionInfo> regions) {

		byte[][] splits = new byte[regions.size() - 1][];

		for (int i = 0; i < regions.size() - 1; i++) {

			splits[i] = regions.get(i).getEndKey().getBytes();

		}

		return splits;

	}

	public void ensureTable(final ClusterName clusterName, final String table, final Set<String> columns) throws IOException, HadoopServiceException {

		try(final HBaseAdmin admin = clients.getHbaseAdmin(clusterName)) {

			if (!admin.tableExists(table)) {

				createTable(clusterName, table, columns);

			} else {

				final HTableDescriptor desc = admin.getTableDescriptor(Bytes.toBytes(table));

				final Set<String> foundColumns = desc.getFamilies().stream()

						.map(d -> d.getNameAsString())

						.collect(Collectors.toCollection(HashSet::new));

				log.info("ensuring columns on table " + table + ": " + columns);

				final Collection<String> missingColumns = Sets.difference(columns, foundColumns);

				if (!missingColumns.isEmpty()) {

					if (admin.isTableEnabled(table)) {

						admin.disableTable(table);

					}

					for (final String column : missingColumns) {

						log.info("hbase table: '" + table + "', adding column: " + column);

						admin.addColumn(table, new HColumnDescriptor(column));

					}

					admin.enableTable(table);

				}

			}

		}

	}

	public void writeOnHBase(final ClusterName clusterName, final String tableName, final List<Put> puts) throws IOException {

		final Configuration conf = configurationEnumerator.get(clusterName);

		final HTable table = new HTable(conf, tableName);

		try {

			table.put(puts);

		} finally {

			table.flushCommits();

			table.close();

		}

	}

	public void deleteFromHBase(final ClusterName clusterName, final String tableName, final List<Delete> deletes) throws IOException {

		final Configuration conf = configurationEnumerator.get(clusterName);

		final HTable table = new HTable(conf, tableName);

		try {

			table.delete(deletes);

		} finally {

			table.flushCommits();

			table.close();

		}

	}

	public void deleteColumnsFromHBase(final ClusterName clusterName, final String tableName, final List<HBaseRowDescriptor> columns) throws IOException {

		final Configuration conf = configurationEnumerator.get(clusterName);

		final HTable table = new HTable(conf, tableName);

		try {

			for(HBaseRowDescriptor desc : columns) {

				final Delete d = new Delete(Bytes.toBytes(desc.getRowKey()));

				d.setWriteToWAL(true);

				for(Column c : desc.getColumns()) {

					for(String qualifier : c.getQualifier()) {

						log.info(String.format("delete from row '%s' cf '%s:%s'", desc.getRowKey(), c.getFamily(), qualifier));

						d.deleteColumns(Bytes.toBytes(c.getFamily()), Bytes.toBytes(qualifier));

					}

				}

				table.delete(d);

			}

		} finally {

			table.flushCommits();

			table.close();

		}

	}

	public Result getRow(final ClusterName clusterName, final String tableName, final byte[] id) throws IOException {

		final Configuration conf = configurationEnumerator.get(clusterName);

		final HTable table = new HTable(conf, tableName);

		try {

			return table.get(new Get(id));

		} finally {

			table.close();

		}

	}

	public Map<String, HBaseRowDescriptor> describeRows(final ClusterName clusterName, final String tableName, final List<String> rowKeys) throws IOException {

		final Map<String, HBaseRowDescriptor> map = Maps.newHashMap();

		for(String rowKey : rowKeys) {

			map.put(rowKey, describeRow(clusterName, tableName, rowKey));

		}

		return map;

	}

	public HBaseRowDescriptor describeRow(final ClusterName clusterName, final String tableName, final String rowKey) throws IOException {

		final Configuration conf = configurationEnumerator.get(clusterName);

		final HTable table = new HTable(conf, tableName);

		final HBaseRowDescriptor desc = new HBaseRowDescriptor();

		try {

			final Result r = table.get(new Get(Bytes.toBytes(rowKey)));

			if (r.isEmpty()) {

				return desc;

			}

			final List<Column> columns = Lists.newArrayList();

			for(Entry<byte[], NavigableMap<byte[], byte[]>> e : r.getNoVersionMap().entrySet()) {

				final Set<byte[]> qualifiers = e.getValue().keySet();

				final String family = new String(e.getKey());

				final Column col = new Column(family);

				for(byte[] q : qualifiers) {

					String qs = new String(q);

					col.getQualifier().add(qs);

				}

				columns.add(col);

			}

			desc.setColumns(columns);

			desc.setRowKey(rowKey);

			return desc;

		} finally {

			table.close();

		}

	}

	public List<Result> getRows(final ClusterName clusterName, final String tableName, final Scan scan) throws IOException {

		final Configuration conf = configurationEnumerator.get(clusterName);

		try(final HTable table = new HTable(conf, tableName)) {

			final ResultScanner rs = table.getScanner(scan);

			try {

				return Lists.newArrayList(rs.iterator());

			} finally {

				rs.close();

			}

		}

	}

	public boolean deleteFromHdfs(final ClusterName clusterName, final String path) throws HadoopServiceException {

		if (StringUtils.isBlank(path))

			throw new HadoopServiceException("Cannot deleteFromHBase an empty HDFS path.");

		final Configuration conf = configurationEnumerator.get(clusterName);

		try(final FileSystem hdfs = FileSystem.get(conf)) {

			final Path absolutePath = new Path(URI.create(conf.get("fs.defaultFS") + path));

			if (hdfs.exists(absolutePath)) {

				log.debug("deleteFromHBase path: " + absolutePath.toString());

				hdfs.delete(absolutePath, true);

				log.info("deleted path: " + absolutePath.toString());

				return true;

			} else {

				log.warn("cannot deleteFromHBase unexisting path: " + absolutePath.toString());

				return false;

			}

		} catch (IOException e) {

			throw new HadoopServiceException(e);

		}

	}

	public boolean createHdfsDir(final ClusterName clusterName, final String path, final boolean force) throws HadoopServiceException {

		if (StringUtils.isBlank(path))

			throw new HadoopServiceException("Cannot create an empty HDFS path.");

		final Configuration conf = configurationEnumerator.get(clusterName);

		try(final FileSystem hdfs = FileSystem.get(conf)) {

			final Path absolutePath = new Path(URI.create(conf.get("fs.defaultFS") + path));

			if (!hdfs.exists(absolutePath)) {

				hdfs.mkdirs(absolutePath);

				log.info("created path: " + absolutePath.toString());

				return true;

			} else if (force) {

				log.info(String.format("found directory '%s', force delete it", absolutePath.toString()));

				hdfs.delete(absolutePath, true);

				hdfs.mkdirs(absolutePath);

				log.info("created path: " + absolutePath.toString());

				return true;

			} else {

				log.info(String.format("directory already exists: '%s', nothing to do", absolutePath.toString()));

				return false;

			}

		} catch (IOException e) {

			throw new HadoopServiceException(e);

		}

	}

	public boolean existHdfsPath(final ClusterName clusterName, final String path) throws HadoopServiceException {

		if (StringUtils.isBlank(path))

			throw new HadoopServiceException("invalid empty path");

		final Configuration conf = configurationEnumerator.get(clusterName);

		try(final FileSystem hdfs = FileSystem.get(conf)) {

			final Path absolutePath = new Path(URI.create(conf.get("fs.defaultFS") + path));

			return hdfs.exists(absolutePath);

		} catch (IOException e) {

			throw new HadoopServiceException(e);

		}

	}

	public Configuration getClusterConiguration(final ClusterName clusterName) {

		return configurationEnumerator.get(clusterName);

	}

	public int getMaxVersions() {

		return maxVersions;

	}

	@Required

	public void setMaxVersions(final int maxVersions) {

		this.maxVersions = maxVersions;

	}

	public HadoopClientMap getClients() {

		return clients;

	}

}