refactor: deprecate graphql request context (#68)

* refactor: deprecate graphql request context

* refactor: migrated all usages in core

Author: aaron-steinfeld

hypertrace-core-graphql-attribute-store/build.gradle.kts

@@ -14,6 +14,7 @@ dependencies {
   implementation("com.google.guava:guava")
 
   implementation("org.hypertrace.core.attribute.service:attribute-service-api")
+  implementation("org.hypertrace.core.grpcutils:grpc-client-rx-utils")
   implementation(project(":hypertrace-core-graphql-grpc-utils"))
   implementation(project(":hypertrace-core-graphql-rx-utils"))
 

hypertrace-core-graphql-attribute-store/src/main/java/org/hypertrace/core/graphql/attributes/AttributeClient.java

@@ -12,20 +12,21 @@
 import org.hypertrace.core.attribute.service.v1.Empty;
 import org.hypertrace.core.graphql.context.GraphQlRequestContext;
 import org.hypertrace.core.graphql.spi.config.GraphQlServiceConfig;
-import org.hypertrace.core.graphql.utils.grpc.GraphQlGrpcContextBuilder;
 import org.hypertrace.core.graphql.utils.grpc.GrpcChannelRegistry;
+import org.hypertrace.core.graphql.utils.grpc.GrpcContextBuilder;
+import org.hypertrace.core.grpcutils.client.rx.GrpcRxExecutionContext;
 
 @Singleton
 class AttributeClient {
   private static final int DEFAULT_DEADLINE_SEC = 10;
   private final AttributeServiceStub attributeServiceClient;
-  private final GraphQlGrpcContextBuilder grpcContextBuilder;
+  private final GrpcContextBuilder grpcContextBuilder;
   private final AttributeModelTranslator translator;
 
   @Inject
   AttributeClient(
       GraphQlServiceConfig serviceConfig,
-      GraphQlGrpcContextBuilder grpcContextBuilder,
+      GrpcContextBuilder grpcContextBuilder,
       GrpcChannelRegistry grpcChannelRegistry,
       CallCredentials credentials,
       AttributeModelTranslator translator) {
@@ -41,9 +42,8 @@ class AttributeClient {
   }
 
   public Observable<AttributeModel> queryAll(GraphQlRequestContext requestContext) {
-    return this.grpcContextBuilder
-        .build(requestContext)
-        .<AttributeMetadata>streamInContext(
+    return GrpcRxExecutionContext.forContext(this.grpcContextBuilder.build(requestContext))
+        .<AttributeMetadata>stream(
             streamObserver ->
                 this.attributeServiceClient
                     .withDeadlineAfter(DEFAULT_DEADLINE_SEC, TimeUnit.SECONDS)

hypertrace-core-graphql-attribute-store/src/main/java/org/hypertrace/core/graphql/attributes/AttributeStoreModule.java

@@ -7,8 +7,8 @@
 import io.reactivex.rxjava3.core.Scheduler;
 import org.hypertrace.core.graphql.rx.BoundedIoScheduler;
 import org.hypertrace.core.graphql.spi.config.GraphQlServiceConfig;
-import org.hypertrace.core.graphql.utils.grpc.GraphQlGrpcContextBuilder;
 import org.hypertrace.core.graphql.utils.grpc.GrpcChannelRegistry;
+import org.hypertrace.core.graphql.utils.grpc.GrpcContextBuilder;
 
 public class AttributeStoreModule extends AbstractModule {
 
@@ -18,7 +18,7 @@ protected void configure() {
     Multibinder.newSetBinder(binder(), IdMapping.class);
     Multibinder.newSetBinder(binder(), IdMappingLoader.class);
     requireBinding(GraphQlServiceConfig.class);
-    requireBinding(GraphQlGrpcContextBuilder.class);
+    requireBinding(GrpcContextBuilder.class);
     requireBinding(CallCredentials.class);
     requireBinding(GrpcChannelRegistry.class);
     requireBinding(Key.get(Scheduler.class, BoundedIoScheduler.class));

hypertrace-core-graphql-grpc-utils/build.gradle.kts

@@ -11,8 +11,8 @@ dependencies {
   api("io.grpc:grpc-core")
   api("io.grpc:grpc-stub")
   api(project(":hypertrace-core-graphql-context"))
+  api("org.hypertrace.core.grpcutils:grpc-context-utils")
 
-  implementation("org.hypertrace.core.grpcutils:grpc-context-utils")
   implementation("org.hypertrace.core.grpcutils:grpc-client-utils")
   implementation("io.grpc:grpc-context")
   implementation("io.reactivex.rxjava3:rxjava")

hypertrace-core-graphql-grpc-utils/src/main/java/org/hypertrace/core/graphql/utils/grpc/DefaultGraphQlGrpcContextBuilder.java

@@ -1,70 +1,40 @@
 package org.hypertrace.core.graphql.utils.grpc;
 
-import static org.hypertrace.core.grpcutils.context.RequestContextConstants.AUTHORIZATION_HEADER;
-import static org.hypertrace.core.grpcutils.context.RequestContextConstants.TENANT_ID_HEADER_KEY;
-
-import io.grpc.Context;
 import io.grpc.stub.StreamObserver;
 import io.reactivex.rxjava3.core.Observable;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.Map;
-import java.util.Map.Entry;
-import java.util.Optional;
 import java.util.concurrent.Callable;
 import java.util.function.Consumer;
-import java.util.stream.Collectors;
 import javax.inject.Inject;
 import org.hypertrace.core.graphql.context.GraphQlRequestContext;
 import org.hypertrace.core.grpcutils.context.RequestContext;
 
 class DefaultGraphQlGrpcContextBuilder implements GraphQlGrpcContextBuilder {
 
-  private final PlatformRequestContextBuilder platformRequestContextBuilder;
+  private final PlatformGrpcContextBuilder platformGrpcContextBuilder;
 
   @Inject
-  DefaultGraphQlGrpcContextBuilder(PlatformRequestContextBuilder platformRequestContextBuilder) {
-    this.platformRequestContextBuilder = platformRequestContextBuilder;
+  DefaultGraphQlGrpcContextBuilder(PlatformGrpcContextBuilder platformGrpcContextBuilder) {
+    this.platformGrpcContextBuilder = platformGrpcContextBuilder;
   }
 
   @Override
   public GraphQlGrpcContext build(GraphQlRequestContext requestContext) {
-    return new DefaultGraphQlGrpcContext(requestContext, this.platformRequestContextBuilder);
+    return new DefaultGraphQlGrpcContext(requestContext, this.platformGrpcContextBuilder);
   }
 
   private static final class DefaultGraphQlGrpcContext implements GraphQlGrpcContext {
-    private final Context grpcContext;
+
+    private final RequestContext grpcContext;
 
     private DefaultGraphQlGrpcContext(
         GraphQlRequestContext requestContext,
-        PlatformRequestContextBuilder platformRequestContextBuilder) {
-      Map<String, String> grpcHeaders =
-          this.mergeMaps(
-              requestContext.getTracingContextHeaders(),
-              this.flattenOptionalMap(
-                  Map.of(
-                      AUTHORIZATION_HEADER, this.extractAuthorizationHeader(requestContext),
-                      TENANT_ID_HEADER_KEY, this.extractTenantId(requestContext))));
-
-      RequestContext platformContext = platformRequestContextBuilder.build(grpcHeaders);
-      this.grpcContext = this.buildGrpcContext(platformContext);
-    }
-
-    private Optional<String> extractAuthorizationHeader(GraphQlRequestContext requestContext) {
-      return requestContext.getAuthorizationHeader().filter(header -> header.length() > 0);
-    }
-
-    private Optional<String> extractTenantId(GraphQlRequestContext requestContext) {
-      return requestContext.getTenantId();
+        PlatformGrpcContextBuilder platformGrpcContextBuilder) {
+      this.grpcContext = platformGrpcContextBuilder.build(requestContext);
     }
 
     @Override
     public <TResp> TResp callInContext(Callable<TResp> callable) {
-      try {
-        return this.grpcContext.call(callable);
-      } catch (Exception e) {
-        throw new RuntimeException(e);
-      }
+      return this.grpcContext.call(callable);
     }
 
     @Override
@@ -80,24 +50,5 @@ public <TResponse> Observable<TResponse> streamInContext(
               this.runInContext(
                   () -> requestExecutor.accept(new StreamingClientResponseObserver<>(emitter))));
     }
-
-    private Context buildGrpcContext(RequestContext platformContext) {
-      return Context.current().withValue(RequestContext.CURRENT, platformContext);
-    }
-
-    private Map<String, String> flattenOptionalMap(Map<String, Optional<String>> optionalMap) {
-      return optionalMap.entrySet().stream()
-          .map(entry -> entry.getValue().map(value -> Map.entry(entry.getKey(), value)))
-          .flatMap(Optional::stream)
-          .collect(Collectors.toUnmodifiableMap(Entry::getKey, Entry::getValue));
-    }
-
-    @SafeVarargs
-    private Map<String, String> mergeMaps(Map<String, String>... maps) {
-      return Arrays.stream(maps)
-          .map(Map::entrySet)
-          .flatMap(Collection::stream)
-          .collect(Collectors.toUnmodifiableMap(Entry::getKey, Entry::getValue));
-    }
   }
 }

hypertrace-core-graphql-grpc-utils/src/main/java/org/hypertrace/core/graphql/utils/grpc/GraphQlGrpcContext.java

@@ -5,6 +5,7 @@
 import java.util.concurrent.Callable;
 import java.util.function.Consumer;
 
+@Deprecated
 public interface GraphQlGrpcContext {
 
   <TResp> TResp callInContext(Callable<TResp> callable);

hypertrace-core-graphql-grpc-utils/src/main/java/org/hypertrace/core/graphql/utils/grpc/GraphQlGrpcContextBuilder.java

@@ -2,6 +2,8 @@
 
 import org.hypertrace.core.graphql.context.GraphQlRequestContext;
 
+/** @deprecated Use {@link GrpcContextBuilder#build(GraphQlRequestContext)} instead */
+@Deprecated
 public interface GraphQlGrpcContextBuilder {
   GraphQlGrpcContext build(GraphQlRequestContext requestContext);
 }

hypertrace-core-graphql-grpc-utils/src/main/java/org/hypertrace/core/graphql/utils/grpc/GraphQlGrpcModule.java

@@ -12,5 +12,6 @@ protected void configure() {
     bind(CallCredentials.class).toInstance(getClientCallCredsProvider().get());
     bind(GraphQlGrpcContextBuilder.class).to(DefaultGraphQlGrpcContextBuilder.class);
     bind(GrpcChannelRegistry.class).to(DefaultGrpcChannelRegistry.class);
+    bind(GrpcContextBuilder.class).to(PlatformGrpcContextBuilder.class);
   }
 }

hypertrace-core-graphql-grpc-utils/src/main/java/org/hypertrace/core/graphql/utils/grpc/GrpcContextBuilder.java

@@ -0,0 +1,8 @@
+package org.hypertrace.core.graphql.utils.grpc;
+
+import org.hypertrace.core.graphql.context.GraphQlRequestContext;
+import org.hypertrace.core.grpcutils.context.RequestContext;
+
+public interface GrpcContextBuilder {
+  RequestContext build(GraphQlRequestContext requestContext);
+}

hypertrace-core-graphql-grpc-utils/src/main/java/org/hypertrace/core/graphql/utils/grpc/PlatformGrpcContextBuilder.java

@@ -0,0 +1,59 @@
+package org.hypertrace.core.graphql.utils.grpc;
+
+import static org.hypertrace.core.grpcutils.context.RequestContextConstants.AUTHORIZATION_HEADER;
+import static org.hypertrace.core.grpcutils.context.RequestContextConstants.TENANT_ID_HEADER_KEY;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Optional;
+import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
+import org.hypertrace.core.graphql.context.GraphQlRequestContext;
+import org.hypertrace.core.grpcutils.context.RequestContext;
+
+class PlatformGrpcContextBuilder implements GrpcContextBuilder {
+
+  @Override
+  public RequestContext build(GraphQlRequestContext requestContext) {
+    Map<String, String> grpcHeaders =
+        this.mergeMaps(
+            requestContext.getTracingContextHeaders(),
+            this.flattenOptionalMap(
+                Map.of(
+                    AUTHORIZATION_HEADER, this.extractAuthorizationHeader(requestContext),
+                    TENANT_ID_HEADER_KEY, this.extractTenantId(requestContext))));
+
+    return this.build(grpcHeaders);
+  }
+
+  private RequestContext build(@Nonnull Map<String, String> headers) {
+    RequestContext platformContext = new RequestContext();
+    headers.forEach(platformContext::add);
+    return platformContext;
+  }
+
+  private Optional<String> extractAuthorizationHeader(GraphQlRequestContext requestContext) {
+    return requestContext.getAuthorizationHeader().filter(header -> header.length() > 0);
+  }
+
+  private Optional<String> extractTenantId(GraphQlRequestContext requestContext) {
+    return requestContext.getTenantId();
+  }
+
+  private Map<String, String> flattenOptionalMap(Map<String, Optional<String>> optionalMap) {
+    return optionalMap.entrySet().stream()
+        .map(entry -> entry.getValue().map(value -> Map.entry(entry.getKey(), value)))
+        .flatMap(Optional::stream)
+        .collect(Collectors.toUnmodifiableMap(Entry::getKey, Entry::getValue));
+  }
+
+  @SafeVarargs
+  private Map<String, String> mergeMaps(Map<String, String>... maps) {
+    return Arrays.stream(maps)
+        .map(Map::entrySet)
+        .flatMap(Collection::stream)
+        .collect(Collectors.toUnmodifiableMap(Entry::getKey, Entry::getValue));
+  }
+}