brianfrankcooper · risdenk · Dec 16, 2016 · Sep 2, 2016 · Sep 2, 2016 · Sep 2, 2016
diff --git a/azuredocumentdb/pom.xml b/azuredocumentdb/pom.xml
@@ -0,0 +1,46 @@
+<?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>
+  <properties>
+      <checkstyle.failOnViolation>false</checkstyle.failOnViolation>
+  </properties>
+  <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>
diff --git a/azuredocumentdb/src/main/java/com/yahoo/ycsb/db/azuredocumentdb/AzureDocumentDBClient.java b/azuredocumentdb/src/main/java/com/yahoo/ycsb/db/azuredocumentdb/AzureDocumentDBClient.java
@@ -0,0 +1,230 @@
+/*
+ * 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.DocumentClient;
+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;
+
+public class AzureDocumentDBClient extends DB {
+  private static String host;
+  private static String masterKey;
+  private static String databaseId;
+  private static String collectionId;
+  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");
+    collectionId = getProperties().getProperty("documentdb.collectionId", "usertable");
+    documentClient = new DocumentClient(host, masterKey,
+            ConnectionPolicy.GetDefault(), ConsistencyLevel.Session);
+    // Initialize test database and collection.
+    getCollection(collectionId);
+
+    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();
+      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();
+      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();
+            }
+        }
+    }
+
+    return database;
+  }
+
+  private DocumentCollection getCollection(String collectionId) {
+    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();
+        }
+      }
+    }
+
+    return collection;
+  }
+
+  private Document getDocumentById(String collectionId, String id) {
+    // Retrieve the document using the DocumentClient.
+    List<Document> documentList = documentClient
+            .queryDocuments(getCollection(collectionId).getSelfLink(),
+                    "SELECT * FROM root r WHERE r.id='" + id + "'", feedOptions)
+            .getQueryIterable().toList();
+
+    if (documentList.size() > 0) {
+      return documentList.get(0);
+    }
+    return null;
+  }
+}
diff --git a/azuredocumentdb/src/main/java/com/yahoo/ycsb/db/azuredocumentdb/package-info.java b/azuredocumentdb/src/main/java/com/yahoo/ycsb/db/azuredocumentdb/package-info.java
@@ -0,0 +1,22 @@
+/*
+ * 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;
+
diff --git a/bin/bindings.properties b/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

diff --git a/bin/ycsb b/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",

diff --git a/distribution/pom.xml b/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>

diff --git a/pom.xml b/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>