From 37ca0a9b31bb0525085739cd3985204ced7b439c Mon Sep 17 00:00:00 2001 From: Shivansh Anand Srivastava Date: Fri, 24 May 2024 18:18:13 +0530 Subject: [PATCH] Add SpanJoiner based on list of spanId --- .../core/graphql/span/SpanSchemaModule.java | 4 +- .../span/joiner/DefaultSpanJoinerBuilder.java | 180 ------------------ ...panJoinerModule.java => JoinerModule.java} | 7 +- .../span/joiner/SourceToSpansProvider.java | 143 ++++++++++++++ .../DefaultMultipleSpanJoinerBuilder.java | 87 +++++++++ .../multiple/span/MultipleSpanJoin.java | 14 ++ .../multiple/span/MultipleSpanJoiner.java | 29 +++ .../span/MultipleSpanJoinerBuilder.java | 15 ++ .../single/span/DefaultSpanJoinerBuilder.java | 95 +++++++++ .../joiner/{ => single/span}/SpanJoin.java | 2 +- .../joiner/{ => single/span}/SpanJoiner.java | 2 +- .../{ => single/span}/SpanJoinerBuilder.java | 2 +- .../span/MultipleSpanJoinerBuilderTest.java | 165 ++++++++++++++++ .../span}/SpanJoinerBuilderTest.java | 17 +- 14 files changed, 567 insertions(+), 195 deletions(-) delete mode 100644 hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/DefaultSpanJoinerBuilder.java rename hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/{SpanJoinerModule.java => JoinerModule.java} (58%) create mode 100644 hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SourceToSpansProvider.java create mode 100644 hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/DefaultMultipleSpanJoinerBuilder.java create mode 100644 hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoin.java create mode 100644 hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoiner.java create mode 100644 hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoinerBuilder.java create mode 100644 hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/single/span/DefaultSpanJoinerBuilder.java rename hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/{ => single/span}/SpanJoin.java (82%) rename hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/{ => single/span}/SpanJoiner.java (91%) rename hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/{ => single/span}/SpanJoinerBuilder.java (88%) create mode 100644 hypertrace-core-graphql-span-schema/src/test/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoinerBuilderTest.java rename hypertrace-core-graphql-span-schema/src/test/java/org/hypertrace/core/graphql/span/joiner/{ => single/span}/SpanJoinerBuilderTest.java (91%) diff --git a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/SpanSchemaModule.java b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/SpanSchemaModule.java index f677143c..81895bef 100644 --- a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/SpanSchemaModule.java +++ b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/SpanSchemaModule.java @@ -4,7 +4,7 @@ import com.google.inject.multibindings.Multibinder; import org.hypertrace.core.graphql.common.request.ResultSetRequestBuilder; import org.hypertrace.core.graphql.span.dao.SpanDaoModule; -import org.hypertrace.core.graphql.span.joiner.SpanJoinerModule; +import org.hypertrace.core.graphql.span.joiner.JoinerModule; import org.hypertrace.core.graphql.span.request.SpanRequestModule; import org.hypertrace.core.graphql.spi.schema.GraphQlSchemaFragment; @@ -18,6 +18,6 @@ protected void configure() { requireBinding(ResultSetRequestBuilder.class); install(new SpanDaoModule()); install(new SpanRequestModule()); - install(new SpanJoinerModule()); + install(new JoinerModule()); } } 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 deleted file mode 100644 index bf4ea147..00000000 --- a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/DefaultSpanJoinerBuilder.java +++ /dev/null @@ -1,180 +0,0 @@ -package org.hypertrace.core.graphql.span.joiner; - -import static com.google.common.collect.ImmutableList.copyOf; -import static com.google.common.collect.Iterables.concat; -import static org.hypertrace.core.graphql.atttributes.scopes.HypertraceCoreAttributeScopeString.SPAN; -import static org.hypertrace.core.graphql.span.joiner.SpanJoin.SPAN_KEY; - -import graphql.schema.DataFetchingFieldSelectionSet; -import graphql.schema.SelectedField; -import io.reactivex.rxjava3.core.Observable; -import io.reactivex.rxjava3.core.Single; -import java.util.Collection; -import java.util.Collections; -import java.util.List; -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 javax.inject.Inject; -import lombok.AllArgsConstructor; -import lombok.Value; -import lombok.experimental.Accessors; -import org.hypertrace.core.graphql.common.request.AttributeAssociation; -import org.hypertrace.core.graphql.common.request.AttributeRequest; -import org.hypertrace.core.graphql.common.request.FilterRequestBuilder; -import org.hypertrace.core.graphql.common.request.ResultSetRequest; -import org.hypertrace.core.graphql.common.request.ResultSetRequestBuilder; -import org.hypertrace.core.graphql.common.schema.arguments.TimeRangeArgument; -import org.hypertrace.core.graphql.common.schema.attributes.AttributeScope; -import org.hypertrace.core.graphql.common.schema.attributes.arguments.AttributeExpression; -import org.hypertrace.core.graphql.common.schema.id.Identifiable; -import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterArgument; -import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterOperatorType; -import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterType; -import org.hypertrace.core.graphql.common.schema.results.arguments.order.OrderArgument; -import org.hypertrace.core.graphql.context.GraphQlRequestContext; -import org.hypertrace.core.graphql.span.dao.SpanDao; -import org.hypertrace.core.graphql.span.request.SpanRequest; -import org.hypertrace.core.graphql.span.schema.Span; -import org.hypertrace.core.graphql.span.schema.SpanResultSet; -import org.hypertrace.core.graphql.utils.schema.GraphQlSelectionFinder; -import org.hypertrace.core.graphql.utils.schema.SelectionQuery; - -public class DefaultSpanJoinerBuilder implements SpanJoinerBuilder { - - private static final int ZERO_OFFSET = 0; - - private final SpanDao spanDao; - private final GraphQlSelectionFinder selectionFinder; - private final ResultSetRequestBuilder resultSetRequestBuilder; - private final FilterRequestBuilder filterRequestBuilder; - - @Inject - DefaultSpanJoinerBuilder( - SpanDao spanDao, - GraphQlSelectionFinder selectionFinder, - ResultSetRequestBuilder resultSetRequestBuilder, - FilterRequestBuilder filterRequestBuilder) { - this.spanDao = spanDao; - this.selectionFinder = selectionFinder; - this.resultSetRequestBuilder = resultSetRequestBuilder; - this.filterRequestBuilder = filterRequestBuilder; - } - - @Override - public Single build( - GraphQlRequestContext context, - TimeRangeArgument timeRange, - DataFetchingFieldSelectionSet selectionSet, - List pathToSpanJoin) { - return Single.just( - new DefaultSpanJoiner( - context, timeRange, this.getSelections(selectionSet, pathToSpanJoin))); - } - - private List getSelections( - DataFetchingFieldSelectionSet selectionSet, List pathToSpanJoin) { - List fullPath = copyOf(concat(pathToSpanJoin, List.of(SPAN_KEY))); - return selectionFinder - .findSelections(selectionSet, SelectionQuery.builder().selectionPath(fullPath).build()) - .collect(Collectors.toUnmodifiableList()); - } - - @AllArgsConstructor - private class DefaultSpanJoiner implements SpanJoiner { - - private final GraphQlRequestContext context; - private final TimeRangeArgument timeRange; - private final List selectedFields; - - @Override - public Single> joinSpans( - Collection joinSources, SpanIdGetter spanIdGetter) { - return this.buildSourceToIdMap(joinSources, spanIdGetter).flatMap(this::joinSpans); - } - - private Single> joinSpans(Map sourceToSpanIdMap) { - return this.buildSpanRequest(sourceToSpanIdMap) - .flatMap(spanDao::getSpans) - .map(this::buildSpanIdToSpanMap) - .map(spanIdToSpanMap -> buildSourceToSpanMap(sourceToSpanIdMap, spanIdToSpanMap)); - } - - private Map buildSourceToSpanMap( - Map sourceToSpanIdMap, Map spanIdToSpanMap) { - return sourceToSpanIdMap.entrySet().stream() - .filter(entry -> spanIdToSpanMap.containsKey(entry.getValue())) - .collect( - Collectors.toUnmodifiableMap( - Entry::getKey, entry -> spanIdToSpanMap.get(entry.getValue()))); - } - - private Map buildSpanIdToSpanMap(SpanResultSet resultSet) { - return resultSet.results().stream() - .collect(Collectors.toUnmodifiableMap(Identifiable::id, Function.identity())); - } - - private Single buildSpanRequest(Map sourceToSpanIdMap) { - Collection spanIds = sourceToSpanIdMap.values(); - return buildSpanIdsFilter(spanIds) - .flatMap(filterArguments -> buildSpanRequest(spanIds.size(), filterArguments)); - } - - private Single buildSpanRequest( - int size, List> filterArguments) { - return resultSetRequestBuilder - .build( - context, - SPAN, - size, - ZERO_OFFSET, - timeRange, - Collections.emptyList(), - filterArguments, - selectedFields.stream(), - Optional.empty()) - .map(spanEventsRequest -> new SpanJoinRequest(context, spanEventsRequest)); - } - - private Single>> buildSpanIdsFilter( - Collection spanIds) { - return filterRequestBuilder.build(context, SPAN, Set.of(new SpanIdFilter(spanIds))); - } - - private Single> buildSourceToIdMap( - Collection joinSources, SpanIdGetter spanIdGetter) { - return Observable.fromIterable(joinSources) - .flatMapSingle(source -> this.maybeBuildMapEntry(source, spanIdGetter)) - .collect(Collectors.toMap(Entry::getKey, Entry::getValue)); - } - - private Single> maybeBuildMapEntry( - T source, SpanIdGetter spanIdGetter) { - return spanIdGetter.getSpanId(source).map(id -> Map.entry(source, id)); - } - } - - @Value - @Accessors(fluent = true) - private static class SpanIdFilter implements FilterArgument { - FilterType type = FilterType.ID; - String key = null; - AttributeExpression keyExpression = null; - FilterOperatorType operator = FilterOperatorType.IN; - Collection value; - AttributeScope idType = null; - String idScope = SPAN; - } - - @Value - @Accessors(fluent = true) - private static class SpanJoinRequest implements SpanRequest { - GraphQlRequestContext context; - ResultSetRequest spanEventsRequest; - Collection logEventAttributes = Collections.emptyList(); - boolean fetchTotal = false; - } -} diff --git a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoinerModule.java b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/JoinerModule.java similarity index 58% rename from hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoinerModule.java rename to hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/JoinerModule.java index 162eb36b..48897fd2 100644 --- a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoinerModule.java +++ b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/JoinerModule.java @@ -4,13 +4,18 @@ import org.hypertrace.core.graphql.common.request.FilterRequestBuilder; import org.hypertrace.core.graphql.common.request.ResultSetRequestBuilder; import org.hypertrace.core.graphql.span.dao.SpanDao; +import org.hypertrace.core.graphql.span.joiner.multiple.span.DefaultMultipleSpanJoinerBuilder; +import org.hypertrace.core.graphql.span.joiner.multiple.span.MultipleSpanJoinerBuilder; +import org.hypertrace.core.graphql.span.joiner.single.span.DefaultSpanJoinerBuilder; +import org.hypertrace.core.graphql.span.joiner.single.span.SpanJoinerBuilder; import org.hypertrace.core.graphql.utils.schema.GraphQlSelectionFinder; -public class SpanJoinerModule extends AbstractModule { +public class JoinerModule extends AbstractModule { @Override protected void configure() { bind(SpanJoinerBuilder.class).to(DefaultSpanJoinerBuilder.class); + bind(MultipleSpanJoinerBuilder.class).to(DefaultMultipleSpanJoinerBuilder.class); requireBinding(SpanDao.class); requireBinding(GraphQlSelectionFinder.class); diff --git a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SourceToSpansProvider.java b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SourceToSpansProvider.java new file mode 100644 index 00000000..1b7ace58 --- /dev/null +++ b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SourceToSpansProvider.java @@ -0,0 +1,143 @@ +package org.hypertrace.core.graphql.span.joiner; + +import static org.hypertrace.core.graphql.atttributes.scopes.HypertraceCoreAttributeScopeString.SPAN; + +import com.google.common.collect.ArrayListMultimap; +import com.google.common.collect.ListMultimap; +import com.google.common.collect.Multimaps; +import com.google.inject.Inject; +import graphql.schema.SelectedField; +import io.reactivex.rxjava3.core.Single; +import java.util.Collection; +import java.util.Collections; +import java.util.List; +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 lombok.AllArgsConstructor; +import lombok.Value; +import lombok.experimental.Accessors; +import org.hypertrace.core.graphql.common.request.AttributeAssociation; +import org.hypertrace.core.graphql.common.request.AttributeRequest; +import org.hypertrace.core.graphql.common.request.FilterRequestBuilder; +import org.hypertrace.core.graphql.common.request.ResultSetRequest; +import org.hypertrace.core.graphql.common.request.ResultSetRequestBuilder; +import org.hypertrace.core.graphql.common.schema.arguments.TimeRangeArgument; +import org.hypertrace.core.graphql.common.schema.attributes.AttributeScope; +import org.hypertrace.core.graphql.common.schema.attributes.arguments.AttributeExpression; +import org.hypertrace.core.graphql.common.schema.id.Identifiable; +import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterArgument; +import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterOperatorType; +import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterType; +import org.hypertrace.core.graphql.common.schema.results.arguments.order.OrderArgument; +import org.hypertrace.core.graphql.context.GraphQlRequestContext; +import org.hypertrace.core.graphql.span.dao.SpanDao; +import org.hypertrace.core.graphql.span.request.SpanRequest; +import org.hypertrace.core.graphql.span.schema.Span; +import org.hypertrace.core.graphql.span.schema.SpanResultSet; + +@AllArgsConstructor(onConstructor_ = {@Inject}) +public class SourceToSpansProvider { + + private static final int ZERO_OFFSET = 0; + + private final SpanDao spanDao; + private final ResultSetRequestBuilder resultSetRequestBuilder; + private final FilterRequestBuilder filterRequestBuilder; + + public Single> joinSpans( + GraphQlRequestContext context, + TimeRangeArgument timeRange, + List selectedFields, + Map> sourceToSpanIdsMap) { + return this.buildSpanRequest(context, timeRange, selectedFields, sourceToSpanIdsMap) + .flatMap(spanDao::getSpans) + .map(this::buildSpanIdToSpanMap) + .map(spanIdToSpanMap -> buildSourceToSpanListMultiMap(sourceToSpanIdsMap, spanIdToSpanMap)); + } + + private ListMultimap buildSourceToSpanListMultiMap( + Map> sourceToSpanIdsMap, Map spanIdToSpanMap) { + ListMultimap listMultimap = ArrayListMultimap.create(); + for (Entry> entry : sourceToSpanIdsMap.entrySet()) { + T source = entry.getKey(); + for (String spanId : entry.getValue()) { + if (spanIdToSpanMap.containsKey(spanId)) { + listMultimap.put(source, spanIdToSpanMap.get(spanId)); + } + } + } + return Multimaps.unmodifiableListMultimap(listMultimap); + } + + private Map buildSpanIdToSpanMap(SpanResultSet resultSet) { + return resultSet.results().stream() + .collect(Collectors.toUnmodifiableMap(Identifiable::id, Function.identity())); + } + + private Single buildSpanRequest( + GraphQlRequestContext context, + TimeRangeArgument timeRange, + List selectedFields, + Map> sourceToSpanIdsMap) { + Collection spanIds = + sourceToSpanIdsMap.values().stream() + .flatMap(List::stream) + .distinct() + .collect(Collectors.toUnmodifiableList()); + return buildSpanIdsFilter(context, spanIds) + .flatMap( + filterArguments -> + buildSpanRequest( + spanIds.size(), context, timeRange, selectedFields, filterArguments)); + } + + private Single buildSpanRequest( + int size, + GraphQlRequestContext context, + TimeRangeArgument timeRange, + List selectedFields, + List> filterArguments) { + return resultSetRequestBuilder + .build( + context, + SPAN, + size, + ZERO_OFFSET, + timeRange, + Collections.emptyList(), + filterArguments, + selectedFields.stream(), + Optional.empty()) + .map(spanEventsRequest -> new SpanJoinRequest(context, spanEventsRequest)); + } + + private Single>> buildSpanIdsFilter( + GraphQlRequestContext context, Collection spanIds) { + return filterRequestBuilder.build(context, SPAN, Set.of(new SpanIdFilter(spanIds))); + } + + @Value + @Accessors(fluent = true) + private static class SpanIdFilter implements FilterArgument { + FilterType type = FilterType.ID; + String key = null; + AttributeExpression keyExpression = null; + FilterOperatorType operator = FilterOperatorType.IN; + Collection value; + AttributeScope idType = null; + String idScope = SPAN; + } + + @Value + @Accessors(fluent = true) + private static class SpanJoinRequest implements SpanRequest { + GraphQlRequestContext context; + ResultSetRequest spanEventsRequest; + Collection logEventAttributes = Collections.emptyList(); + boolean fetchTotal = false; + } +} diff --git a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/DefaultMultipleSpanJoinerBuilder.java b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/DefaultMultipleSpanJoinerBuilder.java new file mode 100644 index 00000000..542c6bad --- /dev/null +++ b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/DefaultMultipleSpanJoinerBuilder.java @@ -0,0 +1,87 @@ +package org.hypertrace.core.graphql.span.joiner.multiple.span; + +import static com.google.common.collect.ImmutableList.copyOf; +import static com.google.common.collect.Iterables.concat; +import static org.hypertrace.core.graphql.span.joiner.multiple.span.MultipleSpanJoin.SPANS_KEY; + +import com.google.common.collect.ListMultimap; +import graphql.schema.DataFetchingFieldSelectionSet; +import graphql.schema.SelectedField; +import io.reactivex.rxjava3.core.Observable; +import io.reactivex.rxjava3.core.Single; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; +import java.util.stream.Collectors; +import javax.inject.Inject; +import lombok.AllArgsConstructor; +import org.hypertrace.core.graphql.common.schema.arguments.TimeRangeArgument; +import org.hypertrace.core.graphql.context.GraphQlRequestContext; +import org.hypertrace.core.graphql.span.joiner.SourceToSpansProvider; +import org.hypertrace.core.graphql.span.schema.Span; +import org.hypertrace.core.graphql.utils.schema.GraphQlSelectionFinder; +import org.hypertrace.core.graphql.utils.schema.SelectionQuery; + +public class DefaultMultipleSpanJoinerBuilder implements MultipleSpanJoinerBuilder { + + private final GraphQlSelectionFinder selectionFinder; + + private final SourceToSpansProvider sourceToSpansProvider; + + @Inject + DefaultMultipleSpanJoinerBuilder( + GraphQlSelectionFinder selectionFinder, SourceToSpansProvider sourceToSpansProvider) { + this.selectionFinder = selectionFinder; + this.sourceToSpansProvider = sourceToSpansProvider; + } + + @Override + public Single build( + GraphQlRequestContext context, + TimeRangeArgument timeRange, + DataFetchingFieldSelectionSet selectionSet, + List pathToSpanJoin) { + return Single.just( + new DefaultMultipleSpanJoiner( + context, timeRange, this.getSelections(selectionSet, pathToSpanJoin))); + } + + private List getSelections( + DataFetchingFieldSelectionSet selectionSet, List pathToSpanJoin) { + List fullPath = copyOf(concat(pathToSpanJoin, List.of(SPANS_KEY))); + return selectionFinder + .findSelections(selectionSet, SelectionQuery.builder().selectionPath(fullPath).build()) + .collect(Collectors.toUnmodifiableList()); + } + + @AllArgsConstructor + private class DefaultMultipleSpanJoiner implements MultipleSpanJoiner { + + private final GraphQlRequestContext context; + private final TimeRangeArgument timeRange; + private final List selectedFields; + + @Override + public Single> joinMultipleSpans( + Collection joinSources, MultipleSpanIdGetter multipleSpanIdGetter) { + return this.buildSourceToIdsMap(joinSources, multipleSpanIdGetter) + .flatMap( + sourceToSpanIdsMap -> + sourceToSpansProvider.joinSpans( + context, timeRange, selectedFields, sourceToSpanIdsMap)); + } + + private Single>> buildSourceToIdsMap( + Collection joinSources, MultipleSpanIdGetter multipleSpanIdGetter) { + return Observable.fromIterable(joinSources) + .flatMapSingle(source -> this.maybeBuildMapEntry(source, multipleSpanIdGetter)) + .collect(Collectors.toMap(Entry::getKey, Entry::getValue)); + } + + private Single>> maybeBuildMapEntry( + T source, MultipleSpanIdGetter multipleSpanIdGetter) { + return multipleSpanIdGetter.getSpanIds(source).map(ids -> Map.entry(source, ids)); + } + } +} diff --git a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoin.java b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoin.java new file mode 100644 index 00000000..a3cd5304 --- /dev/null +++ b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoin.java @@ -0,0 +1,14 @@ +package org.hypertrace.core.graphql.span.joiner.multiple.span; + +import graphql.annotations.annotationTypes.GraphQLField; +import graphql.annotations.annotationTypes.GraphQLName; +import java.util.List; +import org.hypertrace.core.graphql.span.schema.Span; + +public interface MultipleSpanJoin { + String SPANS_KEY = "spans"; + + @GraphQLField + @GraphQLName(SPANS_KEY) + List spans(); +} diff --git a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoiner.java b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoiner.java new file mode 100644 index 00000000..2d1ef966 --- /dev/null +++ b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoiner.java @@ -0,0 +1,29 @@ +package org.hypertrace.core.graphql.span.joiner.multiple.span; + +import com.google.common.collect.ArrayListMultimap; +import com.google.common.collect.ListMultimap; +import io.reactivex.rxjava3.core.Single; +import java.util.Collection; +import java.util.List; +import org.hypertrace.core.graphql.span.schema.Span; + +public interface MultipleSpanJoiner { + + /** A NOOP joiner */ + MultipleSpanJoiner NO_OP_JOINER = + new MultipleSpanJoiner() { + @Override + public Single> joinMultipleSpans( + Collection joinSources, MultipleSpanIdGetter multipleSpanIdGetter) { + return Single.just(ArrayListMultimap.create()); + } + }; + + Single> joinMultipleSpans( + Collection joinSources, MultipleSpanIdGetter multipleSpanIdGetter); + + @FunctionalInterface + interface MultipleSpanIdGetter { + Single> getSpanIds(T source); + } +} diff --git a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoinerBuilder.java b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoinerBuilder.java new file mode 100644 index 00000000..7ee5c01b --- /dev/null +++ b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoinerBuilder.java @@ -0,0 +1,15 @@ +package org.hypertrace.core.graphql.span.joiner.multiple.span; + +import graphql.schema.DataFetchingFieldSelectionSet; +import io.reactivex.rxjava3.core.Single; +import java.util.List; +import org.hypertrace.core.graphql.common.schema.arguments.TimeRangeArgument; +import org.hypertrace.core.graphql.context.GraphQlRequestContext; + +public interface MultipleSpanJoinerBuilder { + Single build( + GraphQlRequestContext context, + TimeRangeArgument timeRange, + DataFetchingFieldSelectionSet selectionSet, + List pathToMultipleSpanJoin); +} diff --git a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/single/span/DefaultSpanJoinerBuilder.java b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/single/span/DefaultSpanJoinerBuilder.java new file mode 100644 index 00000000..6ab057ed --- /dev/null +++ b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/single/span/DefaultSpanJoinerBuilder.java @@ -0,0 +1,95 @@ +package org.hypertrace.core.graphql.span.joiner.single.span; + +import static com.google.common.collect.ImmutableList.copyOf; +import static com.google.common.collect.Iterables.concat; +import static org.hypertrace.core.graphql.span.joiner.single.span.SpanJoin.SPAN_KEY; + +import com.google.common.collect.Multimaps; +import graphql.schema.DataFetchingFieldSelectionSet; +import graphql.schema.SelectedField; +import io.reactivex.rxjava3.core.Observable; +import io.reactivex.rxjava3.core.Single; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; +import java.util.stream.Collectors; +import javax.inject.Inject; +import lombok.AllArgsConstructor; +import org.hypertrace.core.graphql.common.schema.arguments.TimeRangeArgument; +import org.hypertrace.core.graphql.context.GraphQlRequestContext; +import org.hypertrace.core.graphql.span.joiner.SourceToSpansProvider; +import org.hypertrace.core.graphql.span.schema.Span; +import org.hypertrace.core.graphql.utils.schema.GraphQlSelectionFinder; +import org.hypertrace.core.graphql.utils.schema.SelectionQuery; + +public class DefaultSpanJoinerBuilder implements SpanJoinerBuilder { + + private final GraphQlSelectionFinder selectionFinder; + + private final SourceToSpansProvider sourceToSpansProvider; + + @Inject + DefaultSpanJoinerBuilder( + GraphQlSelectionFinder selectionFinder, SourceToSpansProvider sourceToSpansProvider) { + this.selectionFinder = selectionFinder; + this.sourceToSpansProvider = sourceToSpansProvider; + } + + @Override + public Single build( + GraphQlRequestContext context, + TimeRangeArgument timeRange, + DataFetchingFieldSelectionSet selectionSet, + List pathToSpanJoin) { + return Single.just( + new DefaultSpanJoiner( + context, timeRange, this.getSelections(selectionSet, pathToSpanJoin))); + } + + private List getSelections( + DataFetchingFieldSelectionSet selectionSet, List pathToSpanJoin) { + List fullPath = copyOf(concat(pathToSpanJoin, List.of(SPAN_KEY))); + return selectionFinder + .findSelections(selectionSet, SelectionQuery.builder().selectionPath(fullPath).build()) + .collect(Collectors.toUnmodifiableList()); + } + + @AllArgsConstructor + private class DefaultSpanJoiner implements SpanJoiner { + + private final GraphQlRequestContext context; + private final TimeRangeArgument timeRange; + private final List selectedFields; + + @Override + public Single> joinSpans( + Collection joinSources, SpanIdGetter spanIdGetter) { + return this.buildSourceToIdMap(joinSources, spanIdGetter) + .flatMap( + sourceToSpanIdsMap -> + sourceToSpansProvider.joinSpans( + context, timeRange, selectedFields, sourceToSpanIdsMap)) + .map(Multimaps::asMap) + .map(this::reduceMap); + } + + private Map reduceMap(Map> multiMap) { + return multiMap.entrySet().stream() + .filter(entry -> !entry.getValue().isEmpty()) + .collect(Collectors.toUnmodifiableMap(Entry::getKey, entry -> entry.getValue().get(0))); + } + + private Single>> buildSourceToIdMap( + Collection joinSources, SpanIdGetter spanIdGetter) { + return Observable.fromIterable(joinSources) + .flatMapSingle(source -> this.maybeBuildMapEntry(source, spanIdGetter)) + .collect(Collectors.toMap(Entry::getKey, Entry::getValue)); + } + + private Single>> maybeBuildMapEntry( + T source, SpanIdGetter spanIdGetter) { + return spanIdGetter.getSpanId(source).map(List::of).map(ids -> Map.entry(source, ids)); + } + } +} diff --git a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoin.java b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/single/span/SpanJoin.java similarity index 82% rename from hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoin.java rename to hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/single/span/SpanJoin.java index 6611267b..4c06b252 100644 --- a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoin.java +++ b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/single/span/SpanJoin.java @@ -1,4 +1,4 @@ -package org.hypertrace.core.graphql.span.joiner; +package org.hypertrace.core.graphql.span.joiner.single.span; import graphql.annotations.annotationTypes.GraphQLField; import graphql.annotations.annotationTypes.GraphQLName; 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/single/span/SpanJoiner.java similarity index 91% rename from hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoiner.java rename to hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/single/span/SpanJoiner.java index 7cfe87b7..40547e4d 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/single/span/SpanJoiner.java @@ -1,4 +1,4 @@ -package org.hypertrace.core.graphql.span.joiner; +package org.hypertrace.core.graphql.span.joiner.single.span; import io.reactivex.rxjava3.core.Single; import java.util.Collection; diff --git a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoinerBuilder.java b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/single/span/SpanJoinerBuilder.java similarity index 88% rename from hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoinerBuilder.java rename to hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/single/span/SpanJoinerBuilder.java index a06ccf45..6fbc2b17 100644 --- a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoinerBuilder.java +++ b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/single/span/SpanJoinerBuilder.java @@ -1,4 +1,4 @@ -package org.hypertrace.core.graphql.span.joiner; +package org.hypertrace.core.graphql.span.joiner.single.span; import graphql.schema.DataFetchingFieldSelectionSet; import io.reactivex.rxjava3.core.Single; diff --git a/hypertrace-core-graphql-span-schema/src/test/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoinerBuilderTest.java b/hypertrace-core-graphql-span-schema/src/test/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoinerBuilderTest.java new file mode 100644 index 00000000..a93c93b6 --- /dev/null +++ b/hypertrace-core-graphql-span-schema/src/test/java/org/hypertrace/core/graphql/span/joiner/multiple/span/MultipleSpanJoinerBuilderTest.java @@ -0,0 +1,165 @@ +package org.hypertrace.core.graphql.span.joiner.multiple.span; + +import static java.util.Collections.emptyList; +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.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.common.collect.ArrayListMultimap; +import com.google.common.collect.ListMultimap; +import graphql.schema.DataFetchingFieldSelectionSet; +import graphql.schema.SelectedField; +import io.reactivex.rxjava3.core.Single; +import java.util.List; +import java.util.Optional; +import java.util.stream.Stream; +import lombok.Value; +import lombok.experimental.Accessors; +import org.hypertrace.core.graphql.common.request.AttributeAssociation; +import org.hypertrace.core.graphql.common.request.FilterRequestBuilder; +import org.hypertrace.core.graphql.common.request.ResultSetRequest; +import org.hypertrace.core.graphql.common.request.ResultSetRequestBuilder; +import org.hypertrace.core.graphql.common.schema.arguments.TimeRangeArgument; +import org.hypertrace.core.graphql.common.schema.attributes.arguments.AttributeExpression; +import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterArgument; +import org.hypertrace.core.graphql.common.schema.results.arguments.order.OrderArgument; +import org.hypertrace.core.graphql.context.GraphQlRequestContext; +import org.hypertrace.core.graphql.log.event.schema.LogEventResultSet; +import org.hypertrace.core.graphql.span.dao.SpanDao; +import org.hypertrace.core.graphql.span.joiner.SourceToSpansProvider; +import org.hypertrace.core.graphql.span.joiner.multiple.span.MultipleSpanJoiner.MultipleSpanIdGetter; +import org.hypertrace.core.graphql.span.request.SpanRequest; +import org.hypertrace.core.graphql.span.schema.Span; +import org.hypertrace.core.graphql.span.schema.SpanResultSet; +import org.hypertrace.core.graphql.utils.schema.GraphQlSelectionFinder; +import org.hypertrace.core.graphql.utils.schema.SelectionQuery; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +public class MultipleSpanJoinerBuilderTest { + + private static final String SPAN_ID1 = "spanId1"; + private static final String SPAN_ID2 = "spanId2"; + private static final String SPAN_ID3 = "spanId3"; + private static final String SPAN_ID4 = "spanId4"; + + @Mock private SpanDao mockSpanDao; + @Mock private GraphQlSelectionFinder mockSelectionFinder; + @Mock private ResultSetRequestBuilder mockResultSetRequestBuilder; + @Mock private FilterRequestBuilder mockFilterRequestBuilder; + @Mock private DataFetchingFieldSelectionSet mockSelectionSet; + @Mock private GraphQlRequestContext mockRequestContext; + @Mock private ResultSetRequest mockResultSetRequest; + @Mock private AttributeAssociation mockFilter; + @Mock private TimeRangeArgument mockTimeRangeArgument; + + private MultipleSpanJoinerBuilder multipleSpanJoinerBuilder; + + @BeforeEach + void setup() { + multipleSpanJoinerBuilder = + new DefaultMultipleSpanJoinerBuilder( + mockSelectionFinder, + new SourceToSpansProvider( + mockSpanDao, mockResultSetRequestBuilder, mockFilterRequestBuilder)); + } + + @Test + void fetchMultipleSpans() { + Span span1 = new TestSpan(SPAN_ID1); + Span span2 = new TestSpan(SPAN_ID2); + TestJoinSource joinSource1 = new TestJoinSource(List.of(SPAN_ID1, SPAN_ID2)); + TestJoinSource joinSource2 = new TestJoinSource(List.of(SPAN_ID3, SPAN_ID4)); + ListMultimap expected = ArrayListMultimap.create(); + expected.put(joinSource1, span1); + expected.put(joinSource1, span2); + List joinSources = List.of(joinSource1, joinSource2); + mockRequestedSelectionFields(List.of(mock(SelectedField.class), mock(SelectedField.class))); + mockRequestBuilding(); + mockResult(List.of(span1, span2)); + MultipleSpanJoiner joiner = + this.multipleSpanJoinerBuilder + .build( + this.mockRequestContext, + this.mockTimeRangeArgument, + this.mockSelectionSet, + List.of("pathToSpans")) + .blockingGet(); + assertEquals( + expected, + joiner.joinMultipleSpans(joinSources, new TestJoinSourceIdGetter()).blockingGet()); + } + + private void mockRequestBuilding() { + when(mockFilterRequestBuilder.build(eq(mockRequestContext), eq(SPAN), anySet())) + .thenAnswer(invocation -> Single.just(List.of(mockFilter))); + + when(mockResultSetRequestBuilder.build( + eq(mockRequestContext), + eq(SPAN), + eq(4), + eq(0), + eq(mockTimeRangeArgument), + eq(emptyList()), + eq(List.of(mockFilter)), + any(Stream.class), + eq(Optional.empty()))) + .thenAnswer(invocation -> Single.just(mockResultSetRequest)); + } + + private void mockRequestedSelectionFields(List selectedFields) { + when(mockSelectionFinder.findSelections( + mockSelectionSet, + SelectionQuery.builder().selectionPath(List.of("pathToSpans", "spans")).build())) + .thenAnswer(invocation -> selectedFields.stream()); + } + + private void mockResult(List spans) { + when(mockSpanDao.getSpans(any(SpanRequest.class))) + .thenAnswer(invocation -> Single.just(new TestSpanResultSet(spans))); + } + + @Value + private static class TestJoinSource { + List spanIds; + } + + private static class TestJoinSourceIdGetter implements MultipleSpanIdGetter { + @Override + public Single> getSpanIds(TestJoinSource source) { + return Single.just(source.getSpanIds()); + } + } + + @Value + @Accessors(fluent = true) + private static class TestSpanResultSet implements SpanResultSet { + List results; + long count = 0; + long total = 0; + } + + @Value + @Accessors(fluent = true) + private static class TestSpan implements Span { + String id; + + @Override + public Object attribute(AttributeExpression expression) { + return null; + } + + @Override + public LogEventResultSet logEvents() { + return null; + } + } +} 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/single/span/SpanJoinerBuilderTest.java similarity index 91% rename from hypertrace-core-graphql-span-schema/src/test/java/org/hypertrace/core/graphql/span/joiner/SpanJoinerBuilderTest.java rename to hypertrace-core-graphql-span-schema/src/test/java/org/hypertrace/core/graphql/span/joiner/single/span/SpanJoinerBuilderTest.java index d135e32a..114986fe 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/single/span/SpanJoinerBuilderTest.java @@ -1,4 +1,4 @@ -package org.hypertrace.core.graphql.span.joiner; +package org.hypertrace.core.graphql.span.joiner.single.span; import static java.util.Collections.emptyList; import static java.util.Map.entry; @@ -30,7 +30,8 @@ import org.hypertrace.core.graphql.context.GraphQlRequestContext; import org.hypertrace.core.graphql.log.event.schema.LogEventResultSet; import org.hypertrace.core.graphql.span.dao.SpanDao; -import org.hypertrace.core.graphql.span.joiner.SpanJoiner.SpanIdGetter; +import org.hypertrace.core.graphql.span.joiner.SourceToSpansProvider; +import org.hypertrace.core.graphql.span.joiner.single.span.SpanJoiner.SpanIdGetter; import org.hypertrace.core.graphql.span.request.SpanRequest; import org.hypertrace.core.graphql.span.schema.Span; import org.hypertrace.core.graphql.span.schema.SpanResultSet; @@ -64,10 +65,9 @@ public class SpanJoinerBuilderTest { void setup() { spanJoinerBuilder = new DefaultSpanJoinerBuilder( - mockSpanDao, mockSelectionFinder, - mockResultSetRequestBuilder, - mockFilterRequestBuilder); + new SourceToSpansProvider( + mockSpanDao, mockResultSetRequestBuilder, mockFilterRequestBuilder)); } @Test @@ -79,8 +79,7 @@ void fetchSpans() { Map expected = Map.ofEntries(entry(joinSource1, span1), entry(joinSource2, span2)); List joinSources = List.of(joinSource1, joinSource2); - mockRequestedSelectionFields( - List.of(mock(SelectedField.class), mock(SelectedField.class)), "pathToSpan"); + mockRequestedSelectionFields(List.of(mock(SelectedField.class), mock(SelectedField.class))); mockRequestBuilding(); mockResult(List.of(span1, span2)); SpanJoiner joiner = @@ -112,10 +111,10 @@ private void mockRequestBuilding() { .thenReturn(Single.just(mockResultSetRequest)); } - private void mockRequestedSelectionFields(List selectedFields, String location) { + private void mockRequestedSelectionFields(List selectedFields) { when(mockSelectionFinder.findSelections( mockSelectionSet, - SelectionQuery.builder().selectionPath(List.of(location, "span")).build())) + SelectionQuery.builder().selectionPath(List.of("pathToSpan", "span")).build())) .thenReturn(selectedFields.stream()); }