code: /**
   * Copyright (c) 2013 - 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.
   */
  
  package com.yahoo.ycsb.db;
  
  import com.couchbase.client.protocol.views.*;
  import com.fasterxml.jackson.core.JsonFactory;
  import com.fasterxml.jackson.core.JsonGenerator;
  import com.fasterxml.jackson.databind.JsonNode;
  import com.fasterxml.jackson.databind.ObjectMapper;
  import com.fasterxml.jackson.databind.node.ObjectNode;
  import com.yahoo.ycsb.ByteIterator;
  import com.yahoo.ycsb.DB;
  import com.yahoo.ycsb.DBException;
  import com.yahoo.ycsb.Status;
  import com.yahoo.ycsb.StringByteIterator;
  
  import net.spy.memcached.PersistTo;
  import net.spy.memcached.ReplicateTo;
  import net.spy.memcached.internal.OperationFuture;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.StringWriter;
  import java.io.Writer;
  import java.net.URI;
  import java.util.Arrays;
  import java.util.HashMap;
  import java.util.Iterator;
  import java.util.Map;
  import java.util.Properties;
  import java.util.Set;
  import java.util.Vector;
  
  /**
   * A class that wraps the CouchbaseClient to allow it to be interfaced with YCSB.
   * This class extends {@link DB} and implements the database interface used by YCSB client.
   */
  public class CouchbaseClient extends DB {
    public static final String URL_PROPERTY = "couchbase.url";
    public static final String BUCKET_PROPERTY = "couchbase.bucket";
    public static final String PASSWORD_PROPERTY = "couchbase.password";
    public static final String CHECKF_PROPERTY = "couchbase.checkFutures";
    public static final String PERSIST_PROPERTY = "couchbase.persistTo";
    public static final String REPLICATE_PROPERTY = "couchbase.replicateTo";
    public static final String JSON_PROPERTY = "couchbase.json";
    public static final String DESIGN_DOC_PROPERTY = "couchbase.ddoc";
    public static final String VIEW_PROPERTY = "couchbase.view";
    public static final String STALE_PROPERTY = "couchbase.stale";
    public static final String SCAN_PROPERTY = "scanproportion";
  
    public static final String STALE_PROPERTY_DEFAULT = Stale.OK.name();
    public static final String SCAN_PROPERTY_DEFAULT = "0.0";
  
    protected static final ObjectMapper JSON_MAPPER = new ObjectMapper();
  
    private com.couchbase.client.CouchbaseClient client;
    private PersistTo persistTo;
    private ReplicateTo replicateTo;
    private boolean checkFutures;
    private boolean useJson;
    private String designDoc;
    private String viewName;
    private Stale stale;
    private View view;
    private final Logger log = LoggerFactory.getLogger(getClass());
  
    @Override
    public void init() throws DBException {
      Properties props = getProperties();
  
      String url = props.getProperty(URL_PROPERTY, "http://127.0.0.1:8091/pools");
      String bucket = props.getProperty(BUCKET_PROPERTY, "default");
      String password = props.getProperty(PASSWORD_PROPERTY, "");
  
      checkFutures = props.getProperty(CHECKF_PROPERTY, "true").equals("true");
      useJson = props.getProperty(JSON_PROPERTY, "true").equals("true");
  
      persistTo = parsePersistTo(props.getProperty(PERSIST_PROPERTY, "0"));
      replicateTo = parseReplicateTo(props.getProperty(REPLICATE_PROPERTY, "0"));
  
      designDoc = getProperties().getProperty(DESIGN_DOC_PROPERTY);
      viewName = getProperties().getProperty(VIEW_PROPERTY);
      stale = Stale.valueOf(getProperties().getProperty(STALE_PROPERTY, STALE_PROPERTY_DEFAULT).toUpperCase());
  
      Double scanproportion = Double.valueOf(props.getProperty(SCAN_PROPERTY, SCAN_PROPERTY_DEFAULT));
  
      Properties systemProperties = System.getProperties();
      systemProperties.put("net.spy.log.LoggerImpl", "net.spy.memcached.compat.log.SLF4JLogger");
      System.setProperties(systemProperties);
  
      try {
        client = new com.couchbase.client.CouchbaseClient(Arrays.asList(new URI(url)), bucket, password);
      } catch (Exception e) {
        throw new DBException("Could not create CouchbaseClient object.", e);
      }
  
      if (scanproportion > 0) {
        try {
          view = client.getView(designDoc, viewName);
        } catch (Exception e) {
          throw new DBException(String.format("%s=%s and %s=%s provided, unable to connect to view.",
            DESIGN_DOC_PROPERTY, designDoc, VIEW_PROPERTY, viewName), e.getCause());
        }
      }
    }
  
    /**
     * Parse the replicate property into the correct enum.
     *
     * @param property the stringified property value.
     * @throws DBException if parsing the property did fail.
     * @return the correct enum.
     */
    private ReplicateTo parseReplicateTo(final String property) throws DBException {
      int value = Integer.parseInt(property);
  
      switch (value) {
      case 0:
        return ReplicateTo.ZERO;
      case 1:
        return ReplicateTo.ONE;
      case 2:
        return ReplicateTo.TWO;
      case 3:
        return ReplicateTo.THREE;
      default:
        throw new DBException(REPLICATE_PROPERTY + " must be between 0 and 3");
      }
    }
  
    /**
     * Parse the persist property into the correct enum.
     *
     * @param property the stringified property value.
     * @throws DBException if parsing the property did fail.
     * @return the correct enum.
     */
    private PersistTo parsePersistTo(final String property) throws DBException {
      int value = Integer.parseInt(property);
  
      switch (value) {
      case 0:
        return PersistTo.ZERO;
      case 1:
        return PersistTo.ONE;
      case 2:
        return PersistTo.TWO;
      case 3:
        return PersistTo.THREE;
      case 4:
        return PersistTo.FOUR;
      default:
        throw new DBException(PERSIST_PROPERTY + " must be between 0 and 4");
      }
    }
  
    /**
     * Shutdown the client.
     */
    @Override
    public void cleanup() {
      client.shutdown();
    }
  
    @Override
    public Status read(final String table, final String key, final Set<String> fields,
                       final HashMap<String, ByteIterator> result) {
      String formattedKey = formatKey(table, key);
  
      try {
        Object loaded = client.get(formattedKey);
  
        if (loaded == null) {
          return Status.ERROR;
        }
  
        decode(loaded, fields, result);
        return Status.OK;
      } catch (Exception e) {
        if (log.isErrorEnabled()) {
          log.error("Could not read value for key " + formattedKey, e);
        }
        return Status.ERROR;
      }
    }
  
    @Override
    public Status scan(final String table, final String startkey, final int recordcount, final Set<String> fields,
                       final Vector<HashMap<String, ByteIterator>> result) {
      try {
        Query query = new Query().setRangeStart(startkey)
            .setLimit(recordcount)
            .setIncludeDocs(true)
            .setStale(stale);
        ViewResponse response = client.query(view, query);
  
        for (ViewRow row : response) {
          HashMap<String, ByteIterator> rowMap = new HashMap();
          decode(row.getDocument(), fields, rowMap);
          result.add(rowMap);
        }
  
        return Status.OK;
      } catch (Exception e) {
        log.error(e.getMessage());
      }
  
      return Status.ERROR;
    }
  
    @Override
    public Status update(final String table, final String key, final HashMap<String, ByteIterator> values) {
      String formattedKey = formatKey(table, key);
  
      try {
        final OperationFuture<Boolean> future = client.replace(formattedKey, encode(values), persistTo, replicateTo);
        return checkFutureStatus(future);
      } catch (Exception e) {
        if (log.isErrorEnabled()) {
          log.error("Could not update value for key " + formattedKey, e);
        }
        return Status.ERROR;
      }
    }
  
    @Override
    public Status insert(final String table, final String key, final HashMap<String, ByteIterator> values) {
      String formattedKey = formatKey(table, key);
  
      try {
        final OperationFuture<Boolean> future = client.add(formattedKey, encode(values), persistTo, replicateTo);
        return checkFutureStatus(future);
      } catch (Exception e) {
        if (log.isErrorEnabled()) {
          log.error("Could not insert value for key " + formattedKey, e);
        }
        return Status.ERROR;
      }
    }
  
    @Override
    public Status delete(final String table, final String key) {
      String formattedKey = formatKey(table, key);
  
      try {
        final OperationFuture<Boolean> future = client.delete(formattedKey, persistTo, replicateTo);
        return checkFutureStatus(future);
      } catch (Exception e) {
        if (log.isErrorEnabled()) {
          log.error("Could not delete value for key " + formattedKey, e);
        }
        return Status.ERROR;
      }
    }
  
    /**
     * Prefix the key with the given prefix, to establish a unique namespace.
     *
     * @param prefix the prefix to use.
     * @param key the actual key.
     * @return the formatted and prefixed key.
     */
    private String formatKey(final String prefix, final String key) {
      return prefix + ":" + key;
    }
  
    /**
     * Wrapper method that either inspects the future or not.
     *
     * @param future the future to potentially verify.
     * @return the status of the future result.
     */
    private Status checkFutureStatus(final OperationFuture<?> future) {
      if (checkFutures) {
        return future.getStatus().isSuccess() ? Status.OK : Status.ERROR;
      } else {
        return Status.OK;
      }
    }
  
    /**
     * Decode the object from server into the storable result.
     *
     * @param source the loaded object.
     * @param fields the fields to check.
     * @param dest the result passed back to the ycsb core.
     */
    private void decode(final Object source, final Set<String> fields, final HashMap<String, ByteIterator> dest) {
      if (useJson) {
        try {
          JsonNode json = JSON_MAPPER.readTree((String) source);
          boolean checkFields = fields != null && !fields.isEmpty();
          for (Iterator<Map.Entry<String, JsonNode>> jsonFields = json.fields(); jsonFields.hasNext();) {
            Map.Entry<String, JsonNode> jsonField = jsonFields.next();
            String name = jsonField.getKey();
            if (checkFields && fields.contains(name)) {
              continue;
            }
            JsonNode jsonValue = jsonField.getValue();
            if (jsonValue != null && !jsonValue.isNull()) {
              dest.put(name, new StringByteIterator(jsonValue.asText()));
            }
          }
        } catch (Exception e) {
          throw new RuntimeException("Could not decode JSON");
        }
      } else {
        HashMap<String, String> converted = (HashMap<String, String>) source;
        for (Map.Entry<String, String> entry : converted.entrySet()) {
          dest.put(entry.getKey(), new StringByteIterator(entry.getValue()));
        }
      }
    }
  
    /**
     * Encode the object for couchbase storage.
     *
     * @param source the source value.
     * @return the storable object.
     */
    private Object encode(final HashMap<String, ByteIterator> source) {
      HashMap<String, String> stringMap = StringByteIterator.getStringMap(source);
      if (!useJson) {
        return stringMap;
      }
  
      ObjectNode node = JSON_MAPPER.createObjectNode();
      for (Map.Entry<String, String> pair : stringMap.entrySet()) {
        node.put(pair.getKey(), pair.getValue());
      }
      JsonFactory jsonFactory = new JsonFactory();
      Writer writer = new StringWriter();
      try {
        JsonGenerator jsonGenerator = jsonFactory.createGenerator(writer);
        JSON_MAPPER.writeTree(jsonGenerator, node);
      } catch (Exception e) {
        throw new RuntimeException("Could not encode JSON value");
      }
      return writer.toString();
    }
  }
  