Skip to content

Commit f506940

Browse files
ENG-11241: Changes to fetch labels as part of entity joiner (#102)
1 parent 0c3dd36 commit f506940

12 files changed

Lines changed: 336 additions & 78 deletions

File tree

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/dao/GatewayServiceEntityDao.java

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import io.reactivex.rxjava3.core.Scheduler;
77
import io.reactivex.rxjava3.core.Single;
88
import java.util.Collections;
9+
import java.util.List;
910
import java.util.Map;
1011
import javax.inject.Inject;
1112
import javax.inject.Singleton;
@@ -120,15 +121,17 @@ private Single<Map<Entity, LabelResultSet>> buildLabelResultSetMap(
120121
}
121122

122123
private LabelJoiner.LabelIdGetter<Entity> getEntityLabelsGetter(EntityRequest request) {
123-
return entity -> {
124-
Value labelAttributeValue =
125-
entity.getAttributeOrDefault(
126-
request.labelRequest().get().labelIdArrayAttributeRequest().attribute().id(), null);
127-
if (labelAttributeValue == null) {
128-
log.warn("Unable to fetch labels attribute for entity with id {}", entity.getId());
129-
return Single.just(Collections.emptyList());
130-
}
131-
return Single.just(labelAttributeValue.getStringArrayList());
132-
};
124+
return entity -> Single.just(getLabelAttributeValue(request, entity));
125+
}
126+
127+
private List<String> getLabelAttributeValue(EntityRequest request, Entity entity) {
128+
Value labelAttributeValue =
129+
entity.getAttributeOrDefault(
130+
request.labelRequest().get().labelIdArrayAttributeRequest().attribute().id(), null);
131+
if (labelAttributeValue == null) {
132+
log.warn("Unable to fetch labels attribute for entity with id {}", entity.getId());
133+
return Collections.emptyList();
134+
}
135+
return labelAttributeValue.getStringArrayList();
133136
}
134137
}

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/joiner/DefaultEntityJoinerBuilder.java

Lines changed: 36 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import static com.google.common.collect.ImmutableList.copyOf;
44
import static com.google.common.collect.Iterables.concat;
55
import static com.google.common.collect.Tables.immutableCell;
6+
import static io.reactivex.rxjava3.core.Single.zip;
67
import static java.util.Objects.requireNonNull;
78

89
import com.google.common.collect.ImmutableSetMultimap;
@@ -35,6 +36,7 @@
3536
import lombok.extern.slf4j.Slf4j;
3637
import org.hypertrace.core.graphql.common.request.AttributeAssociation;
3738
import org.hypertrace.core.graphql.common.request.AttributeRequest;
39+
import org.hypertrace.core.graphql.common.request.AttributeRequestBuilder;
3840
import org.hypertrace.core.graphql.common.request.FilterRequestBuilder;
3941
import org.hypertrace.core.graphql.common.request.ResultSetRequest;
4042
import org.hypertrace.core.graphql.common.request.ResultSetRequestBuilder;
@@ -50,8 +52,9 @@
5052
import org.hypertrace.core.graphql.utils.schema.SelectionQuery;
5153
import org.hypertrace.graphql.entity.dao.EntityDao;
5254
import org.hypertrace.graphql.entity.request.EdgeSetGroupRequest;
55+
import org.hypertrace.graphql.entity.request.EntityLabelRequest;
56+
import org.hypertrace.graphql.entity.request.EntityLabelRequestBuilder;
5357
import org.hypertrace.graphql.entity.request.EntityRequest;
54-
import org.hypertrace.graphql.entity.request.LabelRequest;
5558
import org.hypertrace.graphql.entity.schema.Entity;
5659
import org.hypertrace.graphql.entity.schema.EntityJoinable;
5760
import org.hypertrace.graphql.entity.schema.EntityResultSet;
@@ -62,15 +65,16 @@
6265

6366
@Slf4j
6467
class DefaultEntityJoinerBuilder implements EntityJoinerBuilder {
65-
6668
private static final int ZERO_OFFSET = 0;
6769

6870
private final EntityDao entityDao;
6971
private final GraphQlSelectionFinder selectionFinder;
7072
private final ArgumentDeserializer argumentDeserializer;
7173
private final ResultSetRequestBuilder resultSetRequestBuilder;
7274
private final FilterRequestBuilder filterRequestBuilder;
75+
private final AttributeRequestBuilder attributeRequestBuilder;
7376
private final Scheduler boundedIoScheduler;
77+
private final EntityLabelRequestBuilder entityLabelRequestBuilder;
7478

7579
@Inject
7680
DefaultEntityJoinerBuilder(
@@ -79,14 +83,18 @@ class DefaultEntityJoinerBuilder implements EntityJoinerBuilder {
7983
ArgumentDeserializer argumentDeserializer,
8084
ResultSetRequestBuilder resultSetRequestBuilder,
8185
FilterRequestBuilder filterRequestBuilder,
82-
@BoundedIoScheduler Scheduler boundedIoScheduler) {
86+
AttributeRequestBuilder attributeRequestBuilder,
87+
@BoundedIoScheduler Scheduler boundedIoScheduler,
88+
EntityLabelRequestBuilder entityLabelRequestBuilder) {
8389

8490
this.entityDao = entityDao;
8591
this.selectionFinder = selectionFinder;
8692
this.argumentDeserializer = argumentDeserializer;
8793
this.resultSetRequestBuilder = resultSetRequestBuilder;
8894
this.filterRequestBuilder = filterRequestBuilder;
95+
this.attributeRequestBuilder = attributeRequestBuilder;
8996
this.boundedIoScheduler = boundedIoScheduler;
97+
this.entityLabelRequestBuilder = entityLabelRequestBuilder;
9098
}
9199

92100
@Override
@@ -241,19 +249,30 @@ Single<EntityRequest> buildEntityRequestForType(
241249
return buildIdFilter(context, entityType, entityIdsToFilter)
242250
.flatMap(
243251
filterArguments ->
244-
resultSetRequestBuilder.build(
245-
context,
246-
entityType,
247-
entityIdsToFilter.size(),
248-
ZERO_OFFSET,
249-
new InstantTimeRange(),
250-
List
251-
.<AttributeAssociation<AggregatableOrderArgument>>
252-
of(), // Order does not matter
253-
filterArguments,
254-
this.entityFieldsByType.get(entityType).stream(),
255-
Optional.empty()))
256-
.map(resultSetRequest -> new DefaultEntityRequest(entityType, resultSetRequest));
252+
buildEntityRequest(
253+
context, entityType, entityIdsToFilter.size(), filterArguments));
254+
}
255+
256+
private Single<EntityRequest> buildEntityRequest(
257+
GraphQlRequestContext context,
258+
String entityType,
259+
int entityIdsToFilterSize,
260+
List<AttributeAssociation<FilterArgument>> filterArguments) {
261+
return zip(
262+
resultSetRequestBuilder.build(
263+
context,
264+
entityType,
265+
entityIdsToFilterSize,
266+
ZERO_OFFSET,
267+
new InstantTimeRange(),
268+
List.<AttributeAssociation<AggregatableOrderArgument>>of(), // Order does not matter
269+
filterArguments,
270+
this.entityFieldsByType.get(entityType).stream(),
271+
Optional.empty()),
272+
entityLabelRequestBuilder.buildLabelRequestIfPresentInAnyEntity(
273+
context, entityType, this.entityFieldsByType.get(entityType)),
274+
(resultSetRequest, optionalLabelRequest) ->
275+
new DefaultEntityRequest(entityType, resultSetRequest, optionalLabelRequest));
257276
}
258277

259278
private Single<List<AttributeAssociation<FilterArgument>>> buildIdFilter(
@@ -273,7 +292,7 @@ private static class DefaultEntityRequest implements EntityRequest {
273292
EdgeSetGroupRequest incomingEdgeRequests = new EmptyEdgeSetGroupRequest();
274293
EdgeSetGroupRequest outgoingEdgeRequests = new EmptyEdgeSetGroupRequest();
275294
boolean includeInactive = true; // When joining we want the entity regardless of time range
276-
Optional<LabelRequest> labelRequest = Optional.empty();
295+
Optional<EntityLabelRequest> labelRequest;
277296
}
278297

279298
@Value
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
package org.hypertrace.graphql.entity.request;
2+
3+
import graphql.schema.DataFetchingFieldSelectionSet;
4+
import graphql.schema.SelectedField;
5+
import io.reactivex.rxjava3.core.Single;
6+
import java.util.Collection;
7+
import java.util.List;
8+
import java.util.Optional;
9+
import javax.inject.Inject;
10+
import lombok.Value;
11+
import lombok.experimental.Accessors;
12+
import org.hypertrace.core.graphql.common.request.AttributeRequest;
13+
import org.hypertrace.core.graphql.common.request.AttributeRequestBuilder;
14+
import org.hypertrace.core.graphql.common.schema.results.ResultSet;
15+
import org.hypertrace.core.graphql.context.GraphQlRequestContext;
16+
import org.hypertrace.core.graphql.utils.schema.GraphQlSelectionFinder;
17+
import org.hypertrace.core.graphql.utils.schema.SelectionQuery;
18+
import org.hypertrace.graphql.entity.schema.Entity;
19+
20+
class DefaultEntityLabelRequestBuilder implements EntityLabelRequestBuilder {
21+
22+
private static final String LABELS_ATTRIBUTE_KEY = "labels";
23+
24+
private static final SelectionQuery RESULT_SET_LABEL_FIELD_QUERY =
25+
SelectionQuery.builder()
26+
.selectionPath(List.of(ResultSet.RESULT_SET_RESULTS_NAME, Entity.LABELS_KEY))
27+
.build();
28+
29+
private static final SelectionQuery LABEL_FIELD_QUERY =
30+
SelectionQuery.namedChild(Entity.LABELS_KEY);
31+
32+
private final AttributeRequestBuilder attributeRequestBuilder;
33+
private final GraphQlSelectionFinder selectionFinder;
34+
35+
@Inject
36+
DefaultEntityLabelRequestBuilder(
37+
AttributeRequestBuilder attributeRequestBuilder, GraphQlSelectionFinder selectionFinder) {
38+
this.attributeRequestBuilder = attributeRequestBuilder;
39+
this.selectionFinder = selectionFinder;
40+
}
41+
42+
@Override
43+
public Single<Optional<EntityLabelRequest>> buildLabelRequestIfPresentInResultSet(
44+
GraphQlRequestContext context, String scope, DataFetchingFieldSelectionSet selectionSet) {
45+
if (isLabelFieldRequested(selectionSet, RESULT_SET_LABEL_FIELD_QUERY)) {
46+
return buildRequest(context, scope);
47+
}
48+
return Single.just(Optional.empty());
49+
}
50+
51+
@Override
52+
public Single<Optional<EntityLabelRequest>> buildLabelRequestIfPresentInAnyEntity(
53+
GraphQlRequestContext context, String scope, Collection<SelectedField> entityFields) {
54+
if (entityFields.stream()
55+
.anyMatch(field -> isLabelFieldRequested(field.getSelectionSet(), LABEL_FIELD_QUERY))) {
56+
return buildRequest(context, scope);
57+
}
58+
return Single.just(Optional.empty());
59+
}
60+
61+
private boolean isLabelFieldRequested(
62+
DataFetchingFieldSelectionSet selectionSet, SelectionQuery labelFieldQuery) {
63+
return this.selectionFinder.findSelections(selectionSet, labelFieldQuery).findAny().isPresent();
64+
}
65+
66+
private Single<Optional<EntityLabelRequest>> buildRequest(
67+
GraphQlRequestContext context, String scope) {
68+
return this.attributeRequestBuilder
69+
.buildForKey(context, scope, LABELS_ATTRIBUTE_KEY)
70+
.map(DefaultLabelRequest::new)
71+
.map(Optional::of);
72+
}
73+
74+
@Value
75+
@Accessors(fluent = true)
76+
private static class DefaultLabelRequest implements EntityLabelRequest {
77+
AttributeRequest labelIdArrayAttributeRequest;
78+
}
79+
}

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/DefaultEntityRequestBuilder.java

Lines changed: 8 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,6 @@
1414
import javax.inject.Inject;
1515
import lombok.Value;
1616
import lombok.experimental.Accessors;
17-
import org.hypertrace.core.graphql.common.request.AttributeRequest;
18-
import org.hypertrace.core.graphql.common.request.AttributeRequestBuilder;
1917
import org.hypertrace.core.graphql.common.request.ResultSetRequest;
2018
import org.hypertrace.core.graphql.common.request.ResultSetRequestBuilder;
2119
import org.hypertrace.core.graphql.common.schema.arguments.TimeRangeArgument;
@@ -25,7 +23,6 @@
2523
import org.hypertrace.core.graphql.deserialization.ArgumentDeserializer;
2624
import org.hypertrace.core.graphql.utils.schema.GraphQlSelectionFinder;
2725
import org.hypertrace.core.graphql.utils.schema.SelectionQuery;
28-
import org.hypertrace.graphql.entity.schema.Entity;
2926
import org.hypertrace.graphql.entity.schema.EntityType;
3027
import org.hypertrace.graphql.entity.schema.argument.EntityScopeArgument;
3128
import org.hypertrace.graphql.entity.schema.argument.EntityTypeArgument;
@@ -36,14 +33,12 @@
3633

3734
class DefaultEntityRequestBuilder implements EntityRequestBuilder {
3835

39-
private static final String LABELS_KEY_NAME = "labels";
40-
4136
private final ResultSetRequestBuilder resultSetRequestBuilder;
4237
private final MetricRequestBuilder metricRequestBuilder;
4338
private final ArgumentDeserializer argumentDeserializer;
4439
private final GraphQlSelectionFinder selectionFinder;
4540
private final EdgeRequestBuilder edgeRequestBuilder;
46-
private final AttributeRequestBuilder attributeRequestBuilder;
41+
private final EntityLabelRequestBuilder entityLabelRequestBuilder;
4742

4843
@Inject
4944
DefaultEntityRequestBuilder(
@@ -52,13 +47,13 @@ class DefaultEntityRequestBuilder implements EntityRequestBuilder {
5247
ArgumentDeserializer argumentDeserializer,
5348
GraphQlSelectionFinder selectionFinder,
5449
EdgeRequestBuilder edgeRequestBuilder,
55-
AttributeRequestBuilder attributeRequestBuilder) {
50+
EntityLabelRequestBuilder entityLabelRequestBuilder) {
5651
this.resultSetRequestBuilder = resultSetRequestBuilder;
5752
this.metricRequestBuilder = metricRequestBuilder;
5853
this.argumentDeserializer = argumentDeserializer;
5954
this.selectionFinder = selectionFinder;
6055
this.edgeRequestBuilder = edgeRequestBuilder;
61-
this.attributeRequestBuilder = attributeRequestBuilder;
56+
this.entityLabelRequestBuilder = entityLabelRequestBuilder;
6257
}
6358

6459
@Override
@@ -110,12 +105,13 @@ private Single<EntityRequest> build(
110105
this.timeRange(arguments),
111106
this.space(arguments),
112107
this.getOutgoingEdges(selectionSet)),
113-
buildLabelRequest(context, scope, selectionSet),
108+
this.entityLabelRequestBuilder.buildLabelRequestIfPresentInResultSet(
109+
context, scope, selectionSet),
114110
(resultSetRequest,
115111
metricRequestList,
116112
incomingEdges,
117113
outgoingEdges,
118-
labelsAttributeRequest) ->
114+
optionalLabelsAttributeRequest) ->
119115
new DefaultEntityRequest(
120116
scope,
121117
resultSetRequest,
@@ -124,28 +120,7 @@ private Single<EntityRequest> build(
124120
outgoingEdges,
125121
includeInactive,
126122
fetchTotal,
127-
labelsAttributeRequest));
128-
}
129-
130-
private Single<Optional<LabelRequest>> buildLabelRequest(
131-
GraphQlRequestContext context, String scope, DataFetchingFieldSelectionSet selectionSet) {
132-
boolean canFetchLabels =
133-
this.selectionFinder
134-
.findSelections(
135-
selectionSet,
136-
SelectionQuery.builder()
137-
.selectionPath(
138-
List.of(ResultSet.RESULT_SET_RESULTS_NAME, Entity.LABELS_KEY))
139-
.build())
140-
.count()
141-
> 0;
142-
143-
return canFetchLabels
144-
? attributeRequestBuilder
145-
.buildForKey(context, scope, LABELS_KEY_NAME)
146-
.map(DefaultLabelRequest::new)
147-
.map(Optional::of)
148-
: Single.just(Optional.empty());
123+
optionalLabelsAttributeRequest));
149124
}
150125

151126
private Stream<SelectedField> getResultSets(DataFetchingFieldSelectionSet selectionSet) {
@@ -189,12 +164,6 @@ private static class DefaultEntityRequest implements EntityRequest {
189164
EdgeSetGroupRequest outgoingEdgeRequests;
190165
boolean includeInactive;
191166
boolean fetchTotal;
192-
Optional<LabelRequest> labelRequest;
193-
}
194-
195-
@Value
196-
@Accessors(fluent = true)
197-
private static class DefaultLabelRequest implements LabelRequest {
198-
AttributeRequest labelIdArrayAttributeRequest;
167+
Optional<EntityLabelRequest> labelRequest;
199168
}
200169
}

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/LabelRequest.java renamed to hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityLabelRequest.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
import org.hypertrace.core.graphql.common.request.AttributeRequest;
44

5-
public interface LabelRequest {
5+
public interface EntityLabelRequest {
66

77
AttributeRequest labelIdArrayAttributeRequest();
88
}
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
package org.hypertrace.graphql.entity.request;
2+
3+
import graphql.schema.DataFetchingFieldSelectionSet;
4+
import graphql.schema.SelectedField;
5+
import io.reactivex.rxjava3.core.Single;
6+
import java.util.Collection;
7+
import java.util.Optional;
8+
import org.hypertrace.core.graphql.context.GraphQlRequestContext;
9+
10+
public interface EntityLabelRequestBuilder {
11+
12+
Single<Optional<EntityLabelRequest>> buildLabelRequestIfPresentInResultSet(
13+
GraphQlRequestContext context, String scope, DataFetchingFieldSelectionSet selectionSet);
14+
15+
Single<Optional<EntityLabelRequest>> buildLabelRequestIfPresentInAnyEntity(
16+
GraphQlRequestContext context, String scope, Collection<SelectedField> entityField);
17+
}

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityRequest.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,5 +21,5 @@ public interface EntityRequest {
2121

2222
boolean fetchTotal();
2323

24-
Optional<LabelRequest> labelRequest();
24+
Optional<EntityLabelRequest> labelRequest();
2525
}

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityRequestModule.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ public class EntityRequestModule extends AbstractModule {
1414
@Override
1515
protected void configure() {
1616
bind(EntityRequestBuilder.class).to(DefaultEntityRequestBuilder.class);
17+
bind(EntityLabelRequestBuilder.class).to(DefaultEntityLabelRequestBuilder.class);
1718
requireBinding(ResultSetRequestBuilder.class);
1819
requireBinding(ArgumentDeserializer.class);
1920
requireBinding(MetricRequestBuilder.class);

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/NeighborEntitiesRequestBuilder.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,6 @@ private static class NeighborEntityRequest implements EntityRequest {
179179
EdgeSetGroupRequest outgoingEdgeRequests;
180180
boolean includeInactive;
181181
boolean fetchTotal;
182-
Optional<LabelRequest> labelRequest;
182+
Optional<EntityLabelRequest> labelRequest;
183183
}
184184
}

0 commit comments

Comments
 (0)