feat: add support for request transformers (#72)

Author: aaron-steinfeld

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

@@ -22,6 +22,7 @@ dependencies {
   implementation(project(":hypertrace-core-graphql-attribute-scope"))
   implementation(project(":hypertrace-core-graphql-rx-utils"))
   implementation(project(":hypertrace-core-graphql-log-event-schema"))
+  implementation(project(":hypertrace-core-graphql-request-transformation"))
 
   implementation("org.slf4j:slf4j-api")
   implementation("com.google.inject:guice")

hypertrace-core-graphql-impl/src/main/java/org/hypertrace/core/graphql/impl/GraphQlModule.java

@@ -9,6 +9,7 @@
 import org.hypertrace.core.graphql.deserialization.GraphQlDeserializationRegistryModule;
 import org.hypertrace.core.graphql.log.event.LogEventSchemaModule;
 import org.hypertrace.core.graphql.metadata.MetadataSchemaModule;
+import org.hypertrace.core.graphql.request.transformation.RequestTransformationModule;
 import org.hypertrace.core.graphql.rx.RxUtilModule;
 import org.hypertrace.core.graphql.schema.registry.GraphQlSchemaRegistryModule;
 import org.hypertrace.core.graphql.span.SpanSchemaModule;
@@ -48,5 +49,6 @@ protected void configure() {
     install(new TraceSchemaModule());
     install(new RxUtilModule());
     install(new LogEventSchemaModule());
+    install(new RequestTransformationModule());
   }
 }

hypertrace-core-graphql-request-transformation/build.gradle.kts

@@ -0,0 +1,19 @@
+plugins {
+  `java-library`
+  jacoco
+  id("org.hypertrace.jacoco-report-plugin")
+}
+
+dependencies {
+  api("com.google.inject:guice")
+  api(project(":hypertrace-core-graphql-common-schema"))
+
+  testAnnotationProcessor("org.projectlombok:lombok")
+  testCompileOnly("org.projectlombok:lombok")
+  testImplementation("org.junit.jupiter:junit-jupiter")
+  testImplementation("org.mockito:mockito-core")
+}
+
+tasks.test {
+  useJUnitPlatform()
+}

hypertrace-core-graphql-request-transformation/src/main/java/org/hypertrace/core/graphql/request/transformation/RequestTransformation.java

@@ -0,0 +1,20 @@
+package org.hypertrace.core.graphql.request.transformation;
+
+import io.reactivex.rxjava3.core.Single;
+import org.hypertrace.core.graphql.common.request.ContextualRequest;
+
+public interface RequestTransformation {
+
+  /**
+   * Any request may be provided to this method, which should return true if this transformation
+   * supports the request and false otherwise.
+   */
+  boolean supportsRequest(ContextualRequest request);
+
+  /**
+   * Applies a transformation to the request. This will only receive requests that have previously
+   * returned true from {@link #supportsRequest(ContextualRequest)}, and must always return a
+   * request, potentially the same request passed in if no transformation is necessary.
+   */
+  <T extends ContextualRequest> Single<T> transform(T request);
+}

hypertrace-core-graphql-request-transformation/src/main/java/org/hypertrace/core/graphql/request/transformation/RequestTransformationModule.java

@@ -0,0 +1,13 @@
+package org.hypertrace.core.graphql.request.transformation;
+
+import com.google.inject.AbstractModule;
+import com.google.inject.multibindings.Multibinder;
+
+public class RequestTransformationModule extends AbstractModule {
+
+  @Override
+  protected void configure() {
+    bind(RequestTransformer.class).to(RequestTransformerImpl.class);
+    Multibinder.newSetBinder(binder(), RequestTransformation.class);
+  }
+}

hypertrace-core-graphql-request-transformation/src/main/java/org/hypertrace/core/graphql/request/transformation/RequestTransformer.java

@@ -0,0 +1,12 @@
+package org.hypertrace.core.graphql.request.transformation;
+
+import io.reactivex.rxjava3.core.Single;
+import org.hypertrace.core.graphql.common.request.ContextualRequest;
+
+/**
+ * A request transformer that can receive any request and will apply all applicable transformations
+ * that have been registered (currently, the order is undefined), returning the result.
+ */
+public interface RequestTransformer {
+  <T extends ContextualRequest> Single<T> transform(T request);
+}

hypertrace-core-graphql-request-transformation/src/main/java/org/hypertrace/core/graphql/request/transformation/RequestTransformerImpl.java

@@ -0,0 +1,27 @@
+package org.hypertrace.core.graphql.request.transformation;
+
+import io.reactivex.rxjava3.core.Observable;
+import io.reactivex.rxjava3.core.Single;
+import java.util.Set;
+import javax.inject.Inject;
+import org.hypertrace.core.graphql.common.request.ContextualRequest;
+
+class RequestTransformerImpl implements RequestTransformer {
+  private final Set<RequestTransformation> requestTransformations;
+
+  @Inject
+  RequestTransformerImpl(Set<RequestTransformation> requestTransformations) {
+    this.requestTransformations = requestTransformations;
+  }
+
+  @Override
+  public <T extends ContextualRequest> Single<T> transform(T request) {
+    return Observable.fromIterable(requestTransformations)
+        .filter(requestTransformation -> requestTransformation.supportsRequest(request))
+        .reduce(
+            Single.just(request), // Work in singles since there's no flat reduce, then unwrap later
+            (currentRequestSingle, requestTransformation) ->
+                currentRequestSingle.flatMap(requestTransformation::transform))
+        .flatMap(wrapped -> wrapped);
+  }
+}

hypertrace-core-graphql-request-transformation/src/test/java/org/hypertrace/core/graphql/request/transformation/RequestTransformerImplTest.java

@@ -0,0 +1,85 @@
+package org.hypertrace.core.graphql.request.transformation;
+
+import static org.mockito.Mockito.mock;
+
+import io.reactivex.rxjava3.core.Single;
+import io.reactivex.rxjava3.observers.TestObserver;
+import java.util.Set;
+import lombok.Value;
+import lombok.experimental.Accessors;
+import org.hypertrace.core.graphql.common.request.ContextualRequest;
+import org.hypertrace.core.graphql.context.GraphQlRequestContext;
+import org.junit.jupiter.api.Test;
+
+class RequestTransformerImplTest {
+
+  @Test
+  void noApplicableTransformerReturnsInput() {
+    RequestTransformer transformer = new RequestTransformerImpl(Set.of());
+
+    ContextualRequest request = mock(ContextualRequest.class);
+
+    TestObserver<ContextualRequest> testObserver = new TestObserver<>();
+    transformer.transform(request).subscribe(testObserver);
+
+    testObserver.assertResult(request);
+  }
+
+  @Test
+  void skipsUnsupportedTransformers() {
+    GraphQlRequestContext context = mock(GraphQlRequestContext.class);
+    RequestTransformer transformer =
+        new RequestTransformerImpl(
+            Set.of(
+                new TestTransformation(true, "prefix-", ""),
+                new TestTransformation(false, "", "-suffix")));
+
+    TestObserver<ContextualRequest> testObserver = new TestObserver<>();
+    transformer.transform(new TestContextualRequest(context, "original")).subscribe(testObserver);
+
+    testObserver.assertResult(new TestContextualRequest(context, "prefix-original"));
+  }
+
+  @Test
+  void appliesAllTransformers() {
+    GraphQlRequestContext context = mock(GraphQlRequestContext.class);
+    RequestTransformer transformer =
+        new RequestTransformerImpl(
+            Set.of(
+                new TestTransformation(true, "prefix-", ""),
+                new TestTransformation(true, "", "-suffix")));
+
+    TestObserver<ContextualRequest> testObserver = new TestObserver<>();
+    transformer.transform(new TestContextualRequest(context, "original")).subscribe(testObserver);
+
+    testObserver.assertResult(new TestContextualRequest(context, "prefix-original-suffix"));
+  }
+
+  @Value
+  private static class TestTransformation implements RequestTransformation {
+    boolean supported;
+    String prefixToAdd;
+    String suffixToAdd;
+
+    @Override
+    public boolean supportsRequest(ContextualRequest request) {
+      return supported;
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public <T extends ContextualRequest> Single<T> transform(T request) {
+      String previousValue = ((TestContextualRequest) request).value();
+      TestContextualRequest newRequest =
+          new TestContextualRequest(request.context(), prefixToAdd + previousValue + suffixToAdd);
+      return Single.just((T) newRequest);
+    }
+  }
+
+  @Value
+  @Accessors(fluent = true)
+  private static class TestContextualRequest implements ContextualRequest {
+    GraphQlRequestContext context;
+    String value;
+  }
+}

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

@@ -28,6 +28,7 @@ dependencies {
   implementation(project(":hypertrace-core-graphql-deserialization"))
   implementation(project(":hypertrace-core-graphql-schema-utils"))
   implementation(project(":hypertrace-core-graphql-attribute-scope-constants"))
+  implementation(project(":hypertrace-core-graphql-request-transformation"))
 
   testImplementation("org.junit.jupiter:junit-jupiter")
   testImplementation("com.fasterxml.jackson.core:jackson-databind")

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

@@ -6,6 +6,7 @@
 import javax.inject.Inject;
 import javax.inject.Singleton;
 import org.hypertrace.core.graphql.context.GraphQlRequestContext;
+import org.hypertrace.core.graphql.request.transformation.RequestTransformer;
 import org.hypertrace.core.graphql.span.request.SpanRequest;
 import org.hypertrace.core.graphql.span.schema.SpanResultSet;
 import org.hypertrace.core.graphql.spi.config.GraphQlServiceConfig;
@@ -23,6 +24,7 @@ class GatewayServiceSpanDao implements SpanDao {
   private final GatewayServiceSpanConverter spanConverter;
   private final SpanLogEventDao spanLogEventDao;
   private final GraphQlServiceConfig serviceConfig;
+  private final RequestTransformer requestTransformer;
 
   @Inject
   GatewayServiceSpanDao(
@@ -31,19 +33,22 @@ class GatewayServiceSpanDao implements SpanDao {
       GrpcContextBuilder grpcContextBuilder,
       GatewayServiceSpanRequestBuilder requestBuilder,
       GatewayServiceSpanConverter spanConverter,
-      SpanLogEventDao spanLogEventDao) {
+      SpanLogEventDao spanLogEventDao,
+      RequestTransformer requestTransformer) {
     this.grpcContextBuilder = grpcContextBuilder;
     this.requestBuilder = requestBuilder;
     this.spanConverter = spanConverter;
     this.spanLogEventDao = spanLogEventDao;
     this.gatewayServiceStub = gatewayServiceFutureStub;
     this.serviceConfig = serviceConfig;
+    this.requestTransformer = requestTransformer;
   }
 
   @Override
   public Single<SpanResultSet> getSpans(SpanRequest request) {
-    return this.requestBuilder
-        .buildRequest(request)
+    return this.requestTransformer
+        .transform(request)
+        .flatMap(this.requestBuilder::buildRequest)
         .flatMap(
             serverRequest -> this.makeRequest(request.spanEventsRequest().context(), serverRequest))
         .flatMap(serverResponse -> spanLogEventDao.fetchLogEvents(request, serverResponse))