PageRenderTime 960ms CodeModel.GetById 24ms RepoModel.GetById 7ms app.codeStats 0ms

/ycsb/YCSB/elasticsearch/src/main/java/com/yahoo/ycsb/db/ElasticsearchClient.java

https://gitlab.com/sudarsunkannan/mapkeeper
Java | 353 lines | 211 code | 33 blank | 109 comment | 22 complexity | 3f46e66151ba0827a5a90fa0a3dcca4f MD5 | raw file
  1. /**
  2. * Copyright (c) 2012 YCSB contributors. All rights reserved.
  3. *
  4. * Licensed under the Apache License, Version 2.0 (the "License"); you
  5. * may not use this file except in compliance with the License. You
  6. * may obtain a copy of the License at
  7. *
  8. * http://www.apache.org/licenses/LICENSE-2.0
  9. *
  10. * Unless required by applicable law or agreed to in writing, software
  11. * distributed under the License is distributed on an "AS IS" BASIS,
  12. * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
  13. * implied. See the License for the specific language governing
  14. * permissions and limitations under the License. See accompanying
  15. * LICENSE file.
  16. */
  17. package com.yahoo.ycsb.db;
  18. import static org.elasticsearch.common.settings.Settings.Builder;
  19. import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
  20. import static org.elasticsearch.index.query.QueryBuilders.rangeQuery;
  21. import static org.elasticsearch.node.NodeBuilder.nodeBuilder;
  22. import com.yahoo.ycsb.ByteIterator;
  23. import com.yahoo.ycsb.DB;
  24. import com.yahoo.ycsb.DBException;
  25. import com.yahoo.ycsb.Status;
  26. import com.yahoo.ycsb.StringByteIterator;
  27. import org.elasticsearch.action.admin.cluster.health.ClusterHealthRequest;
  28. import org.elasticsearch.action.admin.indices.create.CreateIndexRequest;
  29. import org.elasticsearch.action.get.GetResponse;
  30. import org.elasticsearch.action.search.SearchResponse;
  31. import org.elasticsearch.client.Client;
  32. import org.elasticsearch.client.Requests;
  33. import org.elasticsearch.client.transport.TransportClient;
  34. import org.elasticsearch.common.settings.Settings;
  35. import org.elasticsearch.common.transport.InetSocketTransportAddress;
  36. import org.elasticsearch.common.xcontent.XContentBuilder;
  37. import org.elasticsearch.index.query.RangeQueryBuilder;
  38. import org.elasticsearch.node.Node;
  39. import org.elasticsearch.search.SearchHit;
  40. import java.net.InetAddress;
  41. import java.net.UnknownHostException;
  42. import java.util.HashMap;
  43. import java.util.Map.Entry;
  44. import java.util.Properties;
  45. import java.util.Set;
  46. import java.util.Vector;
  47. /**
  48. * Elasticsearch client for YCSB framework.
  49. *
  50. * <p>
  51. * Default properties to set:
  52. * </p>
  53. * <ul>
  54. * <li>cluster.name = es.ycsb.cluster
  55. * <li>es.index.key = es.ycsb
  56. * <li>es.number_of_shards = 1
  57. * <li>es.number_of_replicas = 0
  58. * </ul>
  59. */
  60. public class ElasticsearchClient extends DB {
  61. private static final String DEFAULT_CLUSTER_NAME = "es.ycsb.cluster";
  62. private static final String DEFAULT_INDEX_KEY = "es.ycsb";
  63. private static final String DEFAULT_REMOTE_HOST = "localhost:9300";
  64. private static final int NUMBER_OF_SHARDS = 1;
  65. private static final int NUMBER_OF_REPLICAS = 0;
  66. private Node node;
  67. private Client client;
  68. private String indexKey;
  69. private Boolean remoteMode;
  70. /**
  71. * Initialize any state for this DB. Called once per DB instance; there is one
  72. * DB instance per client thread.
  73. */
  74. @Override
  75. public void init() throws DBException {
  76. Properties props = getProperties();
  77. this.indexKey = props.getProperty("es.index.key", DEFAULT_INDEX_KEY);
  78. int numberOfShards = parseIntegerProperty(props, "es.number_of_shards", NUMBER_OF_SHARDS);
  79. int numberOfReplicas = parseIntegerProperty(props, "es.number_of_replicas", NUMBER_OF_REPLICAS);
  80. // Check if transport client needs to be used (To connect to multiple
  81. // elasticsearch nodes)
  82. remoteMode = Boolean.parseBoolean(props.getProperty("es.remote", "false"));
  83. Boolean newdb = Boolean.parseBoolean(props.getProperty("es.newdb", "false"));
  84. Builder settings = Settings.settingsBuilder()
  85. .put("cluster.name", DEFAULT_CLUSTER_NAME)
  86. .put("node.local", Boolean.toString(!remoteMode))
  87. .put("path.home", System.getProperty("java.io.tmpdir"));
  88. // if properties file contains elasticsearch user defined properties
  89. // add it to the settings file (will overwrite the defaults).
  90. settings.put(props);
  91. final String clusterName = settings.get("cluster.name");
  92. System.err.println("Elasticsearch starting node = " + clusterName);
  93. System.err.println("Elasticsearch node path.home = " + settings.get("path.home"));
  94. System.err.println("Elasticsearch Remote Mode = " + remoteMode);
  95. // Remote mode support for connecting to remote elasticsearch cluster
  96. if (remoteMode) {
  97. settings.put("client.transport.sniff", true)
  98. .put("client.transport.ignore_cluster_name", false)
  99. .put("client.transport.ping_timeout", "30s")
  100. .put("client.transport.nodes_sampler_interval", "30s");
  101. // Default it to localhost:9300
  102. String[] nodeList = props.getProperty("es.hosts.list", DEFAULT_REMOTE_HOST).split(",");
  103. System.out.println("Elasticsearch Remote Hosts = " + props.getProperty("es.hosts.list", DEFAULT_REMOTE_HOST));
  104. TransportClient tClient = TransportClient.builder().settings(settings).build();
  105. for (String h : nodeList) {
  106. String[] nodes = h.split(":");
  107. try {
  108. tClient.addTransportAddress(new InetSocketTransportAddress(
  109. InetAddress.getByName(nodes[0]),
  110. Integer.parseInt(nodes[1])
  111. ));
  112. } catch (NumberFormatException e) {
  113. throw new IllegalArgumentException("Unable to parse port number.", e);
  114. } catch (UnknownHostException e) {
  115. throw new IllegalArgumentException("Unable to Identify host.", e);
  116. }
  117. }
  118. client = tClient;
  119. } else { // Start node only if transport client mode is disabled
  120. node = nodeBuilder().clusterName(clusterName).settings(settings).node();
  121. node.start();
  122. client = node.client();
  123. }
  124. final boolean exists =
  125. client.admin().indices()
  126. .exists(Requests.indicesExistsRequest(indexKey)).actionGet()
  127. .isExists();
  128. if (exists && newdb) {
  129. client.admin().indices().prepareDelete(indexKey).execute().actionGet();
  130. }
  131. if (!exists || newdb) {
  132. client.admin().indices().create(
  133. new CreateIndexRequest(indexKey)
  134. .settings(
  135. Settings.builder()
  136. .put("index.number_of_shards", numberOfShards)
  137. .put("index.number_of_replicas", numberOfReplicas)
  138. .put("index.mapping._id.indexed", true)
  139. )).actionGet();
  140. }
  141. client.admin().cluster().health(new ClusterHealthRequest().waitForGreenStatus()).actionGet();
  142. }
  143. private int parseIntegerProperty(Properties properties, String key, int defaultValue) {
  144. String value = properties.getProperty(key);
  145. return value == null ? defaultValue : Integer.parseInt(value);
  146. }
  147. @Override
  148. public void cleanup() throws DBException {
  149. if (!remoteMode) {
  150. if (!node.isClosed()) {
  151. client.close();
  152. node.close();
  153. }
  154. } else {
  155. client.close();
  156. }
  157. }
  158. /**
  159. * Insert a record in the database. Any field/value pairs in the specified
  160. * values HashMap will be written into the record with the specified record
  161. * key.
  162. *
  163. * @param table
  164. * The name of the table
  165. * @param key
  166. * The record key of the record to insert.
  167. * @param values
  168. * A HashMap of field/value pairs to insert in the record
  169. * @return Zero on success, a non-zero error code on error. See this class's
  170. * description for a discussion of error codes.
  171. */
  172. @Override
  173. public Status insert(String table, String key,
  174. HashMap<String, ByteIterator> values) {
  175. try {
  176. final XContentBuilder doc = jsonBuilder().startObject();
  177. for (Entry<String, String> entry : StringByteIterator.getStringMap(values).entrySet()) {
  178. doc.field(entry.getKey(), entry.getValue());
  179. }
  180. doc.endObject();
  181. client.prepareIndex(indexKey, table, key).setSource(doc).execute().actionGet();
  182. return Status.OK;
  183. } catch (Exception e) {
  184. e.printStackTrace();
  185. }
  186. return Status.ERROR;
  187. }
  188. /**
  189. * Delete a record from the database.
  190. *
  191. * @param table
  192. * The name of the table
  193. * @param key
  194. * The record key of the record to delete.
  195. * @return Zero on success, a non-zero error code on error. See this class's
  196. * description for a discussion of error codes.
  197. */
  198. @Override
  199. public Status delete(String table, String key) {
  200. try {
  201. client.prepareDelete(indexKey, table, key).execute().actionGet();
  202. return Status.OK;
  203. } catch (Exception e) {
  204. e.printStackTrace();
  205. }
  206. return Status.ERROR;
  207. }
  208. /**
  209. * Read a record from the database. Each field/value pair from the result will
  210. * be stored in a HashMap.
  211. *
  212. * @param table
  213. * The name of the table
  214. * @param key
  215. * The record key of the record to read.
  216. * @param fields
  217. * The list of fields to read, or null for all of them
  218. * @param result
  219. * A HashMap of field/value pairs for the result
  220. * @return Zero on success, a non-zero error code on error or "not found".
  221. */
  222. @Override
  223. public Status read(String table, String key, Set<String> fields,
  224. HashMap<String, ByteIterator> result) {
  225. try {
  226. final GetResponse response = client.prepareGet(indexKey, table, key).execute().actionGet();
  227. if (response.isExists()) {
  228. if (fields != null) {
  229. for (String field : fields) {
  230. result.put(field, new StringByteIterator(
  231. (String) response.getSource().get(field)));
  232. }
  233. } else {
  234. for (String field : response.getSource().keySet()) {
  235. result.put(field, new StringByteIterator(
  236. (String) response.getSource().get(field)));
  237. }
  238. }
  239. return Status.OK;
  240. }
  241. } catch (Exception e) {
  242. e.printStackTrace();
  243. }
  244. return Status.ERROR;
  245. }
  246. /**
  247. * Update a record in the database. Any field/value pairs in the specified
  248. * values HashMap will be written into the record with the specified record
  249. * key, overwriting any existing values with the same field name.
  250. *
  251. * @param table
  252. * The name of the table
  253. * @param key
  254. * The record key of the record to write.
  255. * @param values
  256. * A HashMap of field/value pairs to update in the record
  257. * @return Zero on success, a non-zero error code on error. See this class's
  258. * description for a discussion of error codes.
  259. */
  260. @Override
  261. public Status update(String table, String key,
  262. HashMap<String, ByteIterator> values) {
  263. try {
  264. final GetResponse response = client.prepareGet(indexKey, table, key).execute().actionGet();
  265. if (response.isExists()) {
  266. for (Entry<String, String> entry : StringByteIterator.getStringMap(values).entrySet()) {
  267. response.getSource().put(entry.getKey(), entry.getValue());
  268. }
  269. client.prepareIndex(indexKey, table, key).setSource(response.getSource()).execute().actionGet();
  270. return Status.OK;
  271. }
  272. } catch (Exception e) {
  273. e.printStackTrace();
  274. }
  275. return Status.ERROR;
  276. }
  277. /**
  278. * Perform a range scan for a set of records in the database. Each field/value
  279. * pair from the result will be stored in a HashMap.
  280. *
  281. * @param table
  282. * The name of the table
  283. * @param startkey
  284. * The record key of the first record to read.
  285. * @param recordcount
  286. * The number of records to read
  287. * @param fields
  288. * The list of fields to read, or null for all of them
  289. * @param result
  290. * A Vector of HashMaps, where each HashMap is a set field/value
  291. * pairs for one record
  292. * @return Zero on success, a non-zero error code on error. See this class's
  293. * description for a discussion of error codes.
  294. */
  295. @Override
  296. public Status scan(String table, String startkey, int recordcount,
  297. Set<String> fields, Vector<HashMap<String, ByteIterator>> result) {
  298. try {
  299. final RangeQueryBuilder rangeQuery = rangeQuery("_id").gte(startkey);
  300. final SearchResponse response = client.prepareSearch(indexKey)
  301. .setTypes(table)
  302. .setQuery(rangeQuery)
  303. .setSize(recordcount)
  304. .execute()
  305. .actionGet();
  306. HashMap<String, ByteIterator> entry;
  307. for (SearchHit hit : response.getHits()) {
  308. entry = new HashMap<>(fields.size());
  309. for (String field : fields) {
  310. entry.put(field, new StringByteIterator((String) hit.getSource().get(field)));
  311. }
  312. result.add(entry);
  313. }
  314. return Status.OK;
  315. } catch (Exception e) {
  316. e.printStackTrace();
  317. }
  318. return Status.ERROR;
  319. }
  320. }