Instantly share code, notes, and snippets.

Embed
What would you like to do?
Simple change to add auth to YCSB MongoDB (by ScaleGrid, formerly MongoDirector)
// Visit: http://scalegrid.io (Formerly http://mongodirector.com)
// With Reference:
// https://github.com/brianfrankcooper/YCSB
// commit: 5659fc582c8280e1431ebcfa0891979f806c70ed pulled 1 April 2015.
// The current MongoDB driver code @ mongodb/src/main/java/com/yahoo/ycsb/db/MongoDbClient.java doesn't handle MongoDB auth.
// It is also really old. I didn't not want to go and change the entire set of deprecated calls but wanted auth to work for me.
// So I added a new property called mongodb.auth (boolean) and wrote basically 10 lines of code to enable it.
// Putting this out here so that it might help someone else. It's a simple diff.
// Build: mvn clean install
// Remember to comment out/delete the mapkeeper.version entries from the pom.xml to get the mvn build to complete.
// Usage: ./bin/ycsb load mongodb -P workloads/workloada -p mongodb.url=mongodb://user:password@server1,server2,.../database -p mongodb.auth="true"
index 3082e4c..fd0ca83 100644
--- a/mongodb/src/main/java/com/yahoo/ycsb/db/MongoDbClient.java
+++ b/mongodb/src/main/java/com/yahoo/ycsb/db/MongoDbClient.java
@@ -23,6 +23,7 @@ import com.mongodb.DBCollection;
import com.mongodb.DBCursor;
import com.mongodb.DBObject;
import com.mongodb.Mongo;
+import com.mongodb.MongoURI;
import com.mongodb.MongoOptions;
import com.mongodb.WriteConcern;
import com.mongodb.WriteResult;
@@ -43,6 +44,9 @@ import com.yahoo.ycsb.DBException;
*/
public class MongoDbClient extends DB {
+ /** New option for auth*/
+ private static final String MONGO_AUTH = "mongodb.auth";
+
/** Used to include a field in a response. */
protected static final Integer INCLUDE = Integer.valueOf(1);
@@ -72,6 +76,8 @@ public class MongoDbClient extends DB {
// initialize MongoDb driver
Properties props = getProperties();
+ boolean authNeeded = Boolean.parseBoolean(
+ props.getProperty(MONGO_AUTH, "false"));
String url = props.getProperty("mongodb.url",
"mongodb://localhost:27017");
database = props.getProperty("mongodb.database", "ycsb");
@@ -108,6 +114,12 @@ public class MongoDbClient extends DB {
// strip out prefix since Java driver doesn't currently support
// standard connection format URL yet
// http://www.mongodb.org/display/DOCS/Connections
+ if (authNeeded) {
+ System.out.println("Attempting mongo connection with " + url);
+ mongo = new Mongo(new MongoURI(url));
+ System.out.println("Connection successful");
+ return;
+ }
if (url.startsWith("mongodb://")) {
url = url.substring(10);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment