chore | Add support for providing filters in spanJoiner (#181)

* Add support for providing filters in spanJoiner * Resolve PR reviews
hypertrace · Jun 7, 2024 · 36bd206 · 36bd206
1 parent 4202d01
commit 36bd206
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 21 deletions.
diff --git a/...chema/src/main/java/org/hypertrace/core/graphql/span/joiner/DefaultSpanJoinerBuilder.java b/...chema/src/main/java/org/hypertrace/core/graphql/span/joiner/DefaultSpanJoinerBuilder.java
@@ -18,9 +18,9 @@
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Optional;
-import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import javax.inject.Inject;
 import lombok.AllArgsConstructor;
 import lombok.Value;
@@ -86,26 +86,32 @@ private class DefaultSpanJoiner implements SpanJoiner {
 
     @Override
     public <T> Single<Map<T, Span>> joinSpan(
-        Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
+        Collection<T> joinSources,
+        SpanIdGetter<T> spanIdGetter,
+        Collection<FilterArgument> filterArguments) {
       Function<T, Single<List<String>>> idsGetter =
           source -> spanIdGetter.getSpanId(source).map(List::of);
-      return this.joinSpans(joinSources, idsGetter, SPAN_KEY).map(this::reduceMap);
+      return this.joinSpans(joinSources, idsGetter, filterArguments, SPAN_KEY).map(this::reduceMap);
     }
 
     @Override
     public <T> Single<ListMultimap<T, Span>> joinSpans(
-        Collection<T> joinSources, MultipleSpanIdGetter<T> multipleSpanIdGetter) {
-      return this.joinSpans(joinSources, multipleSpanIdGetter::getSpanIds, SPANS_KEY);
+        Collection<T> joinSources,
+        MultipleSpanIdGetter<T> multipleSpanIdGetter,
+        Collection<FilterArgument> filterArguments) {
+      return this.joinSpans(
+          joinSources, multipleSpanIdGetter::getSpanIds, filterArguments, SPANS_KEY);
     }
 
     private <T> Single<ListMultimap<T, Span>> joinSpans(
         Collection<T> joinSources,
         Function<T, Single<List<String>>> idsGetter,
+        Collection<FilterArgument> filterArguments,
         String joinSpanKey) {
       return this.buildSourceToIdsMap(joinSources, idsGetter)
           .flatMap(
               sourceToSpanIdsMap ->
-                  this.buildSpanRequest(sourceToSpanIdsMap, joinSpanKey)
+                  this.buildSpanRequest(sourceToSpanIdsMap, joinSpanKey, filterArguments)
                       .flatMap(spanDao::getSpans)
                       .map(this::buildSpanIdToSpanMap)
                       .map(
@@ -152,15 +158,16 @@ private Map<String, Span> buildSpanIdToSpanMap(SpanResultSet resultSet) {
     }
 
     private <T> Single<SpanRequest> buildSpanRequest(
-        ListMultimap<T, String> sourceToSpanIdsMultimap, String joinSpanKey) {
+        ListMultimap<T, String> sourceToSpanIdsMultimap,
+        String joinSpanKey,
+        Collection<FilterArgument> filterArguments) {
       Collection<String> spanIds =
           sourceToSpanIdsMultimap.values().stream()
               .distinct()
               .collect(Collectors.toUnmodifiableList());
       List<SelectedField> selectedFields = getSelections(joinSpanKey);
-      return buildSpanIdsFilter(context, spanIds)
-          .flatMap(
-              filterArguments -> buildSpanRequest(spanIds.size(), filterArguments, selectedFields));
+      return buildSpanIdsFilter(context, spanIds, filterArguments)
+          .flatMap(filters -> buildSpanRequest(spanIds.size(), filters, selectedFields));
     }
 
     private Single<SpanRequest> buildSpanRequest(
@@ -182,8 +189,14 @@ private Single<SpanRequest> buildSpanRequest(
     }
 
     private Single<List<AttributeAssociation<FilterArgument>>> buildSpanIdsFilter(
-        GraphQlRequestContext context, Collection<String> spanIds) {
-      return filterRequestBuilder.build(context, SPAN, Set.of(new SpanIdFilter(spanIds)));
+        GraphQlRequestContext context,
+        Collection<String> spanIds,
+        Collection<FilterArgument> filterArguments) {
+      return filterRequestBuilder.build(
+          context,
+          SPAN,
+          Stream.concat(filterArguments.stream(), Stream.of(new SpanIdFilter(spanIds)))
+              .collect(Collectors.toUnmodifiableList()));
     }
   }
 

diff --git a/...graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoiner.java b/...graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/joiner/SpanJoiner.java
@@ -7,6 +7,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterArgument;
 import org.hypertrace.core.graphql.span.schema.Span;
 
 public interface SpanJoiner {
@@ -16,21 +17,40 @@ public interface SpanJoiner {
       new SpanJoiner() {
         @Override
         public <T> Single<Map<T, Span>> joinSpan(
-            Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
+            Collection<T> joinSources,
+            SpanIdGetter<T> spanIdGetter,
+            Collection<FilterArgument> filterArguments) {
           return Single.just(Collections.emptyMap());
         }
 
         @Override
         public <T> Single<ListMultimap<T, Span>> joinSpans(
-            Collection<T> joinSources, MultipleSpanIdGetter<T> multipleSpanIdGetter) {
+            Collection<T> joinSources,
+            MultipleSpanIdGetter<T> multipleSpanIdGetter,
+            Collection<FilterArgument> filterArguments) {
           return Single.just(ArrayListMultimap.create());
         }
       };
 
-  <T> Single<Map<T, Span>> joinSpan(Collection<T> joinSources, SpanIdGetter<T> spanIdGetter);
+  default <T> Single<Map<T, Span>> joinSpan(
+      Collection<T> joinSources, SpanIdGetter<T> spanIdGetter) {
+    return joinSpan(joinSources, spanIdGetter, Collections.emptyList());
+  }
+
+  default <T> Single<ListMultimap<T, Span>> joinSpans(
+      Collection<T> joinSources, MultipleSpanIdGetter<T> multipleSpanIdGetter) {
+    return joinSpans(joinSources, multipleSpanIdGetter, Collections.emptyList());
+  }
+
+  <T> Single<Map<T, Span>> joinSpan(
+      Collection<T> joinSources,
+      SpanIdGetter<T> spanIdGetter,
+      Collection<FilterArgument> filterArguments);
 
   <T> Single<ListMultimap<T, Span>> joinSpans(
-      Collection<T> joinSources, MultipleSpanIdGetter<T> multipleSpanIdGetter);
+      Collection<T> joinSources,
+      MultipleSpanIdGetter<T> multipleSpanIdGetter,
+      Collection<FilterArgument> filterArguments);
 
   @FunctionalInterface
   interface SpanIdGetter<T> {

diff --git a/...n-schema/src/test/java/org/hypertrace/core/graphql/span/joiner/SpanJoinerBuilderTest.java b/...n-schema/src/test/java/org/hypertrace/core/graphql/span/joiner/SpanJoinerBuilderTest.java
@@ -5,7 +5,7 @@
 import static org.hypertrace.core.graphql.atttributes.scopes.HypertraceCoreAttributeScopeString.SPAN;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anySet;
+import static org.mockito.ArgumentMatchers.anyList;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -15,6 +15,7 @@
 import graphql.schema.DataFetchingFieldSelectionSet;
 import graphql.schema.SelectedField;
 import io.reactivex.rxjava3.core.Single;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -88,7 +89,7 @@ void fetchSpans() {
             mockSelectionSet,
             SelectionQuery.builder().selectionPath(List.of("pathToSpan", "span")).build()))
         .thenReturn(Stream.of(mock(SelectedField.class), mock(SelectedField.class)));
-    when(mockFilterRequestBuilder.build(eq(mockRequestContext), eq(SPAN), anySet()))
+    when(mockFilterRequestBuilder.build(eq(mockRequestContext), eq(SPAN), anyList()))
         .thenReturn(Single.just(List.of(mockFilter)));
 
     when(mockResultSetRequestBuilder.build(
@@ -112,7 +113,10 @@ void fetchSpans() {
                 List.of("pathToSpan"))
             .blockingGet();
     assertEquals(
-        expected, joiner.joinSpan(joinSources, new TestJoinSourceIdGetter()).blockingGet());
+        expected,
+        joiner
+            .joinSpan(joinSources, new TestJoinSourceIdGetter(), Collections.emptyList())
+            .blockingGet());
   }
 
   @Test
@@ -129,7 +133,7 @@ void fetchMultipleSpans() {
             mockSelectionSet,
             SelectionQuery.builder().selectionPath(List.of("pathToSpans", "spans")).build()))
         .thenReturn(Stream.of(mock(SelectedField.class), mock(SelectedField.class)));
-    when(mockFilterRequestBuilder.build(eq(mockRequestContext), eq(SPAN), anySet()))
+    when(mockFilterRequestBuilder.build(eq(mockRequestContext), eq(SPAN), anyList()))
         .thenReturn(Single.just(List.of(mockFilter)));
 
     when(mockResultSetRequestBuilder.build(
@@ -155,7 +159,9 @@ void fetchMultipleSpans() {
             .blockingGet();
     assertEquals(
         expected,
-        joiner.joinSpans(joinSources, new TestMultipleJoinSourceIdGetter()).blockingGet());
+        joiner
+            .joinSpans(joinSources, new TestMultipleJoinSourceIdGetter(), Collections.emptyList())
+            .blockingGet());
   }
 
   private void mockResult(List<Span> spans) {