io.searchbox.core.SearchScroll Java Examples

@Test
public void testIdsInIndexFunction()
{
    JestResult jestResult = mock(JestResult.class);
    SearchResult searchResult = mock(SearchResult.class);
    List<String> idList = Arrays.asList("{id:1}");
    List<String> emptyList = new ArrayList<>();
    JsonObject jsonObject = new JsonObject();
    jsonObject.addProperty("_scroll_id", "100");
    when(jestClientHelper.execute(any(Search.class))).thenReturn(searchResult);
    when(searchResult.getSourceAsStringList()).thenReturn(idList);
    when(searchResult.getJsonObject()).thenReturn(jsonObject);
    when(jestClientHelper.execute(any(SearchScroll.class))).thenReturn(jestResult);
    when(jestResult.getSourceAsStringList()).thenReturn(emptyList);
    indexFunctionsDao.getIdsInIndex("INDEX_NAME", "DOCUMENT_TYPE");
    verify(jestClientHelper).execute(any(Search.class));
    verify(searchResult, times(2)).getSourceAsStringList();
    verify(searchResult).getJsonObject();
    verify(jestClientHelper).execute(any(SearchScroll.class));
    verify(jestResult).getSourceAsStringList();
}
 

@SuppressWarnings({ "unchecked", "rawtypes" })
private void fetch() throws StorageException {
    try {
        Builder builder = new SearchScroll.Builder(scrollId, "1m");
        SearchScroll scroll = new SearchScroll(builder) {
            @Override
            public JestResult createNewElasticSearchResult(String responseBody, int statusCode,
                    String reasonPhrase, Gson gson) {
                return createNewElasticSearchResult(new SearchResult(gson), responseBody, statusCode, reasonPhrase, gson);
            }
        };
        SearchResult response = (SearchResult) esClient.execute(scroll);
        if (!response.isSucceeded()) {
            throw new StorageException("Scrolled fetch failed " + response.getErrorMessage());
        }
        this.hits = (List) response.getHits(Map.class);
    } catch (IOException e) {
        throw new StorageException(e);
    }
}
 

/**
 * The ids in index function will take as arguments the index name and the document type and will return a list of all the ids in the index.
 */
@Override
public final List<String> getIdsInIndex(String indexName, String documentType)
{
    // Create an array list for storing the ids
    List<String> idList = new ArrayList<>();

    SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
    searchSourceBuilder.query(QueryBuilders.matchAllQuery());
    // Create a search request and set the scroll time and scroll size
    final SearchRequestBuilder searchRequestBuilder = new SearchRequestBuilder(new ElasticsearchClientImpl(), SearchAction.INSTANCE);
    searchRequestBuilder.setIndices(indexName).setTypes(documentType).setScroll(new TimeValue(ELASTIC_SEARCH_SCROLL_KEEP_ALIVE_TIME))
        .setSize(ELASTIC_SEARCH_SCROLL_PAGE_SIZE).setSource(searchSourceBuilder);

    // Retrieve the search response
    final Search.Builder searchBuilder = new Search.Builder(searchRequestBuilder.toString()).addIndex(indexName);

    searchBuilder.setParameter(Parameters.SIZE, ELASTIC_SEARCH_SCROLL_PAGE_SIZE);
    searchBuilder.setParameter(Parameters.SCROLL, new TimeValue(ELASTIC_SEARCH_SCROLL_KEEP_ALIVE_TIME).toString());

    JestResult jestResult = jestClientHelper.execute(searchBuilder.build());

    // While there are hits available, page through the results and add them to the id list
    while (jestResult.getSourceAsStringList().size() != 0)
    {
        for (String jsonString : jestResult.getSourceAsStringList())
        {
            JsonElement root = new JsonParser().parse(jsonString);
            idList.add(root.getAsJsonObject().get("id").getAsString());
        }
        String scrollId = jestResult.getJsonObject().get("_scroll_id").getAsString();
        SearchScroll scroll = new SearchScroll.Builder(scrollId, new TimeValue(ELASTIC_SEARCH_SCROLL_KEEP_ALIVE_TIME).toString()).build();
        jestResult = jestClientHelper.execute(scroll);

    }
    return idList;
}
 

@Override
public void initialize(UimaContext context) throws ResourceInitializationException {
	super.initialize(context);

	logger = context.getLogger();
	// init hoover connection
	client = hooverResource.getClient();
	esIndex = hooverResource.getIndex();

	// query hoover's elasticsearch index
	Search search = new Search.Builder(
			"{\"query\": {\"match_all\" : {}}, \"_source\" : false, \"size\" : " + PARAM_SCROLL_SIZE + "}")
					.addIndex(hooverResource.getIndex()).addType(HooverResource.HOOVER_DOCUMENT_TYPE)
					.setParameter(Parameters.SCROLL, PARAM_SCROLL_TIME).build();

	try {

		// run JEST request
		JestResult result = client.execute(search);

		totalIdList = new ArrayList<String>();

		JsonArray hits = result.getJsonObject().getAsJsonObject("hits").getAsJsonArray("hits");
		Integer total = result.getJsonObject().getAsJsonObject("hits").get("total").getAsInt();

		int nHits = hits.size();

		logger.log(Level.INFO, "Hits first result: " + nHits);
		logger.log(Level.INFO, "Hits total: " + total);

		totalIdList.addAll(hooverResource.getIds(hits));

		String scrollId = result.getJsonObject().get("_scroll_id").getAsString();

		// run scroll request to collect all Ids
		int i = 0;
		while (nHits > 0) {
			SearchScroll scroll = new SearchScroll.Builder(scrollId, PARAM_SCROLL_TIME).build();

			result = client.execute(scroll);

			hits = result.getJsonObject().getAsJsonObject("hits").getAsJsonArray("hits");
			nHits = hits.size();
			logger.log(Level.INFO, "Hits " + ++i + " result: " + nHits);
			totalIdList.addAll(hooverResource.getIds(hits));
			scrollId = result.getJsonObject().getAsJsonPrimitive("_scroll_id").getAsString();
		}

		if (maxRecords > 0 && maxRecords < totalIdList.size()) {
			totalIdList = new ArrayList<String>(totalIdList.subList(0, maxRecords));
		}

		totalRecords = totalIdList.size();
		logger.log(Level.INFO, "Found " + totalRecords + " ids in index " + esIndex);

	} catch (IOException e) {
		throw new ResourceInitializationException(e);
	}

}