Skip to content

Commit

Permalink
Merge pull request #1318 from phillip-kruger/main
Browse files Browse the repository at this point in the history
Small cleanup and fix for context issue
  • Loading branch information
phillip-kruger authored Mar 15, 2022
2 parents 8d6d9a6 + 30ec86c commit 09cc6de
Show file tree
Hide file tree
Showing 10 changed files with 198 additions and 106 deletions.
58 changes: 58 additions & 0 deletions server/api/src/main/java/io/smallrye/graphql/api/Context.java
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,64 @@ default <T> T getArgumentOrDefault(String name, T defaultValue) {
*/
public Map<String, Object> getArguments();

/**
* Check if a certain meta data is available
*
* @param <K>
* @param key
* @return
*/
public <K> boolean hasMetaData(K key);

/**
* Allow getting custom user values.
*
* @param <K> the key
* @param <V> the value
* @param key
* @return
*/
public <K, V> V getMetaData(K key);

/**
* Allow setting custom user values
*
* @param <K>
* @param <V>
* @param key
* @param value
*/
public <K, V> void putMetaData(K key, V value);

/**
* Check if a certain local meta data is available
*
* @param <K>
* @param key
* @return
*/
public <K> boolean hasLocalMetaData(K key);

/**
* Allow getting custom local user values.
*
* @param <K> the key
* @param <V> the value
* @param key
* @return
*/
public <K, V> V getLocalMetaData(K key);

/**
* Allow setting custom local user values
*
* @param <K>
* @param <V>
* @param key
* @param value
*/
public <K, V> void putLocalMetaData(K key, V value);

default boolean hasSource() {
Object o = getSource();
return o != null;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,6 @@
import javax.annotation.Priority;
import javax.enterprise.inject.spi.CDI;

import graphql.ExecutionInput;
import graphql.GraphQLContext;
import graphql.schema.DataFetchingEnvironment;
import graphql.schema.GraphQLList;
import graphql.schema.GraphQLNamedType;
import graphql.schema.GraphQLNonNull;
import graphql.schema.GraphQLType;
import io.opentracing.Scope;
import io.opentracing.Span;
import io.opentracing.Tracer;
Expand Down Expand Up @@ -54,9 +47,8 @@ public class TracingService implements EventingService {

@Override
public void beforeExecute(Context context) {
ExecutionInput executionInput = context.unwrap(ExecutionInput.class);
// FIXME: if operationName is not set in the request explicitly, this is empty
String operationName = getOperationName(executionInput);
String operationName = getOperationName(context);

Span span = getTracer().buildSpan(operationName)
.asChildOf(getTracer().activeSpan())
Expand Down Expand Up @@ -92,10 +84,9 @@ public void errorExecute(String executionId, Throwable t) {

@Override
public void beforeDataFetch(Context context) {
final DataFetchingEnvironment env = context.unwrap(DataFetchingEnvironment.class);
Span parentSpan = getParentSpan(getTracer(), env);
Span parentSpan = getParentSpan(getTracer(), context);

Span span = getTracer().buildSpan(getOperationName(env))
Span span = getTracer().buildSpan(getOperationNameForParentType(context))
.asChildOf(parentSpan)
.withTag("graphql.executionId", context.getExecutionId())
.withTag("graphql.operationType", getOperationNameString(context.getOperationType()))
Expand All @@ -106,9 +97,8 @@ public void beforeDataFetch(Context context) {
.start();
Scope scope = tracer.activateSpan(span);

GraphQLContext graphQLContext = env.getContext();
graphQLContext.put(SPAN_CLASS, parentSpan);
graphQLContext.put(SCOPE_CLASS, scope);
context.putMetaData(SPAN_CLASS, parentSpan);
context.putMetaData(SCOPE_CLASS, scope);
}

// FIXME: is the fetcher is asynchronous, this typically ends its span before
Expand All @@ -118,7 +108,7 @@ public void beforeDataFetch(Context context) {
@Override
public void afterDataFetch(Context context) {
Span span = tracer.activeSpan();
Scope scope = ((GraphQLContext) context.unwrap(DataFetchingEnvironment.class).getContext()).get(SCOPE_CLASS);
Scope scope = context.getMetaData(SCOPE_CLASS);
scope.close();
span.finish();
}
Expand All @@ -143,15 +133,11 @@ private Tracer getTracer() {
return tracer;
}

private Span getParentSpan(Tracer tracer, final DataFetchingEnvironment env) {
final GraphQLContext localContext = env.getLocalContext();
if (localContext != null && localContext.hasKey(SPAN_CLASS)) {
return localContext.get(SPAN_CLASS);
}

final GraphQLContext rootContext = env.getContext();
if (rootContext != null && rootContext.hasKey(SPAN_CLASS)) {
return rootContext.get(SPAN_CLASS);
private Span getParentSpan(Tracer tracer, final Context context) {
if (context != null && context.hasLocalMetaData(SPAN_CLASS)) {
return context.getLocalMetaData(SPAN_CLASS);
} else if (context != null && context.hasMetaData(SPAN_CLASS)) {
return context.getMetaData(SPAN_CLASS);
}

return tracer.activeSpan();
Expand All @@ -172,33 +158,19 @@ private void logError(Span span, Throwable throwable) {
span.setTag("error", true);
}

private String getOperationName(final DataFetchingEnvironment env) {
String parent = getName(env.getParentType());

String name = PREFIX + ":" + parent + "." + env.getField().getName();

private String getOperationNameForParentType(Context context) {
String parent = context.getParentTypeName().orElse(EMPTY);
String name = PREFIX + ":" + parent + "." + context.getFieldName();
return name;
}

private static String getOperationName(ExecutionInput executionInput) {
if (executionInput != null && executionInput.getOperationName() != null
&& !executionInput.getOperationName().isEmpty()) {
return PREFIX + ":" + executionInput.getOperationName();
private static String getOperationName(Context context) {
if (context.getOperationName().isPresent()) {
return PREFIX + ":" + context.getOperationName().get();
}
return PREFIX;
}

private String getName(GraphQLType graphQLType) {
if (graphQLType instanceof GraphQLNamedType) {
return ((GraphQLNamedType) graphQLType).getName();
} else if (graphQLType instanceof GraphQLNonNull) {
return getName(((GraphQLNonNull) graphQLType).getWrappedType());
} else if (graphQLType instanceof GraphQLList) {
return getName(((GraphQLList) graphQLType).getWrappedType());
}
return EMPTY;
}

private String getOperationNameString(List<String> types) {
return String.join(UNDERSCORE, types);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.atomic.AtomicLong;

import javax.json.JsonObject;
Expand All @@ -17,7 +16,6 @@
import org.dataloader.DataLoaderFactory;
import org.dataloader.DataLoaderOptions;
import org.dataloader.DataLoaderRegistry;
import org.eclipse.microprofile.context.ThreadContext;

import graphql.ExecutionInput;
import graphql.ExecutionInput.Builder;
Expand All @@ -27,7 +25,6 @@
import graphql.execution.ExecutionStrategy;
import graphql.execution.SubscriptionExecutionStrategy;
import graphql.schema.GraphQLSchema;
import io.smallrye.graphql.api.Context;
import io.smallrye.graphql.bootstrap.DataFetcherFactory;
import io.smallrye.graphql.execution.context.SmallRyeBatchLoaderContextProvider;
import io.smallrye.graphql.execution.context.SmallRyeContext;
Expand All @@ -38,6 +35,7 @@
import io.smallrye.graphql.schema.model.Operation;
import io.smallrye.graphql.spi.config.Config;
import io.smallrye.graphql.spi.config.LogPayloadOption;
import io.smallrye.mutiny.Uni;

/**
* Executing the GraphQL request
Expand Down Expand Up @@ -159,6 +157,10 @@ public void execute(JsonObject jsonInput, Map<String, Object> context, Execution
// Operation name
smallRyeContext.getOperationName().ifPresent(executionBuilder::operationName);

// Context
context.put(ContextHelper.CONTEXT, smallRyeContext);
executionBuilder.graphQLContext(context);

// DataLoaders
if (batchOperations != null && !batchOperations.isEmpty()) {
DataLoaderRegistry dataLoaderRegistry = getDataLoaderRegistry(batchOperations);
Expand All @@ -169,19 +171,16 @@ public void execute(JsonObject jsonInput, Map<String, Object> context, Execution

// Update context with execution data
smallRyeContext = smallRyeContext.withDataFromExecution(executionInput, queryCache);

// Context
context.put(ContextHelper.CONTEXT, smallRyeContext);
executionInput.getGraphQLContext().putAll(context);
executionInput.getGraphQLContext().put(ContextHelper.CONTEXT, smallRyeContext);

// Notify before
eventEmitter.fireBeforeExecute(smallRyeContext);

// Execute
if (async) {
writeAsync(g, executionInput, context, writer);
writeAsync(g, executionInput, context, smallRyeContext, writer);
} else {
writeSync(g, executionInput, context, writer);
writeSync(g, executionInput, context, smallRyeContext, writer);
}
} else {
log.noGraphQLMethodsFound();
Expand All @@ -195,42 +194,41 @@ public void execute(JsonObject jsonInput, Map<String, Object> context, Execution
private void writeAsync(GraphQL graphQL,
ExecutionInput executionInput,
Map<String, Object> context,
SmallRyeContext smallRyeContext,
ExecutionResponseWriter writer) {

// Execute
ThreadContext threadContext = ThreadContext.builder().build();
CompletionStage<ExecutionResult> executionResult = threadContext
.withContextCapture(graphQL.executeAsync(executionInput));
Uni.createFrom().completionStage(() -> graphQL.executeAsync(executionInput))

executionResult.whenComplete((t, u) -> {
executionInput.getGraphQLContext().putAll(context);
.subscribe().with(executionResult -> {
executionInput.getGraphQLContext().putAll(context);

SmallRyeContext smallryeContext = (SmallRyeContext) context.get(ContextHelper.CONTEXT);
SmallRyeContext.setContext(smallryeContext);
SmallRyeContext.setContext(smallRyeContext);

// Notify after
eventEmitter.fireAfterExecute(smallryeContext);
// Notify after
eventEmitter.fireAfterExecute(smallRyeContext);

ExecutionResponse executionResponse = new ExecutionResponse(t);
if (!payloadOption.equals(LogPayloadOption.off)) {
log.payloadOut(executionResponse.toString());
}
writer.write(executionResponse);
ExecutionResponse executionResponse = new ExecutionResponse(executionResult);
if (!payloadOption.equals(LogPayloadOption.off)) {
log.payloadOut(executionResponse.toString());
}
writer.write(executionResponse);

if (u != null) {
writer.fail(u);
}
});
}, failure -> {
if (failure != null) {
writer.fail(failure);
}
});
}

private void writeSync(GraphQL g,
ExecutionInput executionInput,
Map<String, Object> context,
SmallRyeContext smallRyeContext,
ExecutionResponseWriter writer) {
try {
ExecutionResult executionResult = g.execute(executionInput);
// Notify after
eventEmitter.fireAfterExecute((Context) context.get(ContextHelper.CONTEXT));
eventEmitter.fireAfterExecute(smallRyeContext);

ExecutionResponse executionResponse = new ExecutionResponse(executionResult);
if (!payloadOption.equals(LogPayloadOption.off)) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,6 @@
import org.dataloader.BatchLoaderContextProvider;
import org.dataloader.DataLoader;

import graphql.schema.DataFetchingEnvironment;

// hacky way to pass the GraphQL Context from data from data fetchers to BatchLoaderEnvironment
public class SmallRyeBatchLoaderContextProvider implements BatchLoaderContextProvider {

Expand All @@ -28,14 +26,14 @@ public static SmallRyeBatchLoaderContextProvider getForDataLoader(DataLoader dat
return INSTANCES.get(dataLoader);
}

private AtomicReference<DataFetchingEnvironment> current = new AtomicReference<>();
private AtomicReference<SmallRyeContext> current = new AtomicReference<>();

public void set(DataFetchingEnvironment dfe) {
public void set(SmallRyeContext dfe) {
current.set(dfe);
}

@Override
public DataFetchingEnvironment getContext() {
public SmallRyeContext getContext() {
return current.getAndSet(null);
}
}
Loading

0 comments on commit 09cc6de

Please sign in to comment.