Skip to content

Commit 36bd206

Browse files
chore | Add support for providing filters in spanJoiner (#181)
* Add support for providing filters in spanJoiner * Resolve PR reviews
1 parent 4202d01 commit 36bd206

File tree

3 files changed

+60
-21
lines changed

3 files changed

+60
-21
lines changed

hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/DefaultSpanJoinerBuilder.java

Lines changed: 25 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,9 @@
1818
import java.util.Map;
1919
import java.util.Map.Entry;
2020
import java.util.Optional;
21-
import java.util.Set;
2221
import java.util.function.Function;
2322
import java.util.stream.Collectors;
23+
import java.util.stream.Stream;
2424
import javax.inject.Inject;
2525
import lombok.AllArgsConstructor;
2626
import lombok.Value;
@@ -86,26 +86,32 @@ private class DefaultSpanJoiner implements SpanJoiner {
8686

8787
@Override
8888
public <T> Single<Map<T, Span>> joinSpan(
89-
Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
89+
Collection<T> joinSources,
90+
SpanIdGetter<T> spanIdGetter,
91+
Collection<FilterArgument> filterArguments) {
9092
Function<T, Single<List<String>>> idsGetter =
9193
source -> spanIdGetter.getSpanId(source).map(List::of);
92-
return this.joinSpans(joinSources, idsGetter, SPAN_KEY).map(this::reduceMap);
94+
return this.joinSpans(joinSources, idsGetter, filterArguments, SPAN_KEY).map(this::reduceMap);
9395
}
9496

9597
@Override
9698
public <T> Single<ListMultimap<T, Span>> joinSpans(
97-
Collection<T> joinSources, MultipleSpanIdGetter<T> multipleSpanIdGetter) {
98-
return this.joinSpans(joinSources, multipleSpanIdGetter::getSpanIds, SPANS_KEY);
99+
Collection<T> joinSources,
100+
MultipleSpanIdGetter<T> multipleSpanIdGetter,
101+
Collection<FilterArgument> filterArguments) {
102+
return this.joinSpans(
103+
joinSources, multipleSpanIdGetter::getSpanIds, filterArguments, SPANS_KEY);
99104
}
100105

101106
private <T> Single<ListMultimap<T, Span>> joinSpans(
102107
Collection<T> joinSources,
103108
Function<T, Single<List<String>>> idsGetter,
109+
Collection<FilterArgument> filterArguments,
104110
String joinSpanKey) {
105111
return this.buildSourceToIdsMap(joinSources, idsGetter)
106112
.flatMap(
107113
sourceToSpanIdsMap ->
108-
this.buildSpanRequest(sourceToSpanIdsMap, joinSpanKey)
114+
this.buildSpanRequest(sourceToSpanIdsMap, joinSpanKey, filterArguments)
109115
.flatMap(spanDao::getSpans)
110116
.map(this::buildSpanIdToSpanMap)
111117
.map(
@@ -152,15 +158,16 @@ private Map<String, Span> buildSpanIdToSpanMap(SpanResultSet resultSet) {
152158
}
153159

154160
private <T> Single<SpanRequest> buildSpanRequest(
155-
ListMultimap<T, String> sourceToSpanIdsMultimap, String joinSpanKey) {
161+
ListMultimap<T, String> sourceToSpanIdsMultimap,
162+
String joinSpanKey,
163+
Collection<FilterArgument> filterArguments) {
156164
Collection<String> spanIds =
157165
sourceToSpanIdsMultimap.values().stream()
158166
.distinct()
159167
.collect(Collectors.toUnmodifiableList());
160168
List<SelectedField> selectedFields = getSelections(joinSpanKey);
161-
return buildSpanIdsFilter(context, spanIds)
162-
.flatMap(
163-
filterArguments -> buildSpanRequest(spanIds.size(), filterArguments, selectedFields));
169+
return buildSpanIdsFilter(context, spanIds, filterArguments)
170+
.flatMap(filters -> buildSpanRequest(spanIds.size(), filters, selectedFields));
164171
}
165172

166173
private Single<SpanRequest> buildSpanRequest(
@@ -182,8 +189,14 @@ private Single<SpanRequest> buildSpanRequest(
182189
}
183190

184191
private Single<List<AttributeAssociation<FilterArgument>>> buildSpanIdsFilter(
185-
GraphQlRequestContext context, Collection<String> spanIds) {
186-
return filterRequestBuilder.build(context, SPAN, Set.of(new SpanIdFilter(spanIds)));
192+
GraphQlRequestContext context,
193+
Collection<String> spanIds,
194+
Collection<FilterArgument> filterArguments) {
195+
return filterRequestBuilder.build(
196+
context,
197+
SPAN,
198+
Stream.concat(filterArguments.stream(), Stream.of(new SpanIdFilter(spanIds)))
199+
.collect(Collectors.toUnmodifiableList()));
187200
}
188201
}
189202

hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoiner.java

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import java.util.Collections;
88
import java.util.List;
99
import java.util.Map;
10+
import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterArgument;
1011
import org.hypertrace.core.graphql.span.schema.Span;
1112

1213
public interface SpanJoiner {
@@ -16,21 +17,40 @@ public interface SpanJoiner {
1617
new SpanJoiner() {
1718
@Override
1819
public <T> Single<Map<T, Span>> joinSpan(
19-
Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
20+
Collection<T> joinSources,
21+
SpanIdGetter<T> spanIdGetter,
22+
Collection<FilterArgument> filterArguments) {
2023
return Single.just(Collections.emptyMap());
2124
}
2225

2326
@Override
2427
public <T> Single<ListMultimap<T, Span>> joinSpans(
25-
Collection<T> joinSources, MultipleSpanIdGetter<T> multipleSpanIdGetter) {
28+
Collection<T> joinSources,
29+
MultipleSpanIdGetter<T> multipleSpanIdGetter,
30+
Collection<FilterArgument> filterArguments) {
2631
return Single.just(ArrayListMultimap.create());
2732
}
2833
};
2934

30-
<T> Single<Map<T, Span>> joinSpan(Collection<T> joinSources, SpanIdGetter<T> spanIdGetter);
35+
default <T> Single<Map<T, Span>> joinSpan(
36+
Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
37+
return joinSpan(joinSources, spanIdGetter, Collections.emptyList());
38+
}
39+
40+
default <T> Single<ListMultimap<T, Span>> joinSpans(
41+
Collection<T> joinSources, MultipleSpanIdGetter<T> multipleSpanIdGetter) {
42+
return joinSpans(joinSources, multipleSpanIdGetter, Collections.emptyList());
43+
}
44+
45+
<T> Single<Map<T, Span>> joinSpan(
46+
Collection<T> joinSources,
47+
SpanIdGetter<T> spanIdGetter,
48+
Collection<FilterArgument> filterArguments);
3149

3250
<T> Single<ListMultimap<T, Span>> joinSpans(
33-
Collection<T> joinSources, MultipleSpanIdGetter<T> multipleSpanIdGetter);
51+
Collection<T> joinSources,
52+
MultipleSpanIdGetter<T> multipleSpanIdGetter,
53+
Collection<FilterArgument> filterArguments);
3454

3555
@FunctionalInterface
3656
interface SpanIdGetter<T> {

hypertrace-core-graphql-span-schema/src/test/java/org/hypertrace/core/graphql/span/joiner/SpanJoinerBuilderTest.java

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import static org.hypertrace.core.graphql.atttributes.scopes.HypertraceCoreAttributeScopeString.SPAN;
66
import static org.junit.jupiter.api.Assertions.assertEquals;
77
import static org.mockito.ArgumentMatchers.any;
8-
import static org.mockito.ArgumentMatchers.anySet;
8+
import static org.mockito.ArgumentMatchers.anyList;
99
import static org.mockito.ArgumentMatchers.eq;
1010
import static org.mockito.Mockito.mock;
1111
import static org.mockito.Mockito.when;
@@ -15,6 +15,7 @@
1515
import graphql.schema.DataFetchingFieldSelectionSet;
1616
import graphql.schema.SelectedField;
1717
import io.reactivex.rxjava3.core.Single;
18+
import java.util.Collections;
1819
import java.util.List;
1920
import java.util.Map;
2021
import java.util.Optional;
@@ -88,7 +89,7 @@ void fetchSpans() {
8889
mockSelectionSet,
8990
SelectionQuery.builder().selectionPath(List.of("pathToSpan", "span")).build()))
9091
.thenReturn(Stream.of(mock(SelectedField.class), mock(SelectedField.class)));
91-
when(mockFilterRequestBuilder.build(eq(mockRequestContext), eq(SPAN), anySet()))
92+
when(mockFilterRequestBuilder.build(eq(mockRequestContext), eq(SPAN), anyList()))
9293
.thenReturn(Single.just(List.of(mockFilter)));
9394

9495
when(mockResultSetRequestBuilder.build(
@@ -112,7 +113,10 @@ void fetchSpans() {
112113
List.of("pathToSpan"))
113114
.blockingGet();
114115
assertEquals(
115-
expected, joiner.joinSpan(joinSources, new TestJoinSourceIdGetter()).blockingGet());
116+
expected,
117+
joiner
118+
.joinSpan(joinSources, new TestJoinSourceIdGetter(), Collections.emptyList())
119+
.blockingGet());
116120
}
117121

118122
@Test
@@ -129,7 +133,7 @@ void fetchMultipleSpans() {
129133
mockSelectionSet,
130134
SelectionQuery.builder().selectionPath(List.of("pathToSpans", "spans")).build()))
131135
.thenReturn(Stream.of(mock(SelectedField.class), mock(SelectedField.class)));
132-
when(mockFilterRequestBuilder.build(eq(mockRequestContext), eq(SPAN), anySet()))
136+
when(mockFilterRequestBuilder.build(eq(mockRequestContext), eq(SPAN), anyList()))
133137
.thenReturn(Single.just(List.of(mockFilter)));
134138

135139
when(mockResultSetRequestBuilder.build(
@@ -155,7 +159,9 @@ void fetchMultipleSpans() {
155159
.blockingGet();
156160
assertEquals(
157161
expected,
158-
joiner.joinSpans(joinSources, new TestMultipleJoinSourceIdGetter()).blockingGet());
162+
joiner
163+
.joinSpans(joinSources, new TestMultipleJoinSourceIdGetter(), Collections.emptyList())
164+
.blockingGet());
159165
}
160166

161167
private void mockResult(List<Span> spans) {

0 commit comments

Comments
 (0)