Skip to content

Commit e6a3738

Browse files
author
Rishabh
authored
feat: Add spanId attribute for span-log join query (#60)
1 parent 287f57e commit e6a3738

File tree

4 files changed

+164
-15
lines changed

4 files changed

+164
-15
lines changed

‎hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/dao/SpanLogEventRequestBuilder.java‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,22 +2,26 @@
22

33
import static io.reactivex.rxjava3.core.Single.zip;
44

5+
import io.reactivex.rxjava3.core.Observable;
56
import io.reactivex.rxjava3.core.Single;
67
import java.util.Collection;
78
import java.util.List;
89
import java.util.Set;
910
import java.util.stream.Collectors;
1011
import javax.inject.Inject;
1112
import lombok.experimental.Accessors;
13+
import org.hypertrace.core.graphql.attributes.AttributeStore;
1214
import org.hypertrace.core.graphql.atttributes.scopes.HypertraceCoreAttributeScopeString;
1315
import org.hypertrace.core.graphql.common.request.AttributeAssociation;
1416
import org.hypertrace.core.graphql.common.request.AttributeRequest;
17+
import org.hypertrace.core.graphql.common.request.AttributeRequestBuilder;
1518
import org.hypertrace.core.graphql.common.request.FilterRequestBuilder;
1619
import org.hypertrace.core.graphql.common.schema.attributes.AttributeScope;
1720
import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterArgument;
1821
import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterOperatorType;
1922
import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterType;
2023
import org.hypertrace.core.graphql.common.utils.Converter;
24+
import org.hypertrace.core.graphql.context.GraphQlRequestContext;
2125
import org.hypertrace.core.graphql.span.request.SpanRequest;
2226
import org.hypertrace.gateway.service.v1.common.Expression;
2327
import org.hypertrace.gateway.service.v1.common.Filter;
@@ -29,21 +33,28 @@ class SpanLogEventRequestBuilder {
2933
private final Converter<Collection<AttributeRequest>, Set<Expression>> attributeConverter;
3034
private final Converter<Collection<AttributeAssociation<FilterArgument>>, Filter> filterConverter;
3135
private final FilterRequestBuilder filterRequestBuilder;
36+
private final AttributeStore attributeStore;
37+
private final AttributeRequestBuilder attributeRequestBuilder;
3238

3339
@Inject
3440
SpanLogEventRequestBuilder(
3541
Converter<Collection<AttributeRequest>, Set<Expression>> attributeConverter,
3642
Converter<Collection<AttributeAssociation<FilterArgument>>, Filter> filterConverter,
37-
FilterRequestBuilder filterRequestBuilder) {
43+
FilterRequestBuilder filterRequestBuilder,
44+
AttributeStore attributeStore,
45+
AttributeRequestBuilder attributeRequestBuilder) {
3846
this.attributeConverter = attributeConverter;
3947
this.filterConverter = filterConverter;
4048
this.filterRequestBuilder = filterRequestBuilder;
49+
this.attributeStore = attributeStore;
50+
this.attributeRequestBuilder = attributeRequestBuilder;
4151
}
4252

4353
Single<LogEventsRequest> buildLogEventsRequest(
4454
SpanRequest gqlRequest, SpansResponse spansResponse) {
4555
return zip(
46-
this.attributeConverter.convert(gqlRequest.logEventAttributes()),
56+
getRequestAttributes(
57+
gqlRequest.spanEventsRequest().context(), gqlRequest.logEventAttributes()),
4758
buildLogEventsQueryFilter(gqlRequest, spansResponse).flatMap(filterConverter::convert),
4859
(selections, filter) ->
4960
LogEventsRequest.newBuilder()
@@ -56,6 +67,20 @@ Single<LogEventsRequest> buildLogEventsRequest(
5667
.build());
5768
}
5869

70+
private Single<Set<Expression>> getRequestAttributes(
71+
GraphQlRequestContext requestContext, Collection<AttributeRequest> logEventAttributes) {
72+
return this.attributeStore
73+
.getForeignIdAttribute(
74+
requestContext,
75+
HypertraceCoreAttributeScopeString.LOG_EVENT,
76+
HypertraceCoreAttributeScopeString.SPAN)
77+
.map(attributeRequestBuilder::buildForAttribute)
78+
.toObservable()
79+
.mergeWith(Observable.fromIterable(logEventAttributes))
80+
.collect(Collectors.toSet())
81+
.flatMap(attributeConverter::convert);
82+
}
83+
5984
private Single<List<AttributeAssociation<FilterArgument>>> buildLogEventsQueryFilter(
6085
SpanRequest gqlRequest, SpansResponse spansResponse) {
6186
List<String> spanIds =

‎hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/dao/SpanLogEventResponseConverter.java‎

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import java.util.stream.Collectors;
88
import javax.inject.Inject;
99
import lombok.experimental.Accessors;
10+
import org.hypertrace.core.graphql.attributes.AttributeModel;
1011
import org.hypertrace.core.graphql.attributes.AttributeStore;
1112
import org.hypertrace.core.graphql.atttributes.scopes.HypertraceCoreAttributeScopeString;
1213
import org.hypertrace.core.graphql.common.request.AttributeRequest;
@@ -43,29 +44,44 @@ Single<SpanLogEventsResponse> buildResponse(
4344
HypertraceCoreAttributeScopeString.LOG_EVENT,
4445
HypertraceCoreAttributeScopeString.SPAN)
4546
.flatMap(
46-
spanId ->
47-
buildResponse(spanId.key(), attributeRequests, spansResponse, logEventsResponse));
47+
spanId -> buildResponse(spanId, attributeRequests, spansResponse, logEventsResponse));
4848
}
4949

5050
private Single<SpanLogEventsResponse> buildResponse(
51-
String foreignIdAttribute,
51+
AttributeModel foreignIdAttribute,
5252
Collection<AttributeRequest> attributeRequests,
5353
SpansResponse spansResponse,
5454
LogEventsResponse logEventsResponse) {
5555
return Observable.fromIterable(logEventsResponse.getLogEventsList())
5656
.concatMapSingle(
57-
logEventsResponseVar -> this.convert(attributeRequests, logEventsResponseVar))
58-
.collect(Collectors.groupingBy(logEvent -> (String) logEvent.attribute(foreignIdAttribute)))
57+
logEventsResponseVar ->
58+
this.convert(foreignIdAttribute, attributeRequests, logEventsResponseVar))
59+
.collect(
60+
Collectors.groupingBy(
61+
SpanLogEventPair::spanId,
62+
Collectors.mapping(SpanLogEventPair::logEvent, Collectors.toList())))
5963
.map(
6064
spanIdVsLogEventsMap -> new SpanLogEventsResponse(spansResponse, spanIdVsLogEventsMap));
6165
}
6266

63-
private Single<LogEvent> convert(
67+
private Single<SpanLogEventPair> convert(
68+
AttributeModel foreignIdAttribute,
6469
Collection<AttributeRequest> request,
6570
org.hypertrace.gateway.service.v1.log.events.LogEvent logEvent) {
6671
return this.attributeMapConverter
6772
.convert(request, logEvent.getAttributesMap())
68-
.map(ConvertedLogEvent::new);
73+
.map(
74+
attributeMap ->
75+
new SpanLogEventPair(
76+
logEvent.getAttributesMap().get(foreignIdAttribute.id()).getString(),
77+
new ConvertedLogEvent(attributeMap)));
78+
}
79+
80+
@lombok.Value
81+
@Accessors(fluent = true)
82+
private static class SpanLogEventPair {
83+
String spanId;
84+
LogEvent logEvent;
6985
}
7086

7187
@lombok.Value

‎hypertrace-core-graphql-span-schema/src/test/java/org/hypertrace/core/graphql/span/dao/SpanLogEventRequestBuilderTest.java‎

Lines changed: 83 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,10 @@
77
import static org.junit.jupiter.api.Assertions.assertEquals;
88
import static org.mockito.ArgumentMatchers.any;
99
import static org.mockito.ArgumentMatchers.anyCollection;
10+
import static org.mockito.ArgumentMatchers.anyString;
1011
import static org.mockito.Mockito.doAnswer;
1112
import static org.mockito.Mockito.mock;
13+
import static org.mockito.Mockito.when;
1214

1315
import com.google.inject.AbstractModule;
1416
import com.google.inject.Guice;
@@ -25,12 +27,16 @@
2527
import java.util.Optional;
2628
import java.util.Set;
2729
import java.util.stream.Collectors;
30+
import org.hypertrace.core.graphql.attributes.AttributeModel;
31+
import org.hypertrace.core.graphql.attributes.AttributeStore;
2832
import org.hypertrace.core.graphql.common.request.AttributeAssociation;
2933
import org.hypertrace.core.graphql.common.request.AttributeRequest;
34+
import org.hypertrace.core.graphql.common.request.AttributeRequestBuilder;
3035
import org.hypertrace.core.graphql.common.request.FilterRequestBuilder;
3136
import org.hypertrace.core.graphql.common.request.ResultSetRequest;
3237
import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterArgument;
3338
import org.hypertrace.core.graphql.common.utils.Converter;
39+
import org.hypertrace.core.graphql.span.dao.DaoTestUtil.DefaultAttributeRequest;
3440
import org.hypertrace.core.graphql.span.dao.DaoTestUtil.DefaultResultSetRequest;
3541
import org.hypertrace.core.graphql.span.dao.DaoTestUtil.DefaultSpanRequest;
3642
import org.hypertrace.core.graphql.span.dao.DaoTestUtil.DefaultTimeRange;
@@ -54,6 +60,10 @@ class SpanLogEventRequestBuilderTest {
5460

5561
@Mock private FilterRequestBuilder filterRequestBuilder;
5662

63+
@Mock private AttributeStore attributeStore;
64+
65+
@Mock private AttributeRequestBuilder attributeRequestBuilder;
66+
5767
private SpanLogEventRequestBuilder spanLogEventRequestBuilder;
5868

5969
@BeforeEach
@@ -82,11 +92,13 @@ protected void configure() {
8292
new TypeLiteral<Converter<Collection<AttributeRequest>, Set<Expression>>>() {}));
8393

8494
spanLogEventRequestBuilder =
85-
new SpanLogEventRequestBuilder(attributeConverter, filterConverter, filterRequestBuilder);
86-
}
95+
new SpanLogEventRequestBuilder(
96+
attributeConverter,
97+
filterConverter,
98+
filterRequestBuilder,
99+
attributeStore,
100+
attributeRequestBuilder);
87101

88-
@Test
89-
void testBuildRequest() {
90102
doAnswer(
91103
invocation -> {
92104
Set<FilterArgument> filterArguments = invocation.getArgument(2, Set.class);
@@ -103,6 +115,21 @@ void testBuildRequest() {
103115
.when(filterRequestBuilder)
104116
.build(any(), any(), anyCollection());
105117

118+
when(attributeStore.getForeignIdAttribute(any(), anyString(), anyString()))
119+
.thenReturn(Single.just(spanIdAttribute.attribute()));
120+
121+
doAnswer(
122+
invocation -> {
123+
AttributeModel attributeModel = invocation.getArgument(0, AttributeModel.class);
124+
return new DefaultAttributeRequest(attributeModel);
125+
})
126+
.when(attributeRequestBuilder)
127+
.buildForAttribute(any());
128+
}
129+
130+
@Test
131+
void testBuildRequest() {
132+
106133
long startTime = System.currentTimeMillis();
107134
long endTime = System.currentTimeMillis() + Duration.ofHours(1).toMillis();
108135

@@ -149,7 +176,58 @@ void testBuildRequest() {
149176
assertEquals(
150177
Set.of("attributes", "traceId", "spanId"),
151178
logEventsRequest.getSelectionList().stream()
152-
.map(v -> v.getColumnIdentifier().getColumnName())
179+
.map(expression -> expression.getColumnIdentifier().getColumnName())
180+
.collect(Collectors.toSet()));
181+
}
182+
183+
@Test
184+
void testBuildRequest_addSpanId() {
185+
long startTime = System.currentTimeMillis();
186+
long endTime = System.currentTimeMillis() + Duration.ofHours(1).toMillis();
187+
188+
Collection<AttributeRequest> logAttributeRequests = List.of(traceIdAttribute);
189+
ResultSetRequest resultSetRequest =
190+
new DefaultResultSetRequest(
191+
null,
192+
List.of(DaoTestUtil.eventIdAttribute),
193+
new DefaultTimeRange(Instant.ofEpochMilli(startTime), Instant.ofEpochMilli(endTime)),
194+
DaoTestUtil.eventIdAttribute,
195+
0,
196+
0,
197+
List.of(),
198+
Collections.emptyList(),
199+
Optional.empty());
200+
SpanRequest spanRequest = new DefaultSpanRequest(resultSetRequest, logAttributeRequests);
201+
202+
LogEventsRequest logEventsRequest =
203+
spanLogEventRequestBuilder.buildLogEventsRequest(spanRequest, spansResponse).blockingGet();
204+
205+
assertEquals(Operator.IN, logEventsRequest.getFilter().getChildFilter(0).getOperator());
206+
assertEquals(
207+
spanIdAttribute.attribute().id(),
208+
logEventsRequest
209+
.getFilter()
210+
.getChildFilter(0)
211+
.getLhs()
212+
.getColumnIdentifier()
213+
.getColumnName());
214+
assertEquals(
215+
List.of("span1", "span2", "span3"),
216+
logEventsRequest
217+
.getFilter()
218+
.getChildFilter(0)
219+
.getRhs()
220+
.getLiteral()
221+
.getValue()
222+
.getStringArrayList()
223+
.stream()
224+
.collect(Collectors.toList()));
225+
assertEquals(startTime, logEventsRequest.getStartTimeMillis());
226+
assertEquals(endTime, logEventsRequest.getEndTimeMillis());
227+
assertEquals(
228+
Set.of("traceId", "spanId"),
229+
logEventsRequest.getSelectionList().stream()
230+
.map(expression -> expression.getColumnIdentifier().getColumnName())
153231
.collect(Collectors.toSet()));
154232
}
155233
}

‎hypertrace-core-graphql-span-schema/src/test/java/org/hypertrace/core/graphql/span/dao/SpanLogEventResponseConverterTest.java‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,37 @@ void testBuildResponse() {
6262
Map<String, Value> map = invocation.getArgument(1, Map.class);
6363
return Single.just(
6464
map.entrySet().stream()
65-
.collect(Collectors.toMap(Entry::getKey, v -> v.getValue().getString())));
65+
.collect(
66+
Collectors.toMap(
67+
Entry::getKey, valueEntry -> valueEntry.getValue().getString())));
68+
})
69+
.when(attributeMapConverter)
70+
.convert(anyCollection(), anyMap());
71+
72+
SpanLogEventsResponse response =
73+
spanLogEventResponseConverter
74+
.buildResponse(requestContext, attributeRequests, spansResponse, logEventsResponse)
75+
.blockingGet();
76+
77+
assertEquals(spansResponse, response.spansResponse());
78+
assertEquals(Set.of("span1", "span2"), response.spanIdToLogEvents().keySet());
79+
}
80+
81+
@Test
82+
void testBuildResponse_spanIdNotRequested() {
83+
Collection<AttributeRequest> attributeRequests = List.of(traceIdAttribute, attributesAttribute);
84+
85+
when(attributeStore.getForeignIdAttribute(any(), anyString(), anyString()))
86+
.thenReturn(Single.just(spanIdAttribute.attribute()));
87+
88+
doAnswer(
89+
invocation -> {
90+
Map<String, Value> map = invocation.getArgument(1, Map.class);
91+
return Single.just(
92+
map.entrySet().stream()
93+
.collect(
94+
Collectors.toMap(
95+
Entry::getKey, valueEntry -> valueEntry.getValue().getString())));
6696
})
6797
.when(attributeMapConverter)
6898
.convert(anyCollection(), anyMap());

0 commit comments

Comments
 (0)