Skip to content

Commit 8e9dfae

Browse files
AnandShivanshAnandShivansh
committed
Add SpanJoiner based on list of spanId
1 parent 4df2934 commit 8e9dfae

File tree

4 files changed

+56
-38
lines changed

4 files changed

+56
-38
lines changed

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

Lines changed: 32 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,18 @@
22

33
import static com.google.common.collect.ImmutableList.copyOf;
44
import static com.google.common.collect.Iterables.concat;
5+
import static java.util.Collections.unmodifiableMap;
56
import static org.hypertrace.core.graphql.atttributes.scopes.HypertraceCoreAttributeScopeString.SPAN;
6-
import static org.hypertrace.core.graphql.span.joiner.SpanJoin.SPAN_KEY;
7+
import static org.hypertrace.core.graphql.span.joiner.SpanJoin.SPANS_KEY;
78

89
import graphql.schema.DataFetchingFieldSelectionSet;
910
import graphql.schema.SelectedField;
1011
import io.reactivex.rxjava3.core.Observable;
1112
import io.reactivex.rxjava3.core.Single;
13+
import java.util.ArrayList;
1214
import java.util.Collection;
1315
import java.util.Collections;
16+
import java.util.HashMap;
1417
import java.util.List;
1518
import java.util.Map;
1619
import java.util.Map.Entry;
@@ -77,7 +80,7 @@ public Single<SpanJoiner> build(
7780

7881
private List<SelectedField> getSelections(
7982
DataFetchingFieldSelectionSet selectionSet, List<String> pathToSpanJoin) {
80-
List<String> fullPath = copyOf(concat(pathToSpanJoin, List.of(SPAN_KEY)));
83+
List<String> fullPath = copyOf(concat(pathToSpanJoin, List.of(SPANS_KEY)));
8184
return selectionFinder
8285
.findSelections(selectionSet, SelectionQuery.builder().selectionPath(fullPath).build())
8386
.collect(Collectors.toUnmodifiableList());
@@ -91,34 +94,45 @@ private class DefaultSpanJoiner implements SpanJoiner {
9194
private final List<SelectedField> selectedFields;
9295

9396
@Override
94-
public <T> Single<Map<T, Span>> joinSpans(
97+
public <T> Single<Map<T, Collection<Span>>> joinSpans(
9598
Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
9699
return this.buildSourceToIdMap(joinSources, spanIdGetter).flatMap(this::joinSpans);
97100
}
98101

99-
private <T> Single<Map<T, Span>> joinSpans(Map<T, String> sourceToSpanIdMap) {
100-
return this.buildSpanRequest(sourceToSpanIdMap)
102+
private <T> Single<Map<T, Collection<Span>>> joinSpans(
103+
Map<T, Collection<String>> sourceToSpanIdsMap) {
104+
return this.buildSpanRequest(sourceToSpanIdsMap)
101105
.flatMap(spanDao::getSpans)
102106
.map(this::buildSpanIdToSpanMap)
103-
.map(spanIdToSpanMap -> buildSourceToSpanMap(sourceToSpanIdMap, spanIdToSpanMap));
107+
.map(spanIdToSpanMap -> buildSourceToSpansMap(sourceToSpanIdsMap, spanIdToSpanMap));
104108
}
105109

106-
private <T> Map<T, Span> buildSourceToSpanMap(
107-
Map<T, String> sourceToSpanIdMap, Map<String, Span> spanIdToSpanMap) {
108-
return sourceToSpanIdMap.entrySet().stream()
109-
.filter(entry -> spanIdToSpanMap.containsKey(entry.getValue()))
110-
.collect(
111-
Collectors.toUnmodifiableMap(
112-
Entry::getKey, entry -> spanIdToSpanMap.get(entry.getValue())));
110+
private <T> Map<T, Collection<Span>> buildSourceToSpansMap(
111+
Map<T, Collection<String>> sourceToSpanIdsMap, Map<String, Span> spanIdToSpanMap) {
112+
Map<T, Collection<Span>> sourceToSpansMap = new HashMap<>();
113+
for (Entry<T, Collection<String>> entry : sourceToSpanIdsMap.entrySet()) {
114+
List<Span> spans = new ArrayList<>();
115+
for (String spanId : entry.getValue()) {
116+
if (spanIdToSpanMap.containsKey(spanId)) {
117+
spans.add(spanIdToSpanMap.get(spanId));
118+
}
119+
}
120+
sourceToSpansMap.put(
121+
entry.getKey(), spans.stream().distinct().collect(Collectors.toUnmodifiableList()));
122+
}
123+
return unmodifiableMap(sourceToSpansMap);
113124
}
114125

115126
private Map<String, Span> buildSpanIdToSpanMap(SpanResultSet resultSet) {
116127
return resultSet.results().stream()
117128
.collect(Collectors.toUnmodifiableMap(Identifiable::id, Function.identity()));
118129
}
119130

120-
private <T> Single<SpanRequest> buildSpanRequest(Map<T, String> sourceToSpanIdMap) {
121-
Collection<String> spanIds = sourceToSpanIdMap.values();
131+
private <T> Single<SpanRequest> buildSpanRequest(Map<T, Collection<String>> sourceToSpanIdMap) {
132+
Set<String> spanIds =
133+
sourceToSpanIdMap.values().stream()
134+
.flatMap(Collection::stream)
135+
.collect(Collectors.toUnmodifiableSet());
122136
return buildSpanIdsFilter(spanIds)
123137
.flatMap(filterArguments -> buildSpanRequest(spanIds.size(), filterArguments));
124138
}
@@ -144,16 +158,16 @@ private Single<List<AttributeAssociation<FilterArgument>>> buildSpanIdsFilter(
144158
return filterRequestBuilder.build(context, SPAN, Set.of(new SpanIdFilter(spanIds)));
145159
}
146160

147-
private <T> Single<Map<T, String>> buildSourceToIdMap(
161+
private <T> Single<Map<T, Collection<String>>> buildSourceToIdMap(
148162
Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
149163
return Observable.fromIterable(joinSources)
150164
.flatMapSingle(source -> this.maybeBuildMapEntry(source, spanIdGetter))
151165
.collect(Collectors.toMap(Entry::getKey, Entry::getValue));
152166
}
153167

154-
private <T> Single<Entry<T, String>> maybeBuildMapEntry(
168+
private <T> Single<Entry<T, Collection<String>>> maybeBuildMapEntry(
155169
T source, SpanIdGetter<T> spanIdGetter) {
156-
return spanIdGetter.getSpanId(source).map(id -> Map.entry(source, id));
170+
return spanIdGetter.getSpanIds(source).map(ids -> Map.entry(source, ids));
157171
}
158172
}
159173

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,13 @@
22

33
import graphql.annotations.annotationTypes.GraphQLField;
44
import graphql.annotations.annotationTypes.GraphQLName;
5+
import java.util.List;
56
import org.hypertrace.core.graphql.span.schema.Span;
67

78
public interface SpanJoin {
8-
String SPAN_KEY = "span";
9+
String SPANS_KEY = "spans";
910

1011
@GraphQLField
11-
@GraphQLName(SPAN_KEY)
12-
Span span();
12+
@GraphQLName(SPANS_KEY)
13+
List<Span> spans();
1314
}

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,16 +12,17 @@ public interface SpanJoiner {
1212
SpanJoiner NO_OP_JOINER =
1313
new SpanJoiner() {
1414
@Override
15-
public <T> Single<Map<T, Span>> joinSpans(
15+
public <T> Single<Map<T, Collection<Span>>> joinSpans(
1616
Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
1717
return Single.just(Collections.emptyMap());
1818
}
1919
};
2020

21-
<T> Single<Map<T, Span>> joinSpans(Collection<T> joinSources, SpanIdGetter<T> spanIdGetter);
21+
<T> Single<Map<T, Collection<Span>>> joinSpans(
22+
Collection<T> joinSources, SpanIdGetter<T> spanIdGetter);
2223

2324
@FunctionalInterface
2425
interface SpanIdGetter<T> {
25-
Single<String> getSpanId(T source);
26+
Single<Collection<String>> getSpanIds(T source);
2627
}
2728
}

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

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import graphql.schema.DataFetchingFieldSelectionSet;
1414
import graphql.schema.SelectedField;
1515
import io.reactivex.rxjava3.core.Single;
16+
import java.util.Collection;
1617
import java.util.List;
1718
import java.util.Map;
1819
import java.util.Optional;
@@ -74,13 +75,12 @@ void setup() {
7475
void fetchSpans() {
7576
Span span1 = new TestSpan(FIRST_SPAN_ID);
7677
Span span2 = new TestSpan(SECOND_SPAN_ID);
77-
TestJoinSource joinSource1 = new TestJoinSource(FIRST_SPAN_ID);
78-
TestJoinSource joinSource2 = new TestJoinSource(SECOND_SPAN_ID);
79-
Map<TestJoinSource, Span> expected =
80-
Map.ofEntries(entry(joinSource1, span1), entry(joinSource2, span2));
78+
TestJoinSource joinSource1 = new TestJoinSource(List.of(FIRST_SPAN_ID));
79+
TestJoinSource joinSource2 = new TestJoinSource(List.of(SECOND_SPAN_ID));
80+
Map<TestJoinSource, Collection<Span>> expected =
81+
Map.ofEntries(entry(joinSource1, List.of(span1)), entry(joinSource2, List.of(span2)));
8182
List<TestJoinSource> joinSources = List.of(joinSource1, joinSource2);
82-
mockRequestedSelectionFields(
83-
List.of(mock(SelectedField.class), mock(SelectedField.class)), "pathToSpan");
83+
mockRequestedSelectionFields(List.of(mock(SelectedField.class), mock(SelectedField.class)));
8484
mockRequestBuilding();
8585
mockResult(List.of(span1, span2));
8686
SpanJoiner joiner =
@@ -91,8 +91,10 @@ void fetchSpans() {
9191
this.mockSelectionSet,
9292
List.of("pathToSpan"))
9393
.blockingGet();
94-
assertEquals(
95-
expected, joiner.joinSpans(joinSources, new TestJoinSourceIdGetter()).blockingGet());
94+
Map<TestJoinSource, Collection<Span>> actual =
95+
joiner.joinSpans(joinSources, new TestJoinSourceIdGetter()).blockingGet();
96+
assertEquals(expected.get(joinSource1), actual.get(joinSource1));
97+
assertEquals(expected.get(joinSource2), actual.get(joinSource2));
9698
}
9799

98100
private void mockRequestBuilding() {
@@ -112,10 +114,10 @@ private void mockRequestBuilding() {
112114
.thenReturn(Single.just(mockResultSetRequest));
113115
}
114116

115-
private void mockRequestedSelectionFields(List<SelectedField> selectedFields, String location) {
117+
private void mockRequestedSelectionFields(List<SelectedField> selectedFields) {
116118
when(mockSelectionFinder.findSelections(
117119
mockSelectionSet,
118-
SelectionQuery.builder().selectionPath(List.of(location, "span")).build()))
120+
SelectionQuery.builder().selectionPath(List.of("pathToSpan", "spans")).build()))
119121
.thenReturn(selectedFields.stream());
120122
}
121123

@@ -126,16 +128,16 @@ private void mockResult(List<Span> spans) {
126128

127129
@Value
128130
private static class TestJoinSource {
129-
String spanId;
131+
List<String> spanIds;
130132
}
131133

132134
private static class TestJoinSourceIdGetter implements SpanIdGetter<TestJoinSource> {
133135
@Override
134-
public Single<String> getSpanId(TestJoinSource source) {
135-
if (source.getSpanId() == null || source.getSpanId().isEmpty()) {
136+
public Single<Collection<String>> getSpanIds(TestJoinSource source) {
137+
if (source.getSpanIds() == null || source.getSpanIds().isEmpty()) {
136138
return Single.error(new IllegalArgumentException("Empty spanId"));
137139
}
138-
return Single.just(source.getSpanId());
140+
return Single.just(source.getSpanIds());
139141
}
140142
}
141143

0 commit comments

Comments
 (0)