diff --git a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/dao/GatewayServiceEntityDao.java b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/dao/GatewayServiceEntityDao.java index 2cca0dd1..dba07312 100644 --- a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/dao/GatewayServiceEntityDao.java +++ b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/dao/GatewayServiceEntityDao.java @@ -6,6 +6,7 @@ import io.reactivex.rxjava3.core.Scheduler; import io.reactivex.rxjava3.core.Single; import java.util.Collections; +import java.util.List; import java.util.Map; import javax.inject.Inject; import javax.inject.Singleton; @@ -120,15 +121,17 @@ private Single> buildLabelResultSetMap( } private LabelJoiner.LabelIdGetter getEntityLabelsGetter(EntityRequest request) { - return entity -> { - Value labelAttributeValue = - entity.getAttributeOrDefault( - request.labelRequest().get().labelIdArrayAttributeRequest().attribute().id(), null); - if (labelAttributeValue == null) { - log.warn("Unable to fetch labels attribute for entity with id {}", entity.getId()); - return Single.just(Collections.emptyList()); - } - return Single.just(labelAttributeValue.getStringArrayList()); - }; + return entity -> Single.just(getLabelAttributeValue(request, entity)); + } + + private List getLabelAttributeValue(EntityRequest request, Entity entity) { + Value labelAttributeValue = + entity.getAttributeOrDefault( + request.labelRequest().get().labelIdArrayAttributeRequest().attribute().id(), null); + if (labelAttributeValue == null) { + log.warn("Unable to fetch labels attribute for entity with id {}", entity.getId()); + return Collections.emptyList(); + } + return labelAttributeValue.getStringArrayList(); } } diff --git a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/joiner/DefaultEntityJoinerBuilder.java b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/joiner/DefaultEntityJoinerBuilder.java index 42021fed..ed563004 100644 --- a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/joiner/DefaultEntityJoinerBuilder.java +++ b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/joiner/DefaultEntityJoinerBuilder.java @@ -3,6 +3,7 @@ import static com.google.common.collect.ImmutableList.copyOf; import static com.google.common.collect.Iterables.concat; import static com.google.common.collect.Tables.immutableCell; +import static io.reactivex.rxjava3.core.Single.zip; import static java.util.Objects.requireNonNull; import com.google.common.collect.ImmutableSetMultimap; @@ -35,6 +36,7 @@ import lombok.extern.slf4j.Slf4j; import org.hypertrace.core.graphql.common.request.AttributeAssociation; import org.hypertrace.core.graphql.common.request.AttributeRequest; +import org.hypertrace.core.graphql.common.request.AttributeRequestBuilder; import org.hypertrace.core.graphql.common.request.FilterRequestBuilder; import org.hypertrace.core.graphql.common.request.ResultSetRequest; import org.hypertrace.core.graphql.common.request.ResultSetRequestBuilder; @@ -50,8 +52,9 @@ import org.hypertrace.core.graphql.utils.schema.SelectionQuery; import org.hypertrace.graphql.entity.dao.EntityDao; import org.hypertrace.graphql.entity.request.EdgeSetGroupRequest; +import org.hypertrace.graphql.entity.request.EntityLabelRequest; +import org.hypertrace.graphql.entity.request.EntityLabelRequestBuilder; import org.hypertrace.graphql.entity.request.EntityRequest; -import org.hypertrace.graphql.entity.request.LabelRequest; import org.hypertrace.graphql.entity.schema.Entity; import org.hypertrace.graphql.entity.schema.EntityJoinable; import org.hypertrace.graphql.entity.schema.EntityResultSet; @@ -62,7 +65,6 @@ @Slf4j class DefaultEntityJoinerBuilder implements EntityJoinerBuilder { - private static final int ZERO_OFFSET = 0; private final EntityDao entityDao; @@ -70,7 +72,9 @@ class DefaultEntityJoinerBuilder implements EntityJoinerBuilder { private final ArgumentDeserializer argumentDeserializer; private final ResultSetRequestBuilder resultSetRequestBuilder; private final FilterRequestBuilder filterRequestBuilder; + private final AttributeRequestBuilder attributeRequestBuilder; private final Scheduler boundedIoScheduler; + private final EntityLabelRequestBuilder entityLabelRequestBuilder; @Inject DefaultEntityJoinerBuilder( @@ -79,14 +83,18 @@ class DefaultEntityJoinerBuilder implements EntityJoinerBuilder { ArgumentDeserializer argumentDeserializer, ResultSetRequestBuilder resultSetRequestBuilder, FilterRequestBuilder filterRequestBuilder, - @BoundedIoScheduler Scheduler boundedIoScheduler) { + AttributeRequestBuilder attributeRequestBuilder, + @BoundedIoScheduler Scheduler boundedIoScheduler, + EntityLabelRequestBuilder entityLabelRequestBuilder) { this.entityDao = entityDao; this.selectionFinder = selectionFinder; this.argumentDeserializer = argumentDeserializer; this.resultSetRequestBuilder = resultSetRequestBuilder; this.filterRequestBuilder = filterRequestBuilder; + this.attributeRequestBuilder = attributeRequestBuilder; this.boundedIoScheduler = boundedIoScheduler; + this.entityLabelRequestBuilder = entityLabelRequestBuilder; } @Override @@ -241,19 +249,30 @@ Single buildEntityRequestForType( return buildIdFilter(context, entityType, entityIdsToFilter) .flatMap( filterArguments -> - resultSetRequestBuilder.build( - context, - entityType, - entityIdsToFilter.size(), - ZERO_OFFSET, - new InstantTimeRange(), - List - .> - of(), // Order does not matter - filterArguments, - this.entityFieldsByType.get(entityType).stream(), - Optional.empty())) - .map(resultSetRequest -> new DefaultEntityRequest(entityType, resultSetRequest)); + buildEntityRequest( + context, entityType, entityIdsToFilter.size(), filterArguments)); + } + + private Single buildEntityRequest( + GraphQlRequestContext context, + String entityType, + int entityIdsToFilterSize, + List> filterArguments) { + return zip( + resultSetRequestBuilder.build( + context, + entityType, + entityIdsToFilterSize, + ZERO_OFFSET, + new InstantTimeRange(), + List.>of(), // Order does not matter + filterArguments, + this.entityFieldsByType.get(entityType).stream(), + Optional.empty()), + entityLabelRequestBuilder.buildLabelRequestIfPresentInAnyEntity( + context, entityType, this.entityFieldsByType.get(entityType)), + (resultSetRequest, optionalLabelRequest) -> + new DefaultEntityRequest(entityType, resultSetRequest, optionalLabelRequest)); } private Single>> buildIdFilter( @@ -273,7 +292,7 @@ private static class DefaultEntityRequest implements EntityRequest { EdgeSetGroupRequest incomingEdgeRequests = new EmptyEdgeSetGroupRequest(); EdgeSetGroupRequest outgoingEdgeRequests = new EmptyEdgeSetGroupRequest(); boolean includeInactive = true; // When joining we want the entity regardless of time range - Optional labelRequest = Optional.empty(); + Optional labelRequest; } @Value diff --git a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/DefaultEntityLabelRequestBuilder.java b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/DefaultEntityLabelRequestBuilder.java new file mode 100644 index 00000000..bf8aa986 --- /dev/null +++ b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/DefaultEntityLabelRequestBuilder.java @@ -0,0 +1,79 @@ +package org.hypertrace.graphql.entity.request; + +import graphql.schema.DataFetchingFieldSelectionSet; +import graphql.schema.SelectedField; +import io.reactivex.rxjava3.core.Single; +import java.util.Collection; +import java.util.List; +import java.util.Optional; +import javax.inject.Inject; +import lombok.Value; +import lombok.experimental.Accessors; +import org.hypertrace.core.graphql.common.request.AttributeRequest; +import org.hypertrace.core.graphql.common.request.AttributeRequestBuilder; +import org.hypertrace.core.graphql.common.schema.results.ResultSet; +import org.hypertrace.core.graphql.context.GraphQlRequestContext; +import org.hypertrace.core.graphql.utils.schema.GraphQlSelectionFinder; +import org.hypertrace.core.graphql.utils.schema.SelectionQuery; +import org.hypertrace.graphql.entity.schema.Entity; + +class DefaultEntityLabelRequestBuilder implements EntityLabelRequestBuilder { + + private static final String LABELS_ATTRIBUTE_KEY = "labels"; + + private static final SelectionQuery RESULT_SET_LABEL_FIELD_QUERY = + SelectionQuery.builder() + .selectionPath(List.of(ResultSet.RESULT_SET_RESULTS_NAME, Entity.LABELS_KEY)) + .build(); + + private static final SelectionQuery LABEL_FIELD_QUERY = + SelectionQuery.namedChild(Entity.LABELS_KEY); + + private final AttributeRequestBuilder attributeRequestBuilder; + private final GraphQlSelectionFinder selectionFinder; + + @Inject + DefaultEntityLabelRequestBuilder( + AttributeRequestBuilder attributeRequestBuilder, GraphQlSelectionFinder selectionFinder) { + this.attributeRequestBuilder = attributeRequestBuilder; + this.selectionFinder = selectionFinder; + } + + @Override + public Single> buildLabelRequestIfPresentInResultSet( + GraphQlRequestContext context, String scope, DataFetchingFieldSelectionSet selectionSet) { + if (isLabelFieldRequested(selectionSet, RESULT_SET_LABEL_FIELD_QUERY)) { + return buildRequest(context, scope); + } + return Single.just(Optional.empty()); + } + + @Override + public Single> buildLabelRequestIfPresentInAnyEntity( + GraphQlRequestContext context, String scope, Collection entityFields) { + if (entityFields.stream() + .anyMatch(field -> isLabelFieldRequested(field.getSelectionSet(), LABEL_FIELD_QUERY))) { + return buildRequest(context, scope); + } + return Single.just(Optional.empty()); + } + + private boolean isLabelFieldRequested( + DataFetchingFieldSelectionSet selectionSet, SelectionQuery labelFieldQuery) { + return this.selectionFinder.findSelections(selectionSet, labelFieldQuery).findAny().isPresent(); + } + + private Single> buildRequest( + GraphQlRequestContext context, String scope) { + return this.attributeRequestBuilder + .buildForKey(context, scope, LABELS_ATTRIBUTE_KEY) + .map(DefaultLabelRequest::new) + .map(Optional::of); + } + + @Value + @Accessors(fluent = true) + private static class DefaultLabelRequest implements EntityLabelRequest { + AttributeRequest labelIdArrayAttributeRequest; + } +} diff --git a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/DefaultEntityRequestBuilder.java b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/DefaultEntityRequestBuilder.java index 3ba026ed..39a9e154 100644 --- a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/DefaultEntityRequestBuilder.java +++ b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/DefaultEntityRequestBuilder.java @@ -14,8 +14,6 @@ import javax.inject.Inject; import lombok.Value; import lombok.experimental.Accessors; -import org.hypertrace.core.graphql.common.request.AttributeRequest; -import org.hypertrace.core.graphql.common.request.AttributeRequestBuilder; 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; @@ -25,7 +23,6 @@ import org.hypertrace.core.graphql.deserialization.ArgumentDeserializer; import org.hypertrace.core.graphql.utils.schema.GraphQlSelectionFinder; import org.hypertrace.core.graphql.utils.schema.SelectionQuery; -import org.hypertrace.graphql.entity.schema.Entity; import org.hypertrace.graphql.entity.schema.EntityType; import org.hypertrace.graphql.entity.schema.argument.EntityScopeArgument; import org.hypertrace.graphql.entity.schema.argument.EntityTypeArgument; @@ -36,14 +33,12 @@ class DefaultEntityRequestBuilder implements EntityRequestBuilder { - private static final String LABELS_KEY_NAME = "labels"; - private final ResultSetRequestBuilder resultSetRequestBuilder; private final MetricRequestBuilder metricRequestBuilder; private final ArgumentDeserializer argumentDeserializer; private final GraphQlSelectionFinder selectionFinder; private final EdgeRequestBuilder edgeRequestBuilder; - private final AttributeRequestBuilder attributeRequestBuilder; + private final EntityLabelRequestBuilder entityLabelRequestBuilder; @Inject DefaultEntityRequestBuilder( @@ -52,13 +47,13 @@ class DefaultEntityRequestBuilder implements EntityRequestBuilder { ArgumentDeserializer argumentDeserializer, GraphQlSelectionFinder selectionFinder, EdgeRequestBuilder edgeRequestBuilder, - AttributeRequestBuilder attributeRequestBuilder) { + EntityLabelRequestBuilder entityLabelRequestBuilder) { this.resultSetRequestBuilder = resultSetRequestBuilder; this.metricRequestBuilder = metricRequestBuilder; this.argumentDeserializer = argumentDeserializer; this.selectionFinder = selectionFinder; this.edgeRequestBuilder = edgeRequestBuilder; - this.attributeRequestBuilder = attributeRequestBuilder; + this.entityLabelRequestBuilder = entityLabelRequestBuilder; } @Override @@ -110,12 +105,13 @@ private Single build( this.timeRange(arguments), this.space(arguments), this.getOutgoingEdges(selectionSet)), - buildLabelRequest(context, scope, selectionSet), + this.entityLabelRequestBuilder.buildLabelRequestIfPresentInResultSet( + context, scope, selectionSet), (resultSetRequest, metricRequestList, incomingEdges, outgoingEdges, - labelsAttributeRequest) -> + optionalLabelsAttributeRequest) -> new DefaultEntityRequest( scope, resultSetRequest, @@ -124,28 +120,7 @@ private Single build( outgoingEdges, includeInactive, fetchTotal, - labelsAttributeRequest)); - } - - private Single> buildLabelRequest( - GraphQlRequestContext context, String scope, DataFetchingFieldSelectionSet selectionSet) { - boolean canFetchLabels = - this.selectionFinder - .findSelections( - selectionSet, - SelectionQuery.builder() - .selectionPath( - List.of(ResultSet.RESULT_SET_RESULTS_NAME, Entity.LABELS_KEY)) - .build()) - .count() - > 0; - - return canFetchLabels - ? attributeRequestBuilder - .buildForKey(context, scope, LABELS_KEY_NAME) - .map(DefaultLabelRequest::new) - .map(Optional::of) - : Single.just(Optional.empty()); + optionalLabelsAttributeRequest)); } private Stream getResultSets(DataFetchingFieldSelectionSet selectionSet) { @@ -189,12 +164,6 @@ private static class DefaultEntityRequest implements EntityRequest { EdgeSetGroupRequest outgoingEdgeRequests; boolean includeInactive; boolean fetchTotal; - Optional labelRequest; - } - - @Value - @Accessors(fluent = true) - private static class DefaultLabelRequest implements LabelRequest { - AttributeRequest labelIdArrayAttributeRequest; + Optional labelRequest; } } diff --git a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/LabelRequest.java b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityLabelRequest.java similarity index 81% rename from hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/LabelRequest.java rename to hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityLabelRequest.java index 641e1bf1..9da61a18 100644 --- a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/LabelRequest.java +++ b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityLabelRequest.java @@ -2,7 +2,7 @@ import org.hypertrace.core.graphql.common.request.AttributeRequest; -public interface LabelRequest { +public interface EntityLabelRequest { AttributeRequest labelIdArrayAttributeRequest(); } diff --git a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityLabelRequestBuilder.java b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityLabelRequestBuilder.java new file mode 100644 index 00000000..dcd28785 --- /dev/null +++ b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityLabelRequestBuilder.java @@ -0,0 +1,17 @@ +package org.hypertrace.graphql.entity.request; + +import graphql.schema.DataFetchingFieldSelectionSet; +import graphql.schema.SelectedField; +import io.reactivex.rxjava3.core.Single; +import java.util.Collection; +import java.util.Optional; +import org.hypertrace.core.graphql.context.GraphQlRequestContext; + +public interface EntityLabelRequestBuilder { + + Single> buildLabelRequestIfPresentInResultSet( + GraphQlRequestContext context, String scope, DataFetchingFieldSelectionSet selectionSet); + + Single> buildLabelRequestIfPresentInAnyEntity( + GraphQlRequestContext context, String scope, Collection entityField); +} diff --git a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityRequest.java b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityRequest.java index b3a2048d..7f25cf51 100644 --- a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityRequest.java +++ b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityRequest.java @@ -21,5 +21,5 @@ public interface EntityRequest { boolean fetchTotal(); - Optional labelRequest(); + Optional labelRequest(); } diff --git a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityRequestModule.java b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityRequestModule.java index 85cce0d0..481678f9 100644 --- a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityRequestModule.java +++ b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/EntityRequestModule.java @@ -14,6 +14,7 @@ public class EntityRequestModule extends AbstractModule { @Override protected void configure() { bind(EntityRequestBuilder.class).to(DefaultEntityRequestBuilder.class); + bind(EntityLabelRequestBuilder.class).to(DefaultEntityLabelRequestBuilder.class); requireBinding(ResultSetRequestBuilder.class); requireBinding(ArgumentDeserializer.class); requireBinding(MetricRequestBuilder.class); diff --git a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/NeighborEntitiesRequestBuilder.java b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/NeighborEntitiesRequestBuilder.java index 71e48f4a..be59a4b5 100644 --- a/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/NeighborEntitiesRequestBuilder.java +++ b/hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/request/NeighborEntitiesRequestBuilder.java @@ -179,6 +179,6 @@ private static class NeighborEntityRequest implements EntityRequest { EdgeSetGroupRequest outgoingEdgeRequests; boolean includeInactive; boolean fetchTotal; - Optional labelRequest; + Optional labelRequest; } } diff --git a/hypertrace-graphql-entity-schema/src/test/java/org/hypertrace/graphql/entity/joiner/DefaultEntityJoinerBuilderTest.java b/hypertrace-graphql-entity-schema/src/test/java/org/hypertrace/graphql/entity/joiner/DefaultEntityJoinerBuilderTest.java index 4991ff71..cc92526e 100644 --- a/hypertrace-graphql-entity-schema/src/test/java/org/hypertrace/graphql/entity/joiner/DefaultEntityJoinerBuilderTest.java +++ b/hypertrace-graphql-entity-schema/src/test/java/org/hypertrace/graphql/entity/joiner/DefaultEntityJoinerBuilderTest.java @@ -3,9 +3,11 @@ import static java.util.Collections.emptyList; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -21,10 +23,12 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; 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.AttributeRequestBuilder; import org.hypertrace.core.graphql.common.request.FilterRequestBuilder; import org.hypertrace.core.graphql.common.request.ResultSetRequest; import org.hypertrace.core.graphql.common.request.ResultSetRequestBuilder; @@ -36,6 +40,8 @@ import org.hypertrace.core.graphql.utils.schema.SelectionQuery; import org.hypertrace.graphql.entity.dao.EntityDao; import org.hypertrace.graphql.entity.joiner.EntityJoiner.EntityIdGetter; +import org.hypertrace.graphql.entity.request.EntityLabelRequest; +import org.hypertrace.graphql.entity.request.EntityLabelRequestBuilder; import org.hypertrace.graphql.entity.request.EntityRequest; import org.hypertrace.graphql.entity.schema.EdgeResultSet; import org.hypertrace.graphql.entity.schema.Entity; @@ -60,10 +66,12 @@ class DefaultEntityJoinerBuilderTest { @Mock ArgumentDeserializer mockDeserializer; @Mock ResultSetRequestBuilder mockResultSetRequestBuilder; @Mock FilterRequestBuilder mockFilterRequestBuilder; + @Mock AttributeRequestBuilder attributeRequestBuilder; @Mock GraphQlRequestContext mockRequestContext; @Mock DataFetchingFieldSelectionSet mockSelectionSet; @Mock AttributeAssociation mockFilter; @Mock ResultSetRequest mockResultSetRequest; + @Mock EntityLabelRequestBuilder mockEntityLabelRequestBuilder; Scheduler testScheduler = Schedulers.single(); @@ -78,7 +86,9 @@ void beforeEach() { mockDeserializer, mockResultSetRequestBuilder, mockFilterRequestBuilder, - testScheduler); + attributeRequestBuilder, + testScheduler, + mockEntityLabelRequestBuilder); } @Test @@ -111,6 +121,9 @@ void fetchesEntitiesOfMultipleTypesInBatches() { mock(SelectedField.class)), "pathToEntity"); + when(this.mockEntityLabelRequestBuilder.buildLabelRequestIfPresentInAnyEntity( + eq(mockRequestContext), any(), any())) + .thenReturn(Single.just(Optional.empty())); mockRequestBuilding(); mockResult( Map.of( @@ -167,6 +180,73 @@ void doesNotFetchIfNoEntitiesRequested() { verify(this.mockEntityDao, never()).getEntities(any()); } + @Test + void fetchesEntitiesWithLabels() { + Entity firstTypeFirstIdEntity = new TestEntity(FIRST_ENTITY_TYPE, "first-id-1"); + Entity firstTypeSecondIdEntity = new TestEntity(FIRST_ENTITY_TYPE, "first-id-2"); + Entity secondTypeFirstIdEntity = new TestEntity(SECOND_ENTITY_TYPE, "second-id-1"); + Entity secondTypeSecondIdEntity = new TestEntity(SECOND_ENTITY_TYPE, "second-id-2"); + TestJoinSource firstJoinSource = new TestJoinSource("first-id-1", null); + TestJoinSource secondJoinSource = new TestJoinSource("first-id-2", null); + TestJoinSource thirdJoinSource = new TestJoinSource(null, "second-id-1"); + TestJoinSource fourthJoinSource = new TestJoinSource("first-id-1", "second-id-2"); + Table expected = + ImmutableTable.builder() + .put(firstJoinSource, FIRST_ENTITY_TYPE, firstTypeFirstIdEntity) + .put(secondJoinSource, FIRST_ENTITY_TYPE, firstTypeSecondIdEntity) + .put(thirdJoinSource, SECOND_ENTITY_TYPE, secondTypeFirstIdEntity) + .put(fourthJoinSource, FIRST_ENTITY_TYPE, firstTypeFirstIdEntity) + .put(fourthJoinSource, SECOND_ENTITY_TYPE, secondTypeSecondIdEntity) + .build(); + + List joinSources = + List.of(firstJoinSource, secondJoinSource, thirdJoinSource, fourthJoinSource); + + SelectedField mockSelectionField1 = mock(SelectedField.class); + SelectedField mockSelectionField2 = mock(SelectedField.class); + EntityLabelRequest mockEntityLabelRequest = mock(EntityLabelRequest.class); + when(mockEntityLabelRequestBuilder.buildLabelRequestIfPresentInAnyEntity( + eq(mockRequestContext), eq(FIRST_ENTITY_TYPE), eq(Set.of(mockSelectionField1)))) + .thenReturn(Single.just(Optional.of(mockEntityLabelRequest))); + when(mockEntityLabelRequestBuilder.buildLabelRequestIfPresentInAnyEntity( + eq(mockRequestContext), eq(SECOND_ENTITY_TYPE), eq(Set.of(mockSelectionField2)))) + .thenReturn(Single.just(Optional.empty())); + mockRequestedEntityFields( + Map.of( + FIRST_ENTITY_TYPE, mockSelectionField1, + SECOND_ENTITY_TYPE, mockSelectionField2), + "pathToEntity"); + + mockRequestBuilding(); + mockResult( + Map.of( + FIRST_ENTITY_TYPE, + List.of(firstTypeFirstIdEntity, firstTypeSecondIdEntity), + SECOND_ENTITY_TYPE, + List.of(secondTypeFirstIdEntity, secondTypeSecondIdEntity))); + + EntityJoiner joiner = + this.entityJoinerBuilder + .build(this.mockRequestContext, this.mockSelectionSet, List.of("pathToEntity")) + .blockingGet(); + + assertEquals( + expected, joiner.joinEntities(joinSources, new TestJoinSourceIdGetter()).blockingGet()); + + verify(mockEntityDao, times(1)) + .getEntities( + argThat( + request -> + request.entityType().equals(FIRST_ENTITY_TYPE) + && request.labelRequest().get().equals(mockEntityLabelRequest))); + verify(mockEntityDao, times(1)) + .getEntities( + argThat( + request -> + request.entityType().equals(SECOND_ENTITY_TYPE) + && request.labelRequest().isEmpty())); + } + private void mockRequestedEntityFields( Map selectedFieldsByEntityType, String location) { diff --git a/hypertrace-graphql-entity-schema/src/test/java/org/hypertrace/graphql/entity/request/DefaultEntityLabelRequestBuilderTest.java b/hypertrace-graphql-entity-schema/src/test/java/org/hypertrace/graphql/entity/request/DefaultEntityLabelRequestBuilderTest.java new file mode 100644 index 00000000..da20392b --- /dev/null +++ b/hypertrace-graphql-entity-schema/src/test/java/org/hypertrace/graphql/entity/request/DefaultEntityLabelRequestBuilderTest.java @@ -0,0 +1,89 @@ +package org.hypertrace.graphql.entity.request; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; + +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 org.hypertrace.core.graphql.common.request.AttributeRequest; +import org.hypertrace.core.graphql.common.request.AttributeRequestBuilder; +import org.hypertrace.core.graphql.context.GraphQlRequestContext; +import org.hypertrace.core.graphql.utils.schema.GraphQlSelectionFinder; +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) +class DefaultEntityLabelRequestBuilderTest { + + private static final String SCOPE = "scope"; + + @Mock GraphQlSelectionFinder mockSelectionFinder; + @Mock AttributeRequestBuilder attributeRequestBuilder; + @Mock GraphQlRequestContext mockRequestContext; + + EntityLabelRequestBuilder entityLabelRequestBuilder; + + @BeforeEach + void beforeEach() { + this.entityLabelRequestBuilder = + spy(new DefaultEntityLabelRequestBuilder(attributeRequestBuilder, mockSelectionFinder)); + } + + @Test + void test_buildLabelRequestIfPresentInResultSet_returnsEmpty() { + DataFetchingFieldSelectionSet mockSelectionSet = mock(DataFetchingFieldSelectionSet.class); + when(mockSelectionFinder.findSelections(eq(mockSelectionSet), any())) + .thenReturn(Stream.empty()); + Optional entityLabelRequestOptional = + this.entityLabelRequestBuilder + .buildLabelRequestIfPresentInResultSet(mockRequestContext, SCOPE, mockSelectionSet) + .blockingGet(); + assertTrue(entityLabelRequestOptional.isEmpty()); + } + + @Test + void test_buildLabelRequestIfPresentInResultSet_returnsLabelRequest() { + DataFetchingFieldSelectionSet mockSelectionSet = mock(DataFetchingFieldSelectionSet.class); + when(mockSelectionFinder.findSelections(eq(mockSelectionSet), any())) + .thenReturn(Stream.of(mock(SelectedField.class))); + AttributeRequest mockAttributeRequest = mock(AttributeRequest.class); + when(attributeRequestBuilder.buildForKey(eq(mockRequestContext), any(), eq("labels"))) + .thenReturn(Single.just(mockAttributeRequest)); + Optional entityLabelRequestOptional = + this.entityLabelRequestBuilder + .buildLabelRequestIfPresentInResultSet(mockRequestContext, SCOPE, mockSelectionSet) + .blockingGet(); + assertEquals( + mockAttributeRequest, entityLabelRequestOptional.get().labelIdArrayAttributeRequest()); + } + + @Test + void test_buildLabelRequestIfPresentInAnyEntity_returnsLabelRequest() { + SelectedField mockSelectionField = mock(SelectedField.class); + DataFetchingFieldSelectionSet mockSelectionSet = mock(DataFetchingFieldSelectionSet.class); + when(mockSelectionField.getSelectionSet()).thenReturn(mockSelectionSet); + when(mockSelectionFinder.findSelections(eq(mockSelectionSet), any())) + .thenReturn(Stream.of(mock(SelectedField.class))); + AttributeRequest mockAttributeRequest = mock(AttributeRequest.class); + when(attributeRequestBuilder.buildForKey(eq(mockRequestContext), any(), eq("labels"))) + .thenReturn(Single.just(mockAttributeRequest)); + Optional entityLabelRequestOptional = + this.entityLabelRequestBuilder + .buildLabelRequestIfPresentInAnyEntity( + mockRequestContext, SCOPE, List.of(mockSelectionField)) + .blockingGet(); + assertEquals( + mockAttributeRequest, entityLabelRequestOptional.get().labelIdArrayAttributeRequest()); + } +} diff --git a/hypertrace-graphql-labels-schema/src/main/java/org/hypertrace/graphql/label/joiner/DefaultLabelJoinerBuilder.java b/hypertrace-graphql-labels-schema/src/main/java/org/hypertrace/graphql/label/joiner/DefaultLabelJoinerBuilder.java index c3f7049d..9c8ba845 100644 --- a/hypertrace-graphql-labels-schema/src/main/java/org/hypertrace/graphql/label/joiner/DefaultLabelJoinerBuilder.java +++ b/hypertrace-graphql-labels-schema/src/main/java/org/hypertrace/graphql/label/joiner/DefaultLabelJoinerBuilder.java @@ -79,16 +79,17 @@ private Single> buildLabelResultSetMapEntry( private Single> filterLabels(List labelIds, Map labelMap) { return Single.just( labelIds.stream() - .map( - labelId -> { - Label label = labelMap.get(labelId); - if (label == null) { - log.warn("Label config doesn't exist for label id {}", labelId); - } - return label; - }) + .map(labelId -> getLabelFromMap(labelId, labelMap)) .filter(Objects::nonNull) .collect(Collectors.toUnmodifiableList())); } + + private Label getLabelFromMap(String labelId, Map labelMap) { + Label label = labelMap.get(labelId); + if (label == null) { + log.warn("Label config doesn't exist for label id {}", labelId); + } + return label; + } } }