refactor: replace deprecated context fetch (#135)

Author: aaron-steinfeld

hypertrace-core-graphql-common-schema/src/main/java/org/hypertrace/core/graphql/common/fetcher/InjectableDataFetcher.java

@@ -1,5 +1,7 @@
 package org.hypertrace.core.graphql.common.fetcher;
 
+import static org.hypertrace.core.graphql.context.GraphQlRequestContext.contextFromEnvironment;
+
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.concurrent.CompletableFuture;
@@ -18,7 +20,7 @@ protected InjectableDataFetcher(
 
   @Override
   public final CompletableFuture<T> get(DataFetchingEnvironment environment) throws Exception {
-    return this.getOrCreateImplementation(environment.getContext()).get(environment);
+    return this.getOrCreateImplementation(contextFromEnvironment(environment)).get(environment);
   }
 
   private DataFetcher<CompletableFuture<T>> getOrCreateImplementation(

hypertrace-core-graphql-context/src/main/java/org/hypertrace/core/graphql/context/DefaultGraphQlRequestContextBuilder.java

@@ -51,6 +51,7 @@ private DefaultGraphQlRequestContext(HttpServletRequest request) {
       this.request = request;
       this.cachingKey =
           new DefaultContextualCacheKey(this, this.getTenantId().orElse(DEFAULT_CONTEXT_ID));
+      this.put(GraphQlRequestContext.class, this);
     }
 
     @Override

hypertrace-core-graphql-context/src/main/java/org/hypertrace/core/graphql/context/GraphQlRequestContext.java

@@ -2,6 +2,7 @@
 
 import graphql.kickstart.execution.context.GraphQLKickstartContext;
 import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
@@ -26,4 +27,8 @@ <T> DataFetcher<CompletableFuture<T>> constructDataFetcher(
   ContextualCachingKey getCachingKey();
 
   String getRequestId();
+
+  static GraphQlRequestContext contextFromEnvironment(DataFetchingEnvironment environment) {
+    return environment.getGraphQlContext().get(GraphQlRequestContext.class);
+  }
 }

hypertrace-core-graphql-log-event-schema/src/main/java/org/hypertrace/core/graphql/log/event/fetcher/LogEventFetcher.java

@@ -1,5 +1,7 @@
 package org.hypertrace.core.graphql.log.event.fetcher;
 
+import static org.hypertrace.core.graphql.context.GraphQlRequestContext.contextFromEnvironment;
+
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.concurrent.CompletableFuture;
@@ -30,7 +32,9 @@ static final class LogEventFetcherImpl
     public CompletableFuture<LogEventResultSet> get(DataFetchingEnvironment environment) {
       return this.requestBuilder
           .build(
-              environment.getContext(), environment.getArguments(), environment.getSelectionSet())
+              contextFromEnvironment(environment),
+              environment.getArguments(),
+              environment.getSelectionSet())
           .flatMap(this.logEventDao::getLogEvents)
           .toCompletionStage()
           .toCompletableFuture();

hypertrace-core-graphql-metadata-schema/src/main/java/org/hypertrace/core/graphql/metadata/fetcher/MetadataFetcher.java

@@ -1,5 +1,7 @@
 package org.hypertrace.core.graphql.metadata.fetcher;
 
+import static org.hypertrace.core.graphql.context.GraphQlRequestContext.contextFromEnvironment;
+
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.List;
@@ -30,7 +32,7 @@ static final class MetadataFetcherImpl
     @Override
     public CompletableFuture<List<AttributeMetadata>> get(DataFetchingEnvironment environment) {
       return this.attributeStore
-          .getAllExternal(environment.getContext())
+          .getAllExternal(contextFromEnvironment(environment))
           .flatMap(this.responseBuilder::build)
           .toCompletionStage()
           .toCompletableFuture();

hypertrace-core-graphql-metadata-schema/src/test/java/org/hypertrace/core/graphql/metadata/fetcher/MetadataFetcherTest.java

@@ -17,6 +17,7 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Answers;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 
@@ -25,7 +26,10 @@ class MetadataFetcherTest {
 
   @Mock MetadataResponseBuilder mockResponseBuilder;
   @Mock AttributeStore mockAttributeStore;
-  @Mock DataFetchingEnvironment mockDataFetchingEnvironment;
+
+  @Mock(answer = Answers.RETURNS_DEEP_STUBS)
+  DataFetchingEnvironment mockDataFetchingEnvironment;
+
   @Mock AttributeModel mockModel;
   @Mock AttributeMetadata mockMetadata;
   @Mock GraphQlRequestContext mockContext;
@@ -35,7 +39,8 @@ class MetadataFetcherTest {
   void beforeEach() {
     List<AttributeModel> mockModelResult = List.of(mockModel);
     List<AttributeMetadata> mockMetadataResult = List.of(mockMetadata);
-    when(this.mockDataFetchingEnvironment.getContext()).thenReturn(this.mockContext);
+    when(this.mockDataFetchingEnvironment.getGraphQlContext().get(GraphQlRequestContext.class))
+        .thenReturn(this.mockContext);
     when(this.mockAttributeStore.getAllExternal(eq(this.mockContext)))
         .thenReturn(Single.just(mockModelResult));
     when(this.mockResponseBuilder.build(eq(mockModelResult)))

hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/fetcher/ExportSpanFetcher.java

@@ -1,5 +1,7 @@
 package org.hypertrace.core.graphql.span.fetcher;
 
+import static org.hypertrace.core.graphql.context.GraphQlRequestContext.contextFromEnvironment;
+
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import io.reactivex.rxjava3.core.Single;
@@ -34,7 +36,7 @@ static final class ExportSpanFetcherImpl
     public CompletableFuture<ExportSpanResult> get(DataFetchingEnvironment environment) {
       Single<SpanRequest> spanRequest =
           this.requestBuilder.build(
-              environment.getContext(),
+              contextFromEnvironment(environment),
               environment.getArguments(),
               SpanAttributes.SPAN_ATTRIBUTES,
               LogEventAttributes.LOG_EVENT_ATTRIBUTES);

hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/fetcher/SpanFetcher.java

@@ -1,5 +1,7 @@
 package org.hypertrace.core.graphql.span.fetcher;
 
+import static org.hypertrace.core.graphql.context.GraphQlRequestContext.contextFromEnvironment;
+
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.concurrent.CompletableFuture;
@@ -29,7 +31,9 @@ static final class SpanFetcherImpl implements DataFetcher<CompletableFuture<Span
     public CompletableFuture<SpanResultSet> get(DataFetchingEnvironment environment) {
       return this.requestBuilder
           .build(
-              environment.getContext(), environment.getArguments(), environment.getSelectionSet())
+              contextFromEnvironment(environment),
+              environment.getArguments(),
+              environment.getSelectionSet())
           .flatMap(this.spanDao::getSpans)
           .toCompletionStage()
           .toCompletableFuture();

hypertrace-core-graphql-trace-schema/src/main/java/org/hypertrace/core/graphql/trace/fetcher/TraceFetcher.java

@@ -1,5 +1,7 @@
 package org.hypertrace.core.graphql.trace.fetcher;
 
+import static org.hypertrace.core.graphql.context.GraphQlRequestContext.contextFromEnvironment;
+
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.concurrent.CompletableFuture;
@@ -29,7 +31,9 @@ static final class TraceFetcherImpl implements DataFetcher<CompletableFuture<Tra
     public CompletableFuture<TraceResultSet> get(DataFetchingEnvironment environment) {
       return this.requestBuilder
           .build(
-              environment.getContext(), environment.getArguments(), environment.getSelectionSet())
+              contextFromEnvironment(environment),
+              environment.getArguments(),
+              environment.getSelectionSet())
           .flatMap(this.traceDao::getTraces)
           .toCompletionStage()
           .toCompletableFuture();