Java tutorial
/** * Copyright (c) 2012 - 2015 YCSB contributors. All rights reserved. * * 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. See accompanying * LICENSE file. */ /* * MongoDB client binding for YCSB. * * Submitted by Yen Pai on 5/11/2010. * * https://gist.github.com/000a66b8db2caf42467b#file_mongo_database.java */ package com.yahoo.ycsb.db3; import com.mongodb.MongoClient; import com.mongodb.MongoClientURI; import com.mongodb.ReadPreference; import com.mongodb.WriteConcern; import com.mongodb.client.FindIterable; import com.mongodb.client.MongoCollection; import com.mongodb.client.MongoCursor; import com.mongodb.client.MongoDatabase; import com.mongodb.client.model.InsertManyOptions; import com.mongodb.client.model.UpdateOneModel; import com.mongodb.client.model.UpdateOptions; import com.mongodb.client.result.DeleteResult; import com.mongodb.client.result.UpdateResult; import com.yahoo.ycsb.ByteArrayByteIterator; import com.yahoo.ycsb.ByteIterator; import com.yahoo.ycsb.DB; import com.yahoo.ycsb.DBException; import com.yahoo.ycsb.Status; import org.bson.Document; import org.bson.types.Binary; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Properties; import java.util.Set; import java.util.Vector; import java.util.concurrent.atomic.AtomicInteger; /** * MongoDB binding for YCSB framework using the MongoDB Inc. <a * href="http://docs.mongodb.org/ecosystem/drivers/java/">driver</a> * <p> * See the <code>README.md</code> for configuration information. * </p> * * @author ypai * @see <a href="http://docs.mongodb.org/ecosystem/drivers/java/">MongoDB Inc. * driver</a> */ public class MongoDbClient extends DB { /** Used to include a field in a response. */ private static final Integer INCLUDE = Integer.valueOf(1); /** The options to use for inserting many documents. */ private static final InsertManyOptions INSERT_UNORDERED = new InsertManyOptions().ordered(false); /** The options to use for inserting a single document. */ private static final UpdateOptions UPDATE_WITH_UPSERT = new UpdateOptions().upsert(true); /** * The database name to access. */ private static String databaseName; /** The database name to access. */ private static MongoDatabase database; /** * Count the number of times initialized to teardown on the last * {@link #cleanup()}. */ private static final AtomicInteger INIT_COUNT = new AtomicInteger(0); /** A singleton Mongo instance. */ private static MongoClient mongoClient; /** The default read preference for the test. */ private static ReadPreference readPreference; /** The default write concern for the test. */ private static WriteConcern writeConcern; /** The batch size to use for inserts. */ private static int batchSize; /** If true then use updates with the upsert option for inserts. */ private static boolean useUpsert; /** The bulk inserts pending for the thread. */ private final List<Document> bulkInserts = new ArrayList<Document>(); /** * Cleanup any state for this DB. Called once per DB instance; there is one DB * instance per client thread. */ @Override public void cleanup() throws DBException { if (INIT_COUNT.decrementAndGet() == 0) { try { mongoClient.close(); } catch (Exception e1) { System.err.println("Could not close MongoDB connection pool: " + e1.toString()); e1.printStackTrace(); return; } finally { database = null; mongoClient = null; } } } /** * 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 the {@link DB} * class's description for a discussion of error codes. */ @Override public Status delete(String table, String key) { try { MongoCollection<Document> collection = database.getCollection(table); Document query = new Document("_id", key); DeleteResult result = collection.withWriteConcern(writeConcern).deleteOne(query); if (result.wasAcknowledged() && result.getDeletedCount() == 0) { System.err.println("Nothing deleted for key " + key); return Status.NOT_FOUND; } return Status.OK; } catch (Exception e) { System.err.println(e.toString()); return Status.ERROR; } } /** * Initialize any state for this DB. Called once per DB instance; there is one * DB instance per client thread. */ @Override public void init() throws DBException { INIT_COUNT.incrementAndGet(); synchronized (INCLUDE) { if (mongoClient != null) { return; } Properties props = getProperties(); // Set insert batchsize, default 1 - to be YCSB-original equivalent batchSize = Integer.parseInt(props.getProperty("batchsize", "1")); // Set is inserts are done as upserts. Defaults to false. useUpsert = Boolean.parseBoolean(props.getProperty("mongodb.upsert", "false")); // Just use the standard connection format URL // http://docs.mongodb.org/manual/reference/connection-string/ // to configure the client. String url = props.getProperty("mongodb.url", null); boolean defaultedUrl = false; if (url == null) { defaultedUrl = true; url = "mongodb://localhost:27017/ycsb?w=1"; } url = OptionsSupport.updateUrl(url, props); if (!url.startsWith("mongodb://")) { System.err.println("ERROR: Invalid URL: '" + url + "'. Must be of the form " + "'mongodb://<host1>:<port1>,<host2>:<port2>/database?options'. " + "http://docs.mongodb.org/manual/reference/connection-string/"); System.exit(1); } try { MongoClientURI uri = new MongoClientURI(url); String uriDb = uri.getDatabase(); if (!defaultedUrl && (uriDb != null) && !uriDb.isEmpty() && !"admin".equals(uriDb)) { databaseName = uriDb; } else { // If no database is specified in URI, use "ycsb" databaseName = "ycsb"; } readPreference = uri.getOptions().getReadPreference(); writeConcern = uri.getOptions().getWriteConcern(); mongoClient = new MongoClient(uri); database = mongoClient.getDatabase(databaseName).withReadPreference(readPreference) .withWriteConcern(writeConcern); System.out.println("mongo client connection created with " + url); } catch (Exception e1) { System.err.println("Could not initialize MongoDB connection pool for Loader: " + e1.toString()); e1.printStackTrace(); return; } } } /** * 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 the {@link DB} * class's description for a discussion of error codes. */ @Override public Status insert(String table, String key, HashMap<String, ByteIterator> values) { try { MongoCollection<Document> collection = database.getCollection(table); Document toInsert = new Document("_id", key); for (Map.Entry<String, ByteIterator> entry : values.entrySet()) { toInsert.put(entry.getKey(), entry.getValue().toArray()); } if (batchSize == 1) { if (useUpsert) { // this is effectively an insert, but using an upsert instead due // to current inability of the framework to clean up after itself // between test runs. collection.replaceOne(new Document("_id", toInsert.get("_id")), toInsert, UPDATE_WITH_UPSERT); } else { collection.insertOne(toInsert); } } else { bulkInserts.add(toInsert); if (bulkInserts.size() == batchSize) { if (useUpsert) { List<UpdateOneModel<Document>> updates = new ArrayList<UpdateOneModel<Document>>( bulkInserts.size()); for (Document doc : bulkInserts) { updates.add(new UpdateOneModel<Document>(new Document("_id", doc.get("_id")), doc, UPDATE_WITH_UPSERT)); } collection.bulkWrite(updates); } else { collection.insertMany(bulkInserts, INSERT_UNORDERED); } bulkInserts.clear(); } else { return OptionsSupport.BATCHED_OK; } } return Status.OK; } catch (Exception e) { System.err.println("Exception while trying bulk insert with " + bulkInserts.size()); e.printStackTrace(); return Status.ERROR; } } /** * 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". */ @Override public Status read(String table, String key, Set<String> fields, HashMap<String, ByteIterator> result) { try { MongoCollection<Document> collection = database.getCollection(table); Document query = new Document("_id", key); FindIterable<Document> findIterable = collection.find(query); if (fields != null) { Document projection = new Document(); for (String field : fields) { projection.put(field, INCLUDE); } findIterable.projection(projection); } Document queryResult = findIterable.first(); if (queryResult != null) { fillMap(result, queryResult); } return queryResult != null ? Status.OK : Status.NOT_FOUND; } catch (Exception e) { System.err.println(e.toString()); return Status.ERROR; } } /** * 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 the {@link DB} * class's description for a discussion of error codes. */ @Override public Status scan(String table, String startkey, int recordcount, Set<String> fields, Vector<HashMap<String, ByteIterator>> result) { MongoCursor<Document> cursor = null; try { MongoCollection<Document> collection = database.getCollection(table); Document scanRange = new Document("$gte", startkey); Document query = new Document("_id", scanRange); Document sort = new Document("_id", INCLUDE); FindIterable<Document> findIterable = collection.find(query).sort(sort).limit(recordcount); if (fields != null) { Document projection = new Document(); for (String fieldName : fields) { projection.put(fieldName, INCLUDE); } findIterable.projection(projection); } cursor = findIterable.iterator(); if (!cursor.hasNext()) { System.err.println("Nothing found in scan for key " + startkey); return Status.ERROR; } result.ensureCapacity(recordcount); while (cursor.hasNext()) { HashMap<String, ByteIterator> resultMap = new HashMap<String, ByteIterator>(); Document obj = cursor.next(); fillMap(resultMap, obj); result.add(resultMap); } return Status.OK; } catch (Exception e) { System.err.println(e.toString()); return Status.ERROR; } finally { if (cursor != null) { cursor.close(); } } } /** * 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. */ @Override public Status update(String table, String key, HashMap<String, ByteIterator> values) { try { MongoCollection<Document> collection = database.getCollection(table); Document query = new Document("_id", key); Document fieldsToSet = new Document(); for (Map.Entry<String, ByteIterator> entry : values.entrySet()) { fieldsToSet.put(entry.getKey(), entry.getValue().toArray()); } Document update = new Document("$set", fieldsToSet); UpdateResult result = collection.updateOne(query, update); if (result.wasAcknowledged() && result.getMatchedCount() == 0) { System.err.println("Nothing updated for key " + key); return Status.NOT_FOUND; } return Status.OK; } catch (Exception e) { System.err.println(e.toString()); return Status.ERROR; } } /** * Fills the map with the values from the DBObject. * * @param resultMap * The map to fill/ * @param obj * The object to copy values from. */ protected void fillMap(Map<String, ByteIterator> resultMap, Document obj) { for (Map.Entry<String, Object> entry : obj.entrySet()) { if (entry.getValue() instanceof Binary) { resultMap.put(entry.getKey(), new ByteArrayByteIterator(((Binary) entry.getValue()).getData())); } } } }