From 36bd2064eb18b4607ad64b379ffe1d4c1abda069 Mon Sep 17 00:00:00 2001 From: Shivansh Anand Srivastava Date: Fri, 7 Jun 2024 15:37:30 +0530 Subject: [PATCH] chore | Add support for providing filters in spanJoiner (#181) * Add support for providing filters in spanJoiner * Resolve PR reviews --- .../span/joiner/DefaultSpanJoinerBuilder.java | 37 +++++++++++++------ .../core/graphql/span/joiner/SpanJoiner.java | 28 ++++++++++++-- .../span/joiner/SpanJoinerBuilderTest.java | 16 +++++--- 3 files changed, 60 insertions(+), 21 deletions(-) diff --git a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/DefaultSpanJoinerBuilder.java b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/DefaultSpanJoinerBuilder.java index 46bcd6c3..4cd98626 100644 --- a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/DefaultSpanJoinerBuilder.java +++ b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/DefaultSpanJoinerBuilder.java @@ -18,9 +18,9 @@ import java.util.Map; import java.util.Map.Entry; import java.util.Optional; -import java.util.Set; import java.util.function.Function; import java.util.stream.Collectors; +import java.util.stream.Stream; import javax.inject.Inject; import lombok.AllArgsConstructor; import lombok.Value; @@ -86,26 +86,32 @@ private class DefaultSpanJoiner implements SpanJoiner { @Override public Single> joinSpan( - Collection joinSources, SpanIdGetter spanIdGetter) { + Collection joinSources, + SpanIdGetter spanIdGetter, + Collection filterArguments) { Function>> idsGetter = source -> spanIdGetter.getSpanId(source).map(List::of); - return this.joinSpans(joinSources, idsGetter, SPAN_KEY).map(this::reduceMap); + return this.joinSpans(joinSources, idsGetter, filterArguments, SPAN_KEY).map(this::reduceMap); } @Override public Single> joinSpans( - Collection joinSources, MultipleSpanIdGetter multipleSpanIdGetter) { - return this.joinSpans(joinSources, multipleSpanIdGetter::getSpanIds, SPANS_KEY); + Collection joinSources, + MultipleSpanIdGetter multipleSpanIdGetter, + Collection filterArguments) { + return this.joinSpans( + joinSources, multipleSpanIdGetter::getSpanIds, filterArguments, SPANS_KEY); } private Single> joinSpans( Collection joinSources, Function>> idsGetter, + Collection filterArguments, String joinSpanKey) { return this.buildSourceToIdsMap(joinSources, idsGetter) .flatMap( sourceToSpanIdsMap -> - this.buildSpanRequest(sourceToSpanIdsMap, joinSpanKey) + this.buildSpanRequest(sourceToSpanIdsMap, joinSpanKey, filterArguments) .flatMap(spanDao::getSpans) .map(this::buildSpanIdToSpanMap) .map( @@ -152,15 +158,16 @@ private Map buildSpanIdToSpanMap(SpanResultSet resultSet) { } private Single buildSpanRequest( - ListMultimap sourceToSpanIdsMultimap, String joinSpanKey) { + ListMultimap sourceToSpanIdsMultimap, + String joinSpanKey, + Collection filterArguments) { Collection spanIds = sourceToSpanIdsMultimap.values().stream() .distinct() .collect(Collectors.toUnmodifiableList()); List selectedFields = getSelections(joinSpanKey); - return buildSpanIdsFilter(context, spanIds) - .flatMap( - filterArguments -> buildSpanRequest(spanIds.size(), filterArguments, selectedFields)); + return buildSpanIdsFilter(context, spanIds, filterArguments) + .flatMap(filters -> buildSpanRequest(spanIds.size(), filters, selectedFields)); } private Single buildSpanRequest( @@ -182,8 +189,14 @@ private Single buildSpanRequest( } private Single>> buildSpanIdsFilter( - GraphQlRequestContext context, Collection spanIds) { - return filterRequestBuilder.build(context, SPAN, Set.of(new SpanIdFilter(spanIds))); + GraphQlRequestContext context, + Collection spanIds, + Collection filterArguments) { + return filterRequestBuilder.build( + context, + SPAN, + Stream.concat(filterArguments.stream(), Stream.of(new SpanIdFilter(spanIds))) + .collect(Collectors.toUnmodifiableList())); } } diff --git a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoiner.java b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoiner.java index 9997257a..3cc8c8c9 100644 --- a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoiner.java +++ b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoiner.java @@ -7,6 +7,7 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterArgument; import org.hypertrace.core.graphql.span.schema.Span; public interface SpanJoiner { @@ -16,21 +17,40 @@ public interface SpanJoiner { new SpanJoiner() { @Override public Single> joinSpan( - Collection joinSources, SpanIdGetter spanIdGetter) { + Collection joinSources, + SpanIdGetter spanIdGetter, + Collection filterArguments) { return Single.just(Collections.emptyMap()); } @Override public Single> joinSpans( - Collection joinSources, MultipleSpanIdGetter multipleSpanIdGetter) { + Collection joinSources, + MultipleSpanIdGetter multipleSpanIdGetter, + Collection filterArguments) { return Single.just(ArrayListMultimap.create()); } }; - Single> joinSpan(Collection joinSources, SpanIdGetter spanIdGetter); + default Single> joinSpan( + Collection joinSources, SpanIdGetter spanIdGetter) { + return joinSpan(joinSources, spanIdGetter, Collections.emptyList()); + } + + default Single> joinSpans( + Collection joinSources, MultipleSpanIdGetter multipleSpanIdGetter) { + return joinSpans(joinSources, multipleSpanIdGetter, Collections.emptyList()); + } + + Single> joinSpan( + Collection joinSources, + SpanIdGetter spanIdGetter, + Collection filterArguments); Single> joinSpans( - Collection joinSources, MultipleSpanIdGetter multipleSpanIdGetter); + Collection joinSources, + MultipleSpanIdGetter multipleSpanIdGetter, + Collection filterArguments); @FunctionalInterface interface SpanIdGetter { diff --git a/hypertrace-core-graphql-span-schema/src/test/java/org/hypertrace/core/graphql/span/joiner/SpanJoinerBuilderTest.java b/hypertrace-core-graphql-span-schema/src/test/java/org/hypertrace/core/graphql/span/joiner/SpanJoinerBuilderTest.java index eae7b19a..7242cf9a 100644 --- a/hypertrace-core-graphql-span-schema/src/test/java/org/hypertrace/core/graphql/span/joiner/SpanJoinerBuilderTest.java +++ b/hypertrace-core-graphql-span-schema/src/test/java/org/hypertrace/core/graphql/span/joiner/SpanJoinerBuilderTest.java @@ -5,7 +5,7 @@ import static org.hypertrace.core.graphql.atttributes.scopes.HypertraceCoreAttributeScopeString.SPAN; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anySet; +import static org.mockito.ArgumentMatchers.anyList; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -15,6 +15,7 @@ import graphql.schema.DataFetchingFieldSelectionSet; import graphql.schema.SelectedField; import io.reactivex.rxjava3.core.Single; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Optional; @@ -88,7 +89,7 @@ void fetchSpans() { mockSelectionSet, SelectionQuery.builder().selectionPath(List.of("pathToSpan", "span")).build())) .thenReturn(Stream.of(mock(SelectedField.class), mock(SelectedField.class))); - when(mockFilterRequestBuilder.build(eq(mockRequestContext), eq(SPAN), anySet())) + when(mockFilterRequestBuilder.build(eq(mockRequestContext), eq(SPAN), anyList())) .thenReturn(Single.just(List.of(mockFilter))); when(mockResultSetRequestBuilder.build( @@ -112,7 +113,10 @@ void fetchSpans() { List.of("pathToSpan")) .blockingGet(); assertEquals( - expected, joiner.joinSpan(joinSources, new TestJoinSourceIdGetter()).blockingGet()); + expected, + joiner + .joinSpan(joinSources, new TestJoinSourceIdGetter(), Collections.emptyList()) + .blockingGet()); } @Test @@ -129,7 +133,7 @@ void fetchMultipleSpans() { mockSelectionSet, SelectionQuery.builder().selectionPath(List.of("pathToSpans", "spans")).build())) .thenReturn(Stream.of(mock(SelectedField.class), mock(SelectedField.class))); - when(mockFilterRequestBuilder.build(eq(mockRequestContext), eq(SPAN), anySet())) + when(mockFilterRequestBuilder.build(eq(mockRequestContext), eq(SPAN), anyList())) .thenReturn(Single.just(List.of(mockFilter))); when(mockResultSetRequestBuilder.build( @@ -155,7 +159,9 @@ void fetchMultipleSpans() { .blockingGet(); assertEquals( expected, - joiner.joinSpans(joinSources, new TestMultipleJoinSourceIdGetter()).blockingGet()); + joiner + .joinSpans(joinSources, new TestMultipleJoinSourceIdGetter(), Collections.emptyList()) + .blockingGet()); } private void mockResult(List spans) {