YCSB plocal+hash index

OrientDBClient.java

/**
 * OrientDB client binding for YCSB.
 *
 * Submitted by Luca Garulli on 5/10/2012.
 *
 */

package com.yahoo.ycsb.db;

import java.util.Collection;
import java.util.HashMap;
import java.util.Map.Entry;
import java.util.Properties;
import java.util.Set;
import java.util.Vector;

import com.orientechnologies.orient.core.config.OGlobalConfiguration;
import com.orientechnologies.orient.core.db.document.ODatabaseDocumentTx;
import com.orientechnologies.orient.core.db.record.OIdentifiable;
import com.orientechnologies.orient.core.dictionary.ODictionary;
import com.orientechnologies.orient.core.index.OIndex;
import com.orientechnologies.orient.core.intent.OIntentMassiveInsert;
import com.orientechnologies.orient.core.record.ORecordInternal;
import com.orientechnologies.orient.core.record.impl.ODocument;
import com.yahoo.ycsb.ByteIterator;
import com.yahoo.ycsb.DB;
import com.yahoo.ycsb.DBException;
import com.yahoo.ycsb.StringByteIterator;

/**
 * OrientDB client for YCSB framework.
 * <p/>
 * Properties to set:
 * <p/>
 * orientdb.url=local:C:/temp/databases or remote:localhost:2424 <br>
 * orientdb.database=ycsb <br>
 * orientdb.user=admin <br>
 * orientdb.password=admin <br>
 *
 * @author Luca Garulli
 */
public class OrientDBClient extends DB {

	private ODatabaseDocumentTx db;
	private static final String CLASS = "usertable";
	private OIndex<OIdentifiable> dictionary;

	/**
	 * Initialize any state for this DB. Called once per DB instance; there is one DB instance per client thread.
	 */
	public void init() throws DBException {
		// initialize OrientDB driver
		Properties props = getProperties();

		String url = props.getProperty("orientdb.url", "remote:localhost/ycsb");
		String user = props.getProperty("orientdb.user", "admin");
		String password = props.getProperty("orientdb.password", "admin");

		try {
			System.out.println("OrientDB loading database url = " + url);

			db = new ODatabaseDocumentTx(url);
			db.open(user, password);

			System.out.println("OrientDB connection created with " + url);

			dictionary = (OIndex<OIdentifiable>) db.getMetadata().getIndexManager().getIndex("HASH_DICTIONARY");

			db.declareIntent(new OIntentMassiveInsert());

		} catch (Exception e1) {
			System.err.println("Could not initialize OrientDB connection pool for Loader: " + e1.toString());
			e1.printStackTrace();
			return;
		}

	}

	@Override
	public void cleanup() throws DBException {
		if (db != null) {
			db.close();
			db = null;
		}
	}

	@Override
	/**
	 * Insert a record in the database. Any field/value pairs in the specified values HashMap will be written into the record with the specified
	 * record key.
	 *
	 * @param table The name of the table
	 * @param key The record key of the record to insert.
	 * @param values A HashMap of field/value pairs to insert in the record
	 * @return Zero on success, a non-zero error code on error. See this class's description for a discussion of error codes.
	 */
	public int insert(String table, String key, HashMap<String, ByteIterator> values) {
		try {
			final ODocument document = new ODocument(CLASS);
			for (Entry<String, String> entry : StringByteIterator.getStringMap(values).entrySet())
				document.field(entry.getKey(), entry.getValue());
			document.save();
			dictionary.put(key, document);

			return 0;
		} catch (Exception e) {
			e.printStackTrace();
		}
		return 1;
	}

	@Override
	/**
	 * Delete a record from the database.
	 *
	 * @param table The name of the table
	 * @param key The record key of the record to delete.
	 * @return Zero on success, a non-zero error code on error. See this class's description for a discussion of error codes.
	 */
	public int delete(String table, String key) {
		try {
			dictionary.remove(key);
			return 0;
		} catch (Exception e) {
			e.printStackTrace();
		}
		return 1;
	}

	@Override
	/**
	 * Read a record from the database. Each field/value pair from the result will be stored in a HashMap.
	 *
	 * @param table The name of the table
	 * @param key The record key of the record to read.
	 * @param fields The list of fields to read, or null for all of them
	 * @param result A HashMap of field/value pairs for the result
	 * @return Zero on success, a non-zero error code on error or "not found".
	 */
	public int read(String table, String key, Set<String> fields, HashMap<String, ByteIterator> result) {
		try {
			final OIdentifiable identifiable = dictionary.get(key);
			if (identifiable != null) {
				ODocument document = identifiable.getRecord();
				if (fields != null)
					for (String field : fields)
						result.put(field, new StringByteIterator((String) document.field(field)));
				else
					for (String field : document.fieldNames())
						result.put(field, new StringByteIterator((String) document.field(field)));
				return 0;
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
		return 1;
	}

	@Override
	/**
	 * Update a record in the database. Any field/value pairs in the specified values HashMap will be written into the record with the specified
	 * record key, overwriting any existing values with the same field name.
	 *
	 * @param table The name of the table
	 * @param key The record key of the record to write.
	 * @param values A HashMap of field/value pairs to update in the record
	 * @return Zero on success, a non-zero error code on error. See this class's description for a discussion of error codes.
	 */
	public int update(String table, String key, HashMap<String, ByteIterator> values) {
		try {
			final OIdentifiable identifiable = dictionary.get(key);
			if (identifiable != null) {
				ODocument document = identifiable.getRecord();
				for (Entry<String, String> entry : StringByteIterator.getStringMap(values).entrySet())
					document.field(entry.getKey(), entry.getValue());
				document.save();
				return 0;
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
		return 1;
	}

	@Override
	/**
	 * Perform a range scan for a set of records in the database. Each field/value pair from the result will be stored in a HashMap.
	 *
	 * @param table The name of the table
	 * @param startkey The record key of the first record to read.
	 * @param recordcount The number of records to read
	 * @param fields The list of fields to read, or null for all of them
	 * @param result A Vector of HashMaps, where each HashMap is a set field/value pairs for one record
	 * @return Zero on success, a non-zero error code on error. See this class's description for a discussion of error codes.
	 */
	public int scan(String table, String startkey, int recordcount, Set<String> fields, Vector<HashMap<String, ByteIterator>> result) {
		throw new UnsupportedOperationException();
	}
}