[azuredocumentdb] Added support for Azure DocumentDB (#838)

azuredocumentdb/pom.xml

+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+Copyright (c) 2016 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.
+-->
+
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
+  <modelVersion>4.0.0</modelVersion>
+  <parent>
+    <groupId>com.yahoo.ycsb</groupId>
+    <artifactId>binding-parent</artifactId>
+    <version>0.13.0-SNAPSHOT</version>
+    <relativePath>../binding-parent/</relativePath>
+  </parent>
+
+  <artifactId>azuredocumentdb-binding</artifactId>
+  <name>Azure DocumentDB Binding</name>
+  <dependencies>
+    <dependency>
+      <groupId>com.microsoft.azure</groupId>
+      <artifactId>azure-documentdb</artifactId>
+      <version>${azuredocumentdb.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>com.yahoo.ycsb</groupId>
+      <artifactId>core</artifactId>
+      <version>${project.version}</version>
+      <scope>provided</scope>
+    </dependency>
+  </dependencies>
+</project>

azuredocumentdb/src/main/java/com/yahoo/ycsb/db/azuredocumentdb/AzureDocumentDBClient.java

+/*
+ * Copyright 2016 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.
+ *
+ * This file 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.
+ */
+package com.yahoo.ycsb.db.azuredocumentdb;
+
+import com.yahoo.ycsb.*;
+
+import com.microsoft.azure.documentdb.ConnectionPolicy;
+import com.microsoft.azure.documentdb.ConsistencyLevel;
+import com.microsoft.azure.documentdb.Database;
+import com.microsoft.azure.documentdb.Document;
+import com.microsoft.azure.documentdb.DocumentClient;
+import com.microsoft.azure.documentdb.DocumentClientException;
+import com.microsoft.azure.documentdb.DocumentCollection;
+import com.microsoft.azure.documentdb.FeedOptions;
+
+import java.util.HashMap;
+import java.util.Map.Entry;
+import java.util.Set;
+import java.util.Vector;
+import java.util.List;
+
+/**
+ * Azure DocumentDB client binding.
+ */
+public class AzureDocumentDBClient extends DB {
+  private static String host;
+  private static String masterKey;
+  private static String databaseId;
+  private static Database database;
+  private static DocumentClient documentClient;
+  private static DocumentCollection collection;
+  private static FeedOptions feedOptions;
+
+  @Override
+  public void init() throws DBException {
+    host = getProperties().getProperty("documentdb.host", null);
+    masterKey = getProperties().getProperty("documentdb.masterKey", null);
+
+    if (host == null) {
+      System.err.println("ERROR: 'documentdb.host' must be set!");
+      System.exit(1);
+    }
+
+    if (masterKey == null) {
+      System.err.println("ERROR: 'documentdb.masterKey' must be set!");
+      System.exit(1);
+    }
+
+    databaseId = getProperties().getProperty("documentdb.databaseId", "ycsb");
+    String collectionId =
+        getProperties().getProperty("documentdb.collectionId", "usertable");
+    documentClient =
+        new DocumentClient(host, masterKey, ConnectionPolicy.GetDefault(),
+                           ConsistencyLevel.Session);
+    try {
+      // Initialize test database and collection.
+      collection = getCollection(collectionId);
+    } catch (DocumentClientException e) {
+      throw new DBException("Initialze collection failed", e);
+    }
+
+    feedOptions = new FeedOptions();
+    feedOptions.setEmitVerboseTracesInQuery(false);
+  }
+
+  @Override
+  public Status read(String table, String key, Set<String> fields,
+                     HashMap<String, ByteIterator> result) {
+    Document record = getDocumentById(table, key);
+
+    if (record != null) {
+      Set<String> fieldsToReturn =
+          (fields == null ? record.getHashMap().keySet() : fields);
+
+      for (String field : fieldsToReturn) {
+        if (field.startsWith("_")) {
+          continue;
+        }
+        result.put(field, new StringByteIterator(record.getString(field)));
+      }
+      return Status.OK;
+    }
+    // Unable to find the specidifed document.
+    return Status.ERROR;
+  }
+
+  @Override
+  public Status update(String table, String key,
+                       HashMap<String, ByteIterator> values) {
+    Document record = getDocumentById(table, key);
+
+    if (record == null) {
+      return Status.ERROR;
+    }
+
+    // Update each field.
+    for (Entry<String, ByteIterator> val : values.entrySet()) {
+      record.set(val.getKey(), val.getValue().toString());
+    }
+
+    // Replace the document.
+    try {
+      documentClient.replaceDocument(record, null);
+    } catch (DocumentClientException e) {
+      e.printStackTrace(System.err);
+      return Status.ERROR;
+    }
+
+    return Status.OK;
+  }
+
+  @Override
+  public Status insert(String table, String key,
+                       HashMap<String, ByteIterator> values) {
+    Document record = new Document();
+
+    record.set("id", key);
+
+    for (Entry<String, ByteIterator> val : values.entrySet()) {
+      record.set(val.getKey(), val.getValue().toString());
+    }
+
+    try {
+      documentClient.createDocument(collection.getSelfLink(), record, null,
+                                    false);
+    } catch (DocumentClientException e) {
+      e.printStackTrace(System.err);
+      return Status.ERROR;
+    }
+    return Status.OK;
+  }
+
+  @Override
+  public Status delete(String table, String key) {
+    Document record = getDocumentById(table, key);
+
+    try {
+      // Delete the document by self link.
+      documentClient.deleteDocument(record.getSelfLink(), null);
+    } catch (DocumentClientException e) {
+      e.printStackTrace();
+      return Status.ERROR;
+    }
+
+    return Status.OK;
+  }
+
+  @Override
+  public Status scan(String table, String startkey, int recordcount,
+                     Set<String> fields,
+                     Vector<HashMap<String, ByteIterator>> result) {
+    // TODO: Implement Scan as query on primary key.
+    return Status.NOT_IMPLEMENTED;
+  }
+
+  private Database getDatabase() {
+    if (database == null) {
+      // Get the database if it exists
+      List<Database> databaseList =
+          documentClient
+              .queryDatabases(
+                  "SELECT * FROM root r WHERE r.id='" + databaseId + "'", null)
+              .getQueryIterable()
+              .toList();
+
+      if (databaseList.size() > 0) {
+        // Cache the database object so we won't have to query for it
+        // later to retrieve the selfLink.
+        database = databaseList.get(0);
+      } else {
+        // Create the database if it doesn't exist.
+        try {
+          Database databaseDefinition = new Database();
+          databaseDefinition.setId(databaseId);
+
+          database = documentClient.createDatabase(databaseDefinition, null)
+                         .getResource();
+        } catch (DocumentClientException e) {
+          // TODO: Something has gone terribly wrong - the app wasn't
+          // able to query or create the collection.
+          // Verify your connection, endpoint, and key.
+          e.printStackTrace(System.err);
+        }
+      }
+    }
+
+    return database;
+  }
+
+  private DocumentCollection getCollection(String collectionId)
+      throws DocumentClientException {
+    if (collection == null) {
+      // Get the collection if it exists.
+      List<DocumentCollection> collectionList =
+          documentClient
+              .queryCollections(getDatabase().getSelfLink(),
+                                "SELECT * FROM root r WHERE r.id='" +
+                                    collectionId + "'",
+                                null)
+              .getQueryIterable()
+              .toList();
+
+      if (collectionList.size() > 0) {
+        // Cache the collection object so we won't have to query for it
+        // later to retrieve the selfLink.
+        collection = collectionList.get(0);
+      } else {
+        // Create the collection if it doesn't exist.
+        try {
+          DocumentCollection collectionDefinition = new DocumentCollection();
+          collectionDefinition.setId(collectionId);
+
+          collection = documentClient
+                           .createCollection(getDatabase().getSelfLink(),
+                                             collectionDefinition, null)
+                           .getResource();
+        } catch (DocumentClientException e) {
+          // TODO: Something has gone terribly wrong - the app wasn't
+          // able to query or create the collection.
+          // Verify your connection, endpoint, and key.
+          e.printStackTrace(System.err);
+          throw e;
+        }
+      }
+    }
+
+    return collection;
+  }
+
+  private Document getDocumentById(String collectionId, String id) {
+    if (collection == null) {
+      return null;
+    }
+    // Retrieve the document using the DocumentClient.
+    List<Document> documentList =
+        documentClient
+            .queryDocuments(collection.getSelfLink(),
+                            "SELECT * FROM root r WHERE r.id='" + id + "'",
+                            feedOptions)
+            .getQueryIterable()
+            .toList();
+
+    if (documentList.size() > 0) {
+      return documentList.get(0);
+    }
+    return null;
+  }
+}

azuredocumentdb/src/main/java/com/yahoo/ycsb/db/azuredocumentdb/package-info.java

+/*
+ * Copyright 2016 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.
+ */
+
+/**
+ * The YCSB binding for Azure DocumentDB.
+ */
+package com.yahoo.ycsb.db.azuredocumentdb;
+

bin/bindings.properties

@@ -35,6 +35,7 @@ cassandra-cql:com.yahoo.ycsb.db.CassandraCQLClient
 cassandra2-cql:com.yahoo.ycsb.db.CassandraCQLClient
 couchbase:com.yahoo.ycsb.db.CouchbaseClient
 couchbase2:com.yahoo.ycsb.db.couchbase2.Couchbase2Client
+azuredocumentdb:com.yahoo.ycsb.db.azuredocumentdb.AzureDocumentDBClient
 dynamodb:com.yahoo.ycsb.db.DynamoDBClient
 elasticsearch:com.yahoo.ycsb.db.ElasticsearchClient
 geode:com.yahoo.ycsb.db.GeodeClient

bin/ycsb

@@ -60,6 +60,7 @@ DATABASES = {
     "cassandra2-cql": "com.yahoo.ycsb.db.CassandraCQLClient",
     "couchbase"    : "com.yahoo.ycsb.db.CouchbaseClient",
     "couchbase2"   : "com.yahoo.ycsb.db.couchbase2.Couchbase2Client",
+    "azuredocumentdb"   : "com.yahoo.ycsb.db.azuredocumentdb.AzureDocumentDBClient",
     "dynamodb"     : "com.yahoo.ycsb.db.DynamoDBClient",
     "elasticsearch": "com.yahoo.ycsb.db.ElasticsearchClient",
     "geode"        : "com.yahoo.ycsb.db.GeodeClient",

distribution/pom.xml

@@ -74,6 +74,11 @@ LICENSE file.
       <artifactId>couchbase2-binding</artifactId>
       <version>${project.version}</version>
     </dependency>
+    <dependency>
+      <groupId>com.yahoo.ycsb</groupId>
+      <artifactId>azuredocumentdb-binding</artifactId>
+      <version>${project.version}</version>
+    </dependency>
     <dependency>
       <groupId>com.yahoo.ycsb</groupId>
       <artifactId>dynamodb-binding</artifactId>

pom.xml

@@ -75,6 +75,7 @@ LICENSE file.
     <accumulo.version>1.6.0</accumulo.version>
     <cassandra.cql.version>3.0.0</cassandra.cql.version>
     <geode.version>1.0.0-incubating.M3</geode.version>
+    <azuredocumentdb.version>1.8.1</azuredocumentdb.version>
     <googlebigtable.version>0.2.3</googlebigtable.version>
     <infinispan.version>7.2.2.Final</infinispan.version>
     <kudu.version>1.1.0</kudu.version>
@@ -114,6 +115,7 @@ LICENSE file.
     <module>couchbase</module>
     <module>couchbase2</module>
     <module>distribution</module>
+    <module>azuredocumentdb</module>
     <module>dynamodb</module>
     <module>elasticsearch</module>
     <module>geode</module>