diff --git a/hypertrace-core-graphql-common-schema/src/main/java/org/hypertrace/core/graphql/common/fetcher/InjectableDataFetcher.java b/hypertrace-core-graphql-common-schema/src/main/java/org/hypertrace/core/graphql/common/fetcher/InjectableDataFetcher.java index a910eabd..89561de5 100644 --- a/hypertrace-core-graphql-common-schema/src/main/java/org/hypertrace/core/graphql/common/fetcher/InjectableDataFetcher.java +++ b/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 get(DataFetchingEnvironment environment) throws Exception { - return this.getOrCreateImplementation(environment.getContext()).get(environment); + return this.getOrCreateImplementation(contextFromEnvironment(environment)).get(environment); } private DataFetcher> getOrCreateImplementation( diff --git a/hypertrace-core-graphql-context/src/main/java/org/hypertrace/core/graphql/context/DefaultGraphQlRequestContextBuilder.java b/hypertrace-core-graphql-context/src/main/java/org/hypertrace/core/graphql/context/DefaultGraphQlRequestContextBuilder.java index a5f452c9..91f9fdc1 100644 --- a/hypertrace-core-graphql-context/src/main/java/org/hypertrace/core/graphql/context/DefaultGraphQlRequestContextBuilder.java +++ b/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 diff --git a/hypertrace-core-graphql-context/src/main/java/org/hypertrace/core/graphql/context/GraphQlRequestContext.java b/hypertrace-core-graphql-context/src/main/java/org/hypertrace/core/graphql/context/GraphQlRequestContext.java index 86a62cb7..d943ea0f 100644 --- a/hypertrace-core-graphql-context/src/main/java/org/hypertrace/core/graphql/context/GraphQlRequestContext.java +++ b/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 @@ DataFetcher> constructDataFetcher( ContextualCachingKey getCachingKey(); String getRequestId(); + + static GraphQlRequestContext contextFromEnvironment(DataFetchingEnvironment environment) { + return environment.getGraphQlContext().get(GraphQlRequestContext.class); + } } diff --git a/hypertrace-core-graphql-log-event-schema/src/main/java/org/hypertrace/core/graphql/log/event/fetcher/LogEventFetcher.java b/hypertrace-core-graphql-log-event-schema/src/main/java/org/hypertrace/core/graphql/log/event/fetcher/LogEventFetcher.java index 51701bfd..d6b5a3a5 100644 --- a/hypertrace-core-graphql-log-event-schema/src/main/java/org/hypertrace/core/graphql/log/event/fetcher/LogEventFetcher.java +++ b/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 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(); diff --git a/hypertrace-core-graphql-metadata-schema/src/main/java/org/hypertrace/core/graphql/metadata/fetcher/MetadataFetcher.java b/hypertrace-core-graphql-metadata-schema/src/main/java/org/hypertrace/core/graphql/metadata/fetcher/MetadataFetcher.java index 5ad55b33..260f72ab 100644 --- a/hypertrace-core-graphql-metadata-schema/src/main/java/org/hypertrace/core/graphql/metadata/fetcher/MetadataFetcher.java +++ b/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> get(DataFetchingEnvironment environment) { return this.attributeStore - .getAllExternal(environment.getContext()) + .getAllExternal(contextFromEnvironment(environment)) .flatMap(this.responseBuilder::build) .toCompletionStage() .toCompletableFuture(); diff --git a/hypertrace-core-graphql-metadata-schema/src/test/java/org/hypertrace/core/graphql/metadata/fetcher/MetadataFetcherTest.java b/hypertrace-core-graphql-metadata-schema/src/test/java/org/hypertrace/core/graphql/metadata/fetcher/MetadataFetcherTest.java index 3a677c79..effe854d 100644 --- a/hypertrace-core-graphql-metadata-schema/src/test/java/org/hypertrace/core/graphql/metadata/fetcher/MetadataFetcherTest.java +++ b/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 mockModelResult = List.of(mockModel); List 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))) diff --git a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/fetcher/ExportSpanFetcher.java b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/fetcher/ExportSpanFetcher.java index 96249b8e..761e217e 100644 --- a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/fetcher/ExportSpanFetcher.java +++ b/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 get(DataFetchingEnvironment environment) { Single spanRequest = this.requestBuilder.build( - environment.getContext(), + contextFromEnvironment(environment), environment.getArguments(), SpanAttributes.SPAN_ATTRIBUTES, LogEventAttributes.LOG_EVENT_ATTRIBUTES); diff --git a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/fetcher/SpanFetcher.java b/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/fetcher/SpanFetcher.java index 6bfbf4f8..9079305f 100644 --- a/hypertrace-core-graphql-span-schema/src/main/java/org/hypertrace/core/graphql/span/fetcher/SpanFetcher.java +++ b/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 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(); diff --git a/hypertrace-core-graphql-trace-schema/src/main/java/org/hypertrace/core/graphql/trace/fetcher/TraceFetcher.java b/hypertrace-core-graphql-trace-schema/src/main/java/org/hypertrace/core/graphql/trace/fetcher/TraceFetcher.java index c4b0aed7..8f9f3f02 100644 --- a/hypertrace-core-graphql-trace-schema/src/main/java/org/hypertrace/core/graphql/trace/fetcher/TraceFetcher.java +++ b/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 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();