feat: adding consumer support for new traces/spans api (#140)

* feat: adding consumer support for new traces/spans api

* updated the spans result set

* addressed the comments and updated gateway lib

* fixed owas issue, need to extend supression

Author: kotharironak

hypertrace-core-graphql-platform/build.gradle.kts

@@ -16,7 +16,7 @@ dependencies {
     api("org.hypertrace.core.grpcutils:grpc-context-utils:0.12.1")
     api("org.hypertrace.core.grpcutils:grpc-client-utils:0.12.1")
     api("org.hypertrace.core.grpcutils:grpc-client-rx-utils:0.12.1")
-    api("org.hypertrace.gateway.service:gateway-service-api:0.2.25")
+    api("org.hypertrace.gateway.service:gateway-service-api:0.3.0")
     api("org.hypertrace.core.attribute.service:caching-attribute-service-client:${attributeServiceVersion}")
     api("org.hypertrace.core.attribute.service:attribute-service-api:${attributeServiceVersion}")
 

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

@@ -34,7 +34,7 @@ class GatewayServiceSpanConverter {
   }
 
   public Single<SpanResultSet> convert(SpanRequest request, SpanLogEventsResponse response) {
-    int total = response.spansResponse().getTotal();
+    int total = response.spansResponse().hasTotal() ? response.spansResponse().getTotal() : 0;
 
     return Observable.fromIterable(response.spansResponse().getSpansList())
         .flatMapSingle(spanEvent -> this.convert(request, spanEvent, response.spanIdToLogEvents()))

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

@@ -56,6 +56,7 @@ Single<SpansRequest> buildRequest(SpanRequest gqlRequest) {
                         .spanEventsRequest()
                         .spaceId()
                         .orElse("")) // String proto default value
+                .setFetchTotal(gqlRequest.fetchTotal())
                 .build());
   }
 }

hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/request/DefaultSpanRequestBuilder.java

@@ -15,27 +15,40 @@
 import org.hypertrace.core.graphql.common.request.ResultSetRequest;
 import org.hypertrace.core.graphql.common.request.ResultSetRequestBuilder;
 import org.hypertrace.core.graphql.common.schema.attributes.arguments.AttributeExpression;
+import org.hypertrace.core.graphql.common.schema.results.ResultSet;
 import org.hypertrace.core.graphql.common.schema.results.arguments.order.OrderArgument;
 import org.hypertrace.core.graphql.context.GraphQlRequestContext;
+import org.hypertrace.core.graphql.utils.schema.GraphQlSelectionFinder;
+import org.hypertrace.core.graphql.utils.schema.SelectionQuery;
 
 class DefaultSpanRequestBuilder implements SpanRequestBuilder {
 
   private final ResultSetRequestBuilder resultSetRequestBuilder;
   private final LogEventAttributeRequestBuilder logEventAttributeRequestBuilder;
+  private final GraphQlSelectionFinder selectionFinder;
 
   @Inject
   public DefaultSpanRequestBuilder(
       ResultSetRequestBuilder resultSetRequestBuilder,
-      LogEventAttributeRequestBuilder logEventAttributeRequestBuilder) {
+      LogEventAttributeRequestBuilder logEventAttributeRequestBuilder,
+      GraphQlSelectionFinder selectionFinder) {
     this.resultSetRequestBuilder = resultSetRequestBuilder;
     this.logEventAttributeRequestBuilder = logEventAttributeRequestBuilder;
+    this.selectionFinder = selectionFinder;
   }
 
   @Override
   public Single<SpanRequest> build(
       GraphQlRequestContext context,
       Map<String, Object> arguments,
       DataFetchingFieldSelectionSet selectionSet) {
+    boolean fetchTotal =
+        this.selectionFinder
+                .findSelections(
+                    selectionSet, SelectionQuery.namedChild(ResultSet.RESULT_SET_TOTAL_NAME))
+                .count()
+            > 0;
+
     return zip(
         resultSetRequestBuilder.build(
             context,
@@ -45,7 +58,8 @@ public Single<SpanRequest> build(
             OrderArgument.class),
         logEventAttributeRequestBuilder.buildAttributeRequest(context, selectionSet),
         (resultSetRequest, logEventAttributeRequest) ->
-            new DefaultSpanRequest(context, resultSetRequest, logEventAttributeRequest));
+            new DefaultSpanRequest(
+                context, resultSetRequest, logEventAttributeRequest, fetchTotal));
   }
 
   @Override
@@ -54,12 +68,16 @@ public Single<SpanRequest> build(
       Map<String, Object> arguments,
       List<AttributeExpression> spanAttributeExpressions,
       List<AttributeExpression> logAttributeExpressions) {
+
     return zip(
         resultSetRequestBuilder.build(
             context, HypertraceCoreAttributeScopeString.SPAN, arguments, spanAttributeExpressions),
         logEventAttributeRequestBuilder.buildAttributeRequest(context, logAttributeExpressions),
         (resultSetRequest, logEventAttributeRequest) ->
-            new DefaultSpanRequest(context, resultSetRequest, logEventAttributeRequest));
+            // This build method is utilized in the exportSpans API, which does not accept the total
+            // parameter as an argument. Ref {@link ExportSpanResult}.
+            // So, we explicitly set fetchTotal to false.
+            new DefaultSpanRequest(context, resultSetRequest, logEventAttributeRequest, false));
   }
 
   @Value
@@ -68,5 +86,6 @@ private static class DefaultSpanRequest implements SpanRequest {
     GraphQlRequestContext context;
     ResultSetRequest<OrderArgument> spanEventsRequest;
     Collection<AttributeRequest> logEventAttributes;
+    boolean fetchTotal;
   }
 }

hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/request/SpanRequest.java

@@ -10,4 +10,6 @@ public interface SpanRequest extends ContextualRequest {
   ResultSetRequest<OrderArgument> spanEventsRequest();
 
   Collection<AttributeRequest> logEventAttributes();
+
+  boolean fetchTotal();
 }

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

@@ -72,6 +72,7 @@ static class DefaultSpanRequest implements SpanRequest {
     GraphQlRequestContext context;
     ResultSetRequest<OrderArgument> spanEventsRequest;
     Collection<AttributeRequest> logEventAttributes;
+    boolean fetchTotal;
   }
 
   @Value

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

@@ -152,7 +152,8 @@ void testBuildRequest() {
             List.of(),
             Collections.emptyList(),
             Optional.empty());
-    SpanRequest spanRequest = new DefaultSpanRequest(null, resultSetRequest, logAttributeRequests);
+    SpanRequest spanRequest =
+        new DefaultSpanRequest(null, resultSetRequest, logAttributeRequests, true);
 
     LogEventsRequest expectedLogEventsRequest =
         LogEventsRequest.newBuilder()
@@ -193,7 +194,8 @@ void testBuildRequest_addSpanId() {
             List.of(),
             Collections.emptyList(),
             Optional.empty());
-    SpanRequest spanRequest = new DefaultSpanRequest(null, resultSetRequest, logAttributeRequests);
+    SpanRequest spanRequest =
+        new DefaultSpanRequest(null, resultSetRequest, logAttributeRequests, true);
 
     LogEventsRequest expectedLogEventsRequest =
         LogEventsRequest.newBuilder()

hypertrace-core-graphql-trace-schema/build.gradle.kts

@@ -24,4 +24,5 @@ dependencies {
   implementation(project(":hypertrace-core-graphql-attribute-store"))
   implementation(project(":hypertrace-core-graphql-deserialization"))
   implementation(project(":hypertrace-core-graphql-request-transformation"))
+  implementation(project(":hypertrace-core-graphql-schema-utils"))
 }

hypertrace-core-graphql-trace-schema/src/main/java/org/hypertrace/core/graphql/trace/dao/GatewayServiceTraceConverter.java

@@ -30,8 +30,7 @@ public class GatewayServiceTraceConverter {
   }
 
   public Single<TraceResultSet> convert(ResultSetRequest<?> request, TracesResponse response) {
-    int total = response.getTotal();
-
+    int total = response.hasTotal() ? response.getTotal() : 0;
     return Observable.fromIterable(response.getTracesList())
         .flatMapSingle(trace -> this.convert(request, trace))
         .toList()

hypertrace-core-graphql-trace-schema/src/main/java/org/hypertrace/core/graphql/trace/dao/GatewayServiceTraceRequestBuilder.java

@@ -54,6 +54,7 @@ Single<TracesRequest> buildRequest(TraceRequest request) {
                 .setFilter(filters)
                 .setSpaceId(
                     request.resultSetRequest().spaceId().orElse("")) // String proto default value
+                .setFetchTotal(request.fetchTotal())
                 .build());
   }
 }