Java Code Examples for org.elasticsearch.client.Request#addParameters()

@Override
public boolean advance() throws IOException {
  if (batchIterator.hasNext()) {
    current = batchIterator.next();
    return true;
  } else {
    String requestBody =
        String.format(
            "{\"scroll\" : \"%s\",\"scroll_id\" : \"%s\"}",
            source.spec.getScrollKeepalive(), scrollId);
    HttpEntity scrollEntity = new NStringEntity(requestBody, ContentType.APPLICATION_JSON);
    Request request = new Request("GET", "/_search/scroll");
    request.addParameters(Collections.emptyMap());
    request.setEntity(scrollEntity);
    Response response = restClient.performRequest(request);
    JsonNode searchResult = parseResponse(response.getEntity());
    updateScrollId(searchResult);
    return readNextBatchAndReturnFirstDocument(searchResult);
  }
}
 

@Override
public void close() throws IOException {
  // remove the scroll
  String requestBody = String.format("{\"scroll_id\" : [\"%s\"]}", scrollId);
  HttpEntity entity = new NStringEntity(requestBody, ContentType.APPLICATION_JSON);
  try {
    Request request = new Request("DELETE", "/_search/scroll");
    request.addParameters(Collections.emptyMap());
    request.setEntity(entity);
    restClient.performRequest(request);
  } finally {
    if (restClient != null) {
      restClient.close();
    }
  }
}
 

private static JsonNode getStats(
    ConnectionConfiguration connectionConfiguration, boolean shardLevel) throws IOException {
  HashMap<String, String> params = new HashMap<>();
  if (shardLevel) {
    params.put("level", "shards");
  }
  String endpoint = String.format("/%s/_stats", connectionConfiguration.getIndex());
  try (RestClient restClient = connectionConfiguration.createClient()) {
    Request request = new Request("GET", endpoint);
    request.addParameters(params);
    return parseResponse(restClient.performRequest(request).getEntity());
  }
}
 

@Override
public boolean start() throws IOException {
  restClient = source.spec.getConnectionConfiguration().createClient();

  String query = source.spec.getQuery() != null ? source.spec.getQuery().get() : null;
  if (query == null) {
    query = "{\"query\": { \"match_all\": {} }}";
  }
  if ((source.backendVersion >= 5) && source.numSlices != null && source.numSlices > 1) {
    // if there is more than one slice, add the slice to the user query
    String sliceQuery =
        String.format("\"slice\": {\"id\": %s,\"max\": %s}", source.sliceId, source.numSlices);
    query = query.replaceFirst("\\{", "{" + sliceQuery + ",");
  }
  String endPoint =
      String.format(
          "/%s/%s/_search",
          source.spec.getConnectionConfiguration().getIndex(),
          source.spec.getConnectionConfiguration().getType());
  Map<String, String> params = new HashMap<>();
  params.put("scroll", source.spec.getScrollKeepalive());
  if (source.backendVersion == 2) {
    params.put("size", String.valueOf(source.spec.getBatchSize()));
    if (source.shardPreference != null) {
      params.put("preference", "_shards:" + source.shardPreference);
    }
  }
  HttpEntity queryEntity = new NStringEntity(query, ContentType.APPLICATION_JSON);
  Request request = new Request("GET", endPoint);
  request.addParameters(params);
  request.setEntity(queryEntity);
  Response response = restClient.performRequest(request);
  JsonNode searchResult = parseResponse(response.getEntity());
  updateScrollId(searchResult);
  return readNextBatchAndReturnFirstDocument(searchResult);
}
 

private static void deleteIndex(RestClient restClient, String index) throws IOException {
  try {
    closeIndex(restClient, index);
    Request request = new Request("DELETE", String.format("/%s", index));
    request.addParameters(Collections.singletonMap("refresh", "wait_for"));
    restClient.performRequest(request);
  } catch (IOException e) {
    // it is fine to ignore this expression as deleteIndex occurs in @before,
    // so when the first tests is run, the index does not exist yet
    if (!e.getMessage().contains("index_not_found_exception")) {
      throw e;
    }
  }
}
 

/**
 * Synchronously deletes the target if it exists and then (re)creates it as a copy of source
 * synchronously.
 */
static void copyIndex(RestClient restClient, String source, String target) throws IOException {
  deleteIndex(restClient, target);
  HttpEntity entity =
      new NStringEntity(
          String.format(
              "{\"source\" : { \"index\" : \"%s\" }, \"dest\" : { \"index\" : \"%s\" } }",
              source, target),
          ContentType.APPLICATION_JSON);
  Request request = new Request("POST", "/_reindex");
  request.addParameters(Collections.singletonMap("refresh", "wait_for"));
  request.setEntity(entity);
  restClient.performRequest(request);
}
 

/** Inserts the given number of test documents into Elasticsearch. */
static void insertTestDocuments(
    ConnectionConfiguration connectionConfiguration, long numDocs, RestClient restClient)
    throws IOException {
  List<String> data =
      ElasticsearchIOTestUtils.createDocuments(
          numDocs, ElasticsearchIOTestUtils.InjectionMode.DO_NOT_INJECT_INVALID_DOCS);
  StringBuilder bulkRequest = new StringBuilder();
  int i = 0;
  for (String document : data) {
    bulkRequest.append(
        String.format(
            "{ \"index\" : { \"_index\" : \"%s\", \"_type\" : \"%s\", \"_id\" : \"%s\" } }%n%s%n",
            connectionConfiguration.getIndex(),
            connectionConfiguration.getType(),
            i++,
            document));
  }
  String endPoint =
      String.format(
          "/%s/%s/_bulk", connectionConfiguration.getIndex(), connectionConfiguration.getType());
  HttpEntity requestBody =
      new NStringEntity(bulkRequest.toString(), ContentType.APPLICATION_JSON);
  Request request = new Request("POST", endPoint);
  request.addParameters(Collections.singletonMap("refresh", "wait_for"));
  request.setEntity(requestBody);
  Response response = restClient.performRequest(request);
  ElasticsearchIO.checkForErrors(
      response.getEntity(), ElasticsearchIO.getBackendVersion(connectionConfiguration), false);
}
 

/**
 * Executes a match query for given field/value and returns the count of results.
 *
 * @param connectionConfiguration Specifies the index and type
 * @param restClient To use to execute the call
 * @param field The field to query
 * @param value The value to match
 * @return The count of documents in the search result
 * @throws IOException On error communicating with Elasticsearch
 */
static int countByMatch(
    ConnectionConfiguration connectionConfiguration,
    RestClient restClient,
    String field,
    String value)
    throws IOException {
  String requestBody =
      "{\n"
          + "  \"query\" : {\"match\": {\n"
          + "    \""
          + field
          + "\": \""
          + value
          + "\"\n"
          + "  }}\n"
          + "}\n";
  String endPoint =
      String.format(
          "/%s/%s/_search",
          connectionConfiguration.getIndex(), connectionConfiguration.getType());
  HttpEntity httpEntity = new NStringEntity(requestBody, ContentType.APPLICATION_JSON);

  Request request = new Request("GET", endPoint);
  request.addParameters(Collections.emptyMap());
  request.setEntity(httpEntity);
  Response response = restClient.performRequest(request);
  JsonNode searchResult = parseResponse(response.getEntity());
  if (getBackendVersion(connectionConfiguration) >= 7) {
    return searchResult.path("hits").path("total").path("value").asInt();
  } else {
    return searchResult.path("hits").path("total").asInt();
  }
}
 

/** Test that the default predicate correctly parses chosen error code. */
void testDefaultRetryPredicate(RestClient restClient) throws IOException {

  HttpEntity entity1 = new NStringEntity(BAD_REQUEST, ContentType.APPLICATION_JSON);
  Request request = new Request("POST", "/_bulk");
  request.addParameters(Collections.emptyMap());
  request.setEntity(entity1);
  Response response1 = restClient.performRequest(request);
  assertTrue(CUSTOM_RETRY_PREDICATE.test(response1.getEntity()));

  HttpEntity entity2 = new NStringEntity(OK_REQUEST, ContentType.APPLICATION_JSON);
  request.setEntity(entity2);
  Response response2 = restClient.performRequest(request);
  assertFalse(DEFAULT_RETRY_PREDICATE.test(response2.getEntity()));
}