code: /**
   * 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.
   */
  
  package com.yahoo.ycsb.db;
  
  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 org.apache.solr.client.solrj.SolrQuery;
  import org.apache.solr.client.solrj.impl.CloudSolrClient;
  import org.apache.solr.client.solrj.impl.HttpSolrClient;
  import org.apache.solr.client.solrj.response.QueryResponse;
  import org.apache.solr.client.solrj.response.UpdateResponse;
  import org.apache.solr.client.solrj.SolrServerException;
  import org.apache.solr.common.SolrDocument;
  import org.apache.solr.common.SolrDocumentList;
  import org.apache.solr.common.SolrInputDocument;
  
  import java.io.IOException;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.Map.Entry;
  import java.util.Properties;
  import java.util.Set;
  import java.util.Vector;
  
  /**
   * Solr client for YCSB framework.
   *
   * <p>
   * Default properties to set:
   * </p>
   * <ul>
   * See README.md
   * </ul>
   *
   */
  public class SolrClient extends DB {
  
    public static final String DEFAULT_CLOUD_MODE = "false";
    public static final String DEFAULT_BATCH_MODE = "false";
    public static final String DEFAULT_ZOOKEEPER_HOSTS = "localhost:2181";
    public static final String DEFAULT_SOLR_BASE_URL = "http://localhost:8983/solr";
    public static final String DEFAULT_COMMIT_WITHIN_TIME = "1000";
  
    private org.apache.solr.client.solrj.SolrClient client;
    private Integer commitTime;
    private Boolean batchMode;
  
    /**
     * Initialize any state for this DB. Called once per DB instance; there is one DB instance per
     * client thread.
     */
    @Override
    public void init() throws DBException {
      Properties props = getProperties();
      commitTime = Integer
          .parseInt(props.getProperty("solr.commit.within.time", DEFAULT_COMMIT_WITHIN_TIME));
      batchMode = Boolean.parseBoolean(props.getProperty("solr.batch.mode", DEFAULT_BATCH_MODE));
  
      // Check if Solr cluster is running in SolrCloud or Stand-alone mode
      Boolean cloudMode = Boolean.parseBoolean(props.getProperty("solr.cloud", DEFAULT_CLOUD_MODE));
      System.err.println("Solr Cloud Mode = " + cloudMode);
      if (cloudMode) {
        System.err.println("Solr Zookeeper Remote Hosts = "
            + props.getProperty("solr.zookeeper.hosts", DEFAULT_ZOOKEEPER_HOSTS));
        client = new CloudSolrClient(
            props.getProperty("solr.zookeeper.hosts", DEFAULT_ZOOKEEPER_HOSTS));
      } else {
        client = new HttpSolrClient(props.getProperty("solr.base.url", DEFAULT_SOLR_BASE_URL));
      }
    }
  
    @Override
    public void cleanup() throws DBException {
      try {
        client.close();
      } catch (IOException e) {
        throw new DBException(e);
      }
    }
  
    /**
     * 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 this class's description for a
     *         discussion of error codes.
     */
    @Override
    public Status insert(String table, String key, HashMap<String, ByteIterator> values) {
      try {
        SolrInputDocument doc = new SolrInputDocument();
  
        doc.addField("id", key);
        for (Entry<String, String> entry : StringByteIterator.getStringMap(values).entrySet()) {
          doc.addField(entry.getKey(), entry.getValue());
        }
        UpdateResponse response;
        if (batchMode) {
          response = client.add(table, doc, commitTime);
        } else {
          response = client.add(table, doc);
          client.commit(table);
        }
        return checkStatus(response.getStatus());
  
      } catch (IOException | SolrServerException e) {
        e.printStackTrace();
      }
      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 this class's description for a
     *         discussion of error codes.
     */
    @Override
    public Status delete(String table, String key) {
      try {
        UpdateResponse response;
        if (batchMode) {
          response = client.deleteById(table, key, commitTime);
        } else {
          response = client.deleteById(table, key);
          client.commit(table);
        }
        return checkStatus(response.getStatus());
      } catch (IOException | SolrServerException e) {
        e.printStackTrace();
      }
      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".
     */
    @Override
    public Status read(String table, String key, Set<String> fields,
        HashMap<String, ByteIterator> result) {
      try {
        Boolean returnFields = false;
        String[] fieldList = null;
        if (fields != null) {
          returnFields = true;
          fieldList = fields.toArray(new String[fields.size()]);
        }
        SolrQuery query = new SolrQuery();
        query.setQuery("id:" + key);
        if (returnFields) {
          query.setFields(fieldList);
        }
        final QueryResponse response = client.query(table, query);
        SolrDocumentList results = response.getResults();
        if ((results != null) && (results.getNumFound() > 0)) {
          for (String field : results.get(0).getFieldNames()) {
            result.put(field,
                new StringByteIterator(String.valueOf(results.get(0).getFirstValue(field))));
          }
        }
        return checkStatus(response.getStatus());
      } catch (IOException | SolrServerException e) {
        e.printStackTrace();
      }
      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 {
        SolrInputDocument updatedDoc = new SolrInputDocument();
        updatedDoc.addField("id", key);
  
        for (Entry<String, String> entry : StringByteIterator.getStringMap(values).entrySet()) {
          updatedDoc.addField(entry.getKey(), Collections.singletonMap("set", entry.getValue()));
        }
  
        UpdateResponse writeResponse;
        if (batchMode) {
          writeResponse = client.add(table, updatedDoc, commitTime);
        } else {
          writeResponse = client.add(table, updatedDoc);
          client.commit(table);
        }
        return checkStatus(writeResponse.getStatus());
      } catch (IOException | SolrServerException e) {
        e.printStackTrace();
      }
      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 this 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) {
      try {
        Boolean returnFields = false;
        String[] fieldList = null;
        if (fields != null) {
          returnFields = true;
          fieldList = fields.toArray(new String[fields.size()]);
        }
        SolrQuery query = new SolrQuery();
        query.setQuery("*:*");
        query.setParam("fq", "id:[ " + startkey + " TO * ]");
        if (returnFields) {
          query.setFields(fieldList);
        }
        query.setRows(recordcount);
        final QueryResponse response = client.query(table, query);
        SolrDocumentList results = response.getResults();
  
        HashMap<String, ByteIterator> entry;
  
        for (SolrDocument hit : results) {
          entry = new HashMap<String, ByteIterator>((int) results.getNumFound());
          for (String field : hit.getFieldNames()) {
            entry.put(field, new StringByteIterator(String.valueOf(hit.getFirstValue(field))));
          }
          result.add(entry);
        }
        return checkStatus(response.getStatus());
      } catch (IOException | SolrServerException e) {
        e.printStackTrace();
      }
      return Status.ERROR;
    }
  
    private Status checkStatus(int status) {
      Status responseStatus;
      switch (status) {
      case 0:
        responseStatus = Status.OK;
        break;
      case 400:
        responseStatus = Status.BAD_REQUEST;
        break;
      case 403:
        responseStatus = Status.FORBIDDEN;
        break;
      case 404:
        responseStatus = Status.NOT_FOUND;
        break;
      case 500:
        responseStatus = Status.ERROR;
        break;
      case 503:
        responseStatus = Status.SERVICE_UNAVAILABLE;
        break;
      default:
        responseStatus = Status.UNEXPECTED_STATE;
        break;
      }
      return responseStatus;
    }
  
  }
  