Examples with SearchSourceBuilder org.graylog.shaded.elasticsearch7.org.elasticsearch.search.builder.SearchSourceBuilder used on opensource projects

Example 61 with SearchSourceBuilder

use of org.graylog.shaded.elasticsearch7.org.elasticsearch.search.builder.SearchSourceBuilder in project graylog2-server by Graylog2.

the class ESPivotTest method rowPivotsShouldBeNested.

@Test
public void rowPivotsShouldBeNested() {
    final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
    final ESPivotBucketSpecHandler<? extends BucketSpec, ? extends Aggregation> valuesBucketHandler = mock(ESValuesHandler.class);
    mockBucketSpecGeneratesComparableString(valuesBucketHandler);
    final ESPivotBucketSpecHandler<? extends BucketSpec, ? extends Aggregation> timeBucketHandler = mock(ESTimeHandler.class);
    mockBucketSpecGeneratesComparableString(timeBucketHandler);
    bucketHandlers.put(Values.NAME, valuesBucketHandler);
    bucketHandlers.put(Time.NAME, timeBucketHandler);
    when(queryContext.searchSourceBuilder(pivot)).thenReturn(searchSourceBuilder);
    when(queryContext.nextName()).thenReturn("time-agg", "values-agg");
    final Time time = Time.builder().field("timestamp").interval(AutoInterval.create()).build();
    final Values values = Values.builder().field("action").limit(10).build();
    when(pivot.rowGroups()).thenReturn(ImmutableList.of(time, values));
    this.esPivot.doGenerateQueryPart(job, query, pivot, queryContext);
    verify(valuesBucketHandler, times(1)).createAggregation(eq("values-agg"), eq(pivot), eq(values), eq(this.esPivot), eq(queryContext), eq(query));
    verify(timeBucketHandler, times(1)).createAggregation(eq("time-agg"), eq(pivot), eq(time), eq(this.esPivot), eq(queryContext), eq(query));
    final DocumentContext context = JsonPath.parse(searchSourceBuilder.toString());
    extractAggregation(context, "time-agg").isEqualTo("Time{type=time, field=timestamp, interval=AutoInterval{type=auto, scaling=1.0}}");
    extractAggregation(context, "time-agg.values-agg").isEqualTo("Values{type=values, field=action, limit=10}");
}

Example 62 with SearchSourceBuilder

use of org.graylog.shaded.elasticsearch7.org.elasticsearch.search.builder.SearchSourceBuilder in project graylog2-server by Graylog2.

the class ESPivotTest method mixedPivotsAndSeriesShouldBeNested.

@Test
public void mixedPivotsAndSeriesShouldBeNested() {
    final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
    final ESPivotBucketSpecHandler<? extends BucketSpec, ? extends Aggregation> valuesBucketHandler = mock(ESValuesHandler.class);
    mockBucketSpecGeneratesComparableString(valuesBucketHandler);
    final ESPivotBucketSpecHandler<? extends BucketSpec, ? extends Aggregation> timeBucketHandler = mock(ESTimeHandler.class);
    mockBucketSpecGeneratesComparableString(timeBucketHandler);
    final ESPivotSeriesSpecHandler<? extends SeriesSpec, ? extends Aggregation> countHandler = mock(ESCountHandler.class);
    mockSeriesSpecGeneratesComparableString(countHandler);
    bucketHandlers.put(Values.NAME, valuesBucketHandler);
    bucketHandlers.put(Time.NAME, timeBucketHandler);
    seriesHandlers.put(Count.NAME, countHandler);
    when(queryContext.searchSourceBuilder(pivot)).thenReturn(searchSourceBuilder);
    when(queryContext.nextName()).thenReturn("rowPivot1", "rowPivot2", "columnPivot1", "columnPivot2");
    final BucketSpec rowPivot1 = Time.builder().field("timestamp").interval(AutoInterval.create()).build();
    final BucketSpec rowPivot2 = Values.builder().field("http_method").limit(10).build();
    final BucketSpec columnPivot1 = Values.builder().field("controller").limit(10).build();
    final BucketSpec columnPivot2 = Values.builder().field("action").limit(10).build();
    final Count count = Count.builder().build();
    when(pivot.rowGroups()).thenReturn(ImmutableList.of(rowPivot1, rowPivot2));
    when(pivot.columnGroups()).thenReturn(ImmutableList.of(columnPivot1, columnPivot2));
    when(pivot.series()).thenReturn(Collections.singletonList(count));
    when(pivot.rollup()).thenReturn(false);
    when(queryContext.seriesName(any(), any())).thenCallRealMethod();
    this.esPivot.doGenerateQueryPart(job, query, pivot, queryContext);
    verify(timeBucketHandler).createAggregation(eq("rowPivot1"), eq(pivot), eq(rowPivot1), eq(this.esPivot), eq(queryContext), eq(query));
    verify(valuesBucketHandler).createAggregation(eq("rowPivot2"), eq(pivot), eq(rowPivot2), eq(this.esPivot), eq(queryContext), eq(query));
    verify(valuesBucketHandler).createAggregation(eq("columnPivot1"), eq(pivot), eq(columnPivot1), eq(this.esPivot), eq(queryContext), eq(query));
    verify(valuesBucketHandler).createAggregation(eq("columnPivot2"), eq(pivot), eq(columnPivot2), eq(this.esPivot), eq(queryContext), eq(query));
    final DocumentContext context = JsonPath.parse(searchSourceBuilder.toString());
    extractAggregation(context, "rowPivot1").isEqualTo("Time{type=time, field=timestamp, interval=AutoInterval{type=auto, scaling=1.0}}");
    extractAggregation(context, "rowPivot1.rowPivot2").isEqualTo("Values{type=values, field=http_method, limit=10}");
    extractAggregation(context, "rowPivot1.rowPivot2.columnPivot1").isEqualTo("Values{type=values, field=controller, limit=10}");
    extractAggregation(context, "rowPivot1.rowPivot2.columnPivot1.columnPivot2").isEqualTo("Values{type=values, field=action, limit=10}");
    extractAggregation(context, "rowPivot1.rowPivot2.dummypivot-series-count()").isEqualTo("Count{type=count, id=count(), field=null}");
    extractAggregation(context, "rowPivot1.rowPivot2.columnPivot1.columnPivot2.dummypivot-series-count()").isEqualTo("Count{type=count, id=count(), field=null}");
}

Example 63 with SearchSourceBuilder

use of org.graylog.shaded.elasticsearch7.org.elasticsearch.search.builder.SearchSourceBuilder in project graylog2-server by Graylog2.

the class ElasticsearchBackendErrorHandlingTest method setUp.

@Before
public void setUp() throws Exception {
    final FieldTypesLookup fieldTypesLookup = mock(FieldTypesLookup.class);
    this.backend = new ElasticsearchBackend(ImmutableMap.of("dummy", () -> mock(DummyHandler.class)), jestClient, indexLookup, new QueryStringDecorators(Collections.emptySet()), (elasticsearchBackend, ssb, job, query) -> new ESGeneratedQueryContext(elasticsearchBackend, ssb, job, query, fieldTypesLookup), false, objectMapper);
    when(indexLookup.indexNamesForStreamsInTimeRange(any(), any())).thenReturn(Collections.emptySet());
    final SearchType searchType1 = mock(SearchType.class);
    when(searchType1.id()).thenReturn("deadbeef");
    when(searchType1.type()).thenReturn("dummy");
    final SearchType searchType2 = mock(SearchType.class);
    when(searchType2.id()).thenReturn("cafeaffe");
    when(searchType2.type()).thenReturn("dummy");
    final Set<SearchType> searchTypes = ImmutableSet.of(searchType1, searchType2);
    this.query = Query.builder().id("query1").timerange(RelativeRange.create(300)).query(ElasticsearchQueryString.of("*")).searchTypes(searchTypes).build();
    final Search search = Search.builder().id("search1").queries(ImmutableSet.of(query)).build();
    this.searchJob = new SearchJob("job1", search, "admin");
    this.queryContext = new ESGeneratedQueryContext(this.backend, new SearchSourceBuilder(), searchJob, query, mock(FieldTypesLookup.class));
    searchTypes.forEach(queryContext::searchSourceBuilder);
    when(jestClient.execute(any())).thenReturn(result);
}

Example 64 with SearchSourceBuilder

use of org.graylog.shaded.elasticsearch7.org.elasticsearch.search.builder.SearchSourceBuilder in project graylog2-server by Graylog2.

the class CountsAdapterES7 method totalCount.

@Override
public long totalCount(List<String> indices) {
    final SearchSourceBuilder query = new SearchSourceBuilder().query(QueryBuilders.matchAllQuery()).size(0).trackTotalHits(true);
    final SearchRequest searchRequest = new SearchRequest(indices.toArray(new String[0])).source(query);
    final SearchResponse result = client.search(searchRequest, "Fetching message count failed for indices ");
    return result.getHits().getTotalHits().value;
}

Example 65 with SearchSourceBuilder

use of org.graylog.shaded.elasticsearch7.org.elasticsearch.search.builder.SearchSourceBuilder in project graylog2-server by Graylog2.

the class QuerySuggestionsES7 method suggest.

@Override
public SuggestionResponse suggest(SuggestionRequest req) {
    final Set<String> affectedIndices = indexLookup.indexNamesForStreamsInTimeRange(req.streams(), req.timerange());
    final TermSuggestionBuilder suggestionBuilder = SuggestBuilders.termSuggestion(req.field()).text(req.input()).size(req.size());
    final SearchSourceBuilder search = new SearchSourceBuilder().query(QueryBuilders.prefixQuery(req.field(), req.input())).size(0).aggregation(AggregationBuilders.terms("fieldvalues").field(req.field()).size(req.size())).suggest(new SuggestBuilder().addSuggestion("corrections", suggestionBuilder));
    try {
        final SearchResponse result = client.singleSearch(new SearchRequest(affectedIndices.toArray(new String[] {})).source(search), "Failed to execute aggregation");
        final ParsedStringTerms fieldValues = result.getAggregations().get("fieldvalues");
        final List<SuggestionEntry> entries = fieldValues.getBuckets().stream().map(b -> new SuggestionEntry(b.getKeyAsString(), b.getDocCount())).collect(Collectors.toList());
        if (!entries.isEmpty()) {
            return SuggestionResponse.forSuggestions(req.field(), req.input(), entries, fieldValues.getSumOfOtherDocCounts());
        } else {
            TermSuggestion suggestion = result.getSuggest().getSuggestion("corrections");
            final List<SuggestionEntry> corrections = suggestion.getEntries().stream().flatMap(e -> e.getOptions().stream()).map(o -> new SuggestionEntry(o.getText().string(), o.getFreq())).collect(Collectors.toList());
            return SuggestionResponse.forSuggestions(req.field(), req.input(), corrections, null);
        }
    } catch (org.graylog.shaded.elasticsearch7.org.elasticsearch.ElasticsearchException exception) {
        final SuggestionError err = tryResponseException(exception).orElseGet(() -> parseException(exception));
        return SuggestionResponse.forError(req.field(), req.input(), err);
    }
}