code: /**
   * 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.
   */
  package com.yahoo.ycsb.db;
  
  import com.arangodb.ArangoConfigure;
  import com.arangodb.ArangoDriver;
  import com.arangodb.ArangoException;
  import com.arangodb.ArangoHost;
  import com.arangodb.DocumentCursor;
  import com.arangodb.ErrorNums;
  import com.arangodb.entity.BaseDocument;
  import com.arangodb.entity.DocumentEntity;
  import com.arangodb.entity.EntityFactory;
  import com.arangodb.entity.TransactionEntity;
  import com.arangodb.util.MapBuilder;
  
  import com.yahoo.ycsb.DB;
  import com.yahoo.ycsb.Status;
  import com.yahoo.ycsb.DBException;
  import com.yahoo.ycsb.ByteIterator;
  import com.yahoo.ycsb.StringByteIterator;
  
  import java.util.HashMap;
  import java.util.Iterator;
  import java.util.Map;
  import java.util.Properties;
  import java.util.Set;
  import java.util.Vector;
  import java.util.concurrent.atomic.AtomicInteger;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * ArangoDB binding for YCSB framework using the ArangoDB Inc. <a
   * href="https://github.com/arangodb/arangodb-java-driver">driver</a>
   * <p>
   * See the <code>README.md</code> for configuration information.
   * </p>
   * 
   * @see <a href="https://github.com/arangodb/arangodb-java-driver">ArangoDB Inc.
   *      driver</a>
   */
  public class ArangoDBClient extends DB {
  
    private static Logger logger = LoggerFactory.getLogger(ArangoDBClient.class);
    
    /**
     * The database name to access.
     */
    private static String databaseName = "ycsb";
  
    /**
     * Count the number of times initialized to teardown on the last
     * {@link #cleanup()}.
     */
    private static final AtomicInteger INIT_COUNT = new AtomicInteger(0);
  
    /** ArangoDB Driver related, Singleton. */
    private static ArangoDriver arangoDriver;
    private static Boolean dropDBBeforeRun;
    private static Boolean waitForSync = true;
    private static Boolean transactionUpdate = false;
  
    /**
     * Initialize any state for this DB. Called once per DB instance; there is
     * one DB instance per client thread.
     * 
     * Actually, one client process will share one DB instance here.(Coincide to
     * mongoDB driver)
     */
    @Override
    public void init() throws DBException {
      INIT_COUNT.incrementAndGet();
      synchronized (ArangoDBClient.class) {
        if (arangoDriver != null) {
          return;
        }
  
        Properties props = getProperties();
  
        // Set the DB address
        String ip = props.getProperty("arangodb.ip", "localhost");
        String portStr = props.getProperty("arangodb.port", "8529");
        int port = Integer.parseInt(portStr);
  
        // If clear db before run
        String dropDBBeforeRunStr = props.getProperty("arangodb.dropDBBeforeRun", "false");
        dropDBBeforeRun = Boolean.parseBoolean(dropDBBeforeRunStr);
        
        // Set the sync mode
        String waitForSyncStr = props.getProperty("arangodb.waitForSync", "false");
        waitForSync = Boolean.parseBoolean(waitForSyncStr);
        
        // Set if transaction for update
        String transactionUpdateStr = props.getProperty("arangodb.transactionUpdate", "false");
        transactionUpdate = Boolean.parseBoolean(transactionUpdateStr);
        
        // Init ArangoDB connection
        try {
          ArangoConfigure arangoConfigure = new ArangoConfigure();
          arangoConfigure.setArangoHost(new ArangoHost(ip, port));
          arangoConfigure.init();
          arangoDriver = new ArangoDriver(arangoConfigure);
        } catch (Exception e) {
          logger.error("Failed to initialize ArangoDB", e);
          System.exit(-1);
        }
  
        // Init the database
        if (dropDBBeforeRun) {
          // Try delete first
          try {
            arangoDriver.deleteDatabase(databaseName);
          } catch (ArangoException e) {
            if (e.getErrorNumber() != ErrorNums.ERROR_ARANGO_DATABASE_NOT_FOUND) {
              logger.error("Failed to delete database: {} with ex: {}", databaseName, e.toString());
              System.exit(-1);
            } else {
              logger.info("Fail to delete DB, already deleted: {}", databaseName);
            }
          }
        }
        try {
          arangoDriver.createDatabase(databaseName);
          logger.info("Database created: " + databaseName);
        } catch (ArangoException e) {
          if (e.getErrorNumber() != ErrorNums.ERROR_ARANGO_DUPLICATE_NAME) {
            logger.error("Failed to create database: {} with ex: {}", databaseName, e.toString());
            System.exit(-1);
          } else {
            logger.info("DB already exists: {}", databaseName);
          }
        }
        // Always set the default db
        arangoDriver.setDefaultDatabase(databaseName);
        logger.info("ArangoDB client connection created to {}:{}", ip, port);
        
        // Log the configuration
        logger.info("Arango Configuration: dropDBBeforeRun: {}; address: {}:{}; databaseName: {};"
                    + " waitForSync: {}; transactionUpdate: {};",
                    dropDBBeforeRun, ip, port, databaseName, waitForSync, transactionUpdate);
      }
    }
  
    /**
     * Cleanup any state for this DB. Called once per DB instance; there is one
     * DB instance per client thread.
     * 
     * Actually, one client process will share one DB instance here.(Coincide to
     * mongoDB driver)
     */
    @Override
    public void cleanup() throws DBException {
      if (INIT_COUNT.decrementAndGet() == 0) {
        arangoDriver = null;
        logger.info("Local cleaned up.");
      }
    }
  
    /**
     * 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 {
        BaseDocument toInsert = new BaseDocument(key);
        for (Map.Entry<String, ByteIterator> entry : values.entrySet()) {
          toInsert.addAttribute(entry.getKey(), byteIteratorToString(entry.getValue()));
        }
        arangoDriver.createDocument(table, toInsert, true/*create collection if not exist*/,
                                    waitForSync);
        return Status.OK;
      } catch (ArangoException e) {
        if (e.getErrorNumber() != ErrorNums.ERROR_ARANGO_UNIQUE_CONSTRAINT_VIOLATED) {
          logger.error("Fail to insert: {} {} with ex {}", table, key, e.toString());
        } else {
          logger.debug("Trying to create document with duplicate key: {} {}", table, key);
          return Status.BAD_REQUEST;
        }
      }  catch (RuntimeException e) {
        logger.error("Exception while trying insert {} {} with ex {}", table, key, e.toString());
      }
      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".
     */
    @SuppressWarnings("unchecked")
    @Override
    public Status read(String table, String key, Set<String> fields, HashMap<String, ByteIterator> result) {
      try {
        DocumentEntity<BaseDocument> targetDoc = arangoDriver.getDocument(table, key, BaseDocument.class);
        BaseDocument aDocument = targetDoc.getEntity();
        if (!this.fillMap(result, aDocument.getProperties(), fields)) {
          return Status.ERROR;
        }
        return Status.OK;
      } catch (ArangoException e) {
        if (e.getErrorNumber() != ErrorNums.ERROR_ARANGO_DOCUMENT_NOT_FOUND) {
          logger.error("Fail to read: {} {} with ex {}", table, key, e.toString());
        } else {
          logger.debug("Trying to read document not exist: {} {}", table, key);
          return Status.NOT_FOUND;
        }
      } catch (RuntimeException e) {
        logger.error("Exception while trying read {} {} with ex {}", table, key, e.toString());
      }
      return Status.ERROR;
    }
  
    /**
     * 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 {
        
        if (!transactionUpdate) {
          BaseDocument updateDoc = new BaseDocument();
          for (String field : values.keySet()) {
            updateDoc.addAttribute(field, byteIteratorToString(values.get(field)));
          }
          arangoDriver.updateDocument(table, key, updateDoc);
          return Status.OK;
        } else {
          // id for documentHandle
          String transactionAction = "function (id) {"
                 // use internal database functions
              + "var db = require('internal').db;"
                // collection.update(document, data, overwrite, keepNull, waitForSync)
              + String.format("db._update(id, %s, true, false, %s);}",
                  mapToJson(values), Boolean.toString(waitForSync).toLowerCase());
          TransactionEntity transaction = arangoDriver.createTransaction(transactionAction);
          transaction.addWriteCollection(table);
          transaction.setParams(createDocumentHandle(table, key));
          arangoDriver.executeTransaction(transaction);
          return Status.OK;
        }
      } catch (ArangoException e) {
        if (e.getErrorNumber() != ErrorNums.ERROR_ARANGO_DOCUMENT_NOT_FOUND) {
          logger.error("Fail to update: {} {} with ex {}", table, key, e.toString());
        } else {
          logger.debug("Trying to update document not exist: {} {}", table, key);
          return Status.NOT_FOUND;
        }
      } catch (RuntimeException e) {
        logger.error("Exception while trying update {} {} with ex {}", table, key, e.toString());
      }
      return Status.ERROR;
    }
  
    /**
     * 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 {
        arangoDriver.deleteDocument(table, key);
        return Status.OK;
      } catch (ArangoException e) {
        if (e.getErrorNumber() != ErrorNums.ERROR_ARANGO_DOCUMENT_NOT_FOUND) {
          logger.error("Fail to delete: {} {} with ex {}", table, key, e.toString());
        } else {
          logger.debug("Trying to delete document not exist: {} {}", table, key);
          return Status.NOT_FOUND;
        }
      } catch (RuntimeException e) {
        logger.error("Exception while trying delete {} {} with ex {}", table, key, 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) {
      DocumentCursor<BaseDocument> cursor = null;
      try {
        String aqlQuery = String.format(
            "FOR target IN %s FILTER target._key >= @key SORT target._key ASC LIMIT %d RETURN %s ", table,
            recordcount, constructReturnForAQL(fields, "target"));
  
        Map<String, Object> bindVars = new MapBuilder().put("key", startkey).get();
        cursor = arangoDriver.executeDocumentQuery(aqlQuery, bindVars, null, BaseDocument.class);
        Iterator<BaseDocument> iterator = cursor.entityIterator();
        while (iterator.hasNext()) {
          BaseDocument aDocument = iterator.next();
          HashMap<String, ByteIterator> aMap = new HashMap<String, ByteIterator>(aDocument.getProperties().size());
          if (!this.fillMap(aMap, aDocument.getProperties())) {
            return Status.ERROR;
          }
          result.add(aMap);
        }
        return Status.OK;
      } catch (Exception e) {
        logger.error("Exception while trying scan {} {} {} with ex {}", table, startkey, recordcount, e.toString());
      } finally {
        if (cursor != null) {
          try {
            cursor.close();
          } catch (ArangoException e) {
            logger.error("Fail to close cursor", e);
          }
        }
      }
      return Status.ERROR;
    }
  
    private String createDocumentHandle(String collectionName, String documentKey) throws ArangoException {
      validateCollectionName(collectionName);
      return collectionName + "/" + documentKey;
    }
  
    private void validateCollectionName(String name) throws ArangoException {
      if (name.indexOf('/') != -1) {
        throw new ArangoException("does not allow '/' in name.");
      }
    }
  
    
    private String constructReturnForAQL(Set<String> fields, String targetName) {
      // Construct the AQL query string.
      String resultDes = targetName;
      if (fields != null && fields.size() != 0) {
        StringBuilder builder = new StringBuilder("{");
        for (String field : fields) {
          builder.append(String.format("\n\"%s\" : %s.%s,", field, targetName, field));
        }
        //Replace last ',' to newline.
        builder.setCharAt(builder.length() - 1, '\n');
        builder.append("}");
        resultDes = builder.toString();
      }
      return resultDes;
    }
    
    private boolean fillMap(Map<String, ByteIterator> resultMap, Map<String, Object> properties) {
      return fillMap(resultMap, properties, null);
    }
    
    /**
     * Fills the map with the properties from the BaseDocument.
     * 
     * @param resultMap
     *      The map to fill/
     * @param obj
     *      The object to copy values from.
     * @return isSuccess
     */
    @SuppressWarnings("unchecked")
    private boolean fillMap(Map<String, ByteIterator> resultMap, Map<String, Object> properties, Set<String> fields) {
      if (fields == null || fields.size() == 0) {
        for (Map.Entry<String, Object> entry : properties.entrySet()) {
          if (entry.getValue() instanceof String) {
            resultMap.put(entry.getKey(),
                stringToByteIterator((String)(entry.getValue())));
          } else {
            logger.error("Error! Not the format expected! Actually is {}",
                entry.getValue().getClass().getName());
            return false;
          }
        }
      } else {
        for (String field : fields) {
          if (properties.get(field) instanceof String) {
            resultMap.put(field, stringToByteIterator((String)(properties.get(field))));
          } else {
            logger.error("Error! Not the format expected! Actually is {}",
                properties.get(field).getClass().getName());
            return false;
          }
        }
      }
      return true;
    }
    
    private String byteIteratorToString(ByteIterator byteIter) {
      return new String(byteIter.toArray());
    }
  
    private ByteIterator stringToByteIterator(String content) {
      return new StringByteIterator(content);
    }
    
    private String mapToJson(HashMap<String, ByteIterator> values) {
      HashMap<String, String> intervalRst = new HashMap<String, String>();
      for (Map.Entry<String, ByteIterator> entry : values.entrySet()) {
        intervalRst.put(entry.getKey(), byteIteratorToString(entry.getValue()));
      }
      return EntityFactory.toJsonString(intervalRst);
    }
    
  }
  