org.dataloader.DataLoaderRegistry Java Examples

@Override
public void setUp() throws Exception {
  super.setUp();

  BatchLoaderWithContext<String, User> userBatchLoader = VertxBatchLoader.create(
    (keys, environment, listPromise) -> {
      if (batchloaderInvoked.compareAndSet(false, true)) {
        listPromise.complete(keys
          .stream()
          .map(testData.users::get)
          .collect(toList())
        );
      } else {
        listPromise.fail(new IllegalStateException());
      }
    }
  );

  graphQLHandler.dataLoaderRegistry(rc -> {
    DataLoader<String, User> userDataLoader = DataLoader.newDataLoader(userBatchLoader);
    return new DataLoaderRegistry().register("user", userDataLoader);
  });
}
 

@Override
public void setUp() throws Exception {
  super.setUp();

  MappedBatchLoaderWithContext<String, User> userBatchLoader = VertxMappedBatchLoader.create(
    (keys, environment, mapPromise) -> {
      if (batchloaderInvoked.compareAndSet(false, true)) {
        mapPromise.complete(keys
          .stream()
          .map(testData.users::get)
          .collect(toMap(User::getId, Function.identity()))
        );
      } else {
        mapPromise.fail(new IllegalStateException());
      }
    }
  );

  graphQLHandler.dataLoaderRegistry(rc -> {
    DataLoader<String, User> userDataLoader = DataLoader.newMappedDataLoader(userBatchLoader);
    return new DataLoaderRegistry().register("user", userDataLoader);
  });
}
 

@Provides
DataLoaderRegistry dataLoaderRegistry(BookServiceGrpc.BookServiceFutureStub bookService) {

  // TODO: Use multibinder to modularize this, or automate this somehow
  BatchLoader<String, Book> bookBatchLoader =
      keys -> {
        ListenableFuture<List<Book>> listenableFuture =
            Futures.transform(
                bookService.listBooks(
                    ListBooksRequest.newBuilder()
                        .addAllIds(keys)
                        .setPageSize(keys.size())
                        .build()),
                resp -> resp.getBooksList(),
                MoreExecutors.directExecutor());
        return FutureConverter.toCompletableFuture(listenableFuture);
      };

  DataLoaderRegistry registry = new DataLoaderRegistry();
  registry.register("books", new DataLoader<>(bookBatchLoader));
  return registry;
}
 

static CompletableFuture<ExecutionResult> executeGraphQLRequest(
    GraphQLRequest request,
    ByteBuf byteBuf,
    DataLoaderRegistry registry,
    GraphQLSchema graphQLSchema,
    Instrumentation instrumentation) {
  ExecutionInput.Builder builder =
      ExecutionInput.newExecutionInput()
          .query(request.getQuery())
          .operationName(request.getOperationName())
          .variables(request.getVariables())
          .context(byteBuf)
          .executionId(ExecutionId.generate());

  if (registry != null) {
    builder.dataLoaderRegistry(registry);
  }

  ExecutionInput executionInput = builder.build();

  return GraphQL.newGraphQL(graphQLSchema)
      .instrumentation(instrumentation)
      .build()
      .executeAsync(executionInput);
}
 

@Query("getBook")
@RelayNode
ListenableFuture<Book> getBook(
    GetBookRequest request, DataFetchingEnvironment dataFetchingEnvironment) {
  return FutureConverter.toListenableFuture(
      dataFetchingEnvironment
          .<DataLoaderRegistry>getContext()
          .<String, Book>getDataLoader("books")
          .load(request.getId()));
}
 

private Future<JsonObject> execute(RoutingContext rc, GraphQLQuery query) {
  ExecutionInput.Builder builder = ExecutionInput.newExecutionInput();

  builder.query(query.getQuery());
  String operationName = query.getOperationName();
  if (operationName != null) {
    builder.operationName(operationName);
  }
  Map<String, Object> variables = query.getVariables();
  if (variables != null) {
    builder.variables(variables);
  }

  Function<RoutingContext, Object> qc;
  synchronized (this) {
    qc = queryContextFactory;
  }
  builder.context(qc.apply(rc));

  Function<RoutingContext, DataLoaderRegistry> dlr;
  synchronized (this) {
    dlr = dataLoaderRegistryFactory;
  }
  DataLoaderRegistry registry = dlr.apply(rc);
  if (registry != null) {
    builder.dataLoaderRegistry(registry);
  }

  Function<RoutingContext, Locale> l;
  synchronized (this) {
    l = localeFactory;
  }
  Locale locale = l.apply(rc);
  if (locale != null) {
    builder.locale(locale);
  }

  return Future.fromCompletionStage(graphQL.executeAsync(builder.build()), rc.vertx().getOrCreateContext())
    .map(executionResult -> new JsonObject(executionResult.toSpecification()));
}
 

@Override
public void doPost(HttpServletRequest req, HttpServletResponse resp) throws IOException {
  Map<String, Object> json = readJson(req);
  String query = (String) json.get("query");
  if (query == null) {
    resp.setStatus(400);
    return;
  }
  String operationName = (String) json.get("operationName");
  Map<String, Object> variables = getVariables(json.get("variables"));

  DataLoaderRegistry dataLoaderRegistry = registryProvider.get();
  GraphQL graphql = GraphQL.newGraphQL(schema).instrumentation(INSTRUMENTATION).build();

  ExecutionInput executionInput =
      ExecutionInput.newExecutionInput()
          .query(query)
          .operationName(operationName)
          .variables(variables)
          .context(dataLoaderRegistry)
          .dataLoaderRegistry(dataLoaderRegistry)
          .build();
  ExecutionResult executionResult = graphql.execute(executionInput);
  resp.setContentType("application/json");
  resp.setStatus(HttpServletResponse.SC_OK);
  GSON.toJson(executionResult.toSpecification(), resp.getWriter());
  logger.info("stats: " + dataLoaderRegistry.getStatistics());
}
 

/**
 * Default constructor.
 *
 * @param graphQL                         the {@link GraphQL} instance
 * @param graphQLExecutionInputCustomizer the {@link GraphQLExecutionInputCustomizer} instance
 * @param dataLoaderRegistry              the {@link DataLoaderRegistry} instance
 */
public DefaultGraphQLInvocation(
        GraphQL graphQL,
        @Nullable GraphQLExecutionInputCustomizer graphQLExecutionInputCustomizer,
        @Nullable Provider<DataLoaderRegistry> dataLoaderRegistry) {
    this.graphQL = graphQL;
    this.graphQLExecutionInputCustomizer = graphQLExecutionInputCustomizer;
    this.dataLoaderRegistry = dataLoaderRegistry;
}
 

@SchemaModification(addField = "books", onType = Shelf.class)
ListenableFuture<List<Book>> shelfToBooks(Shelf shelf, DataFetchingEnvironment environment) {
  return FutureConverter.toListenableFuture(
      environment
          .<DataLoaderRegistry>getContext()
          .<String, Book>getDataLoader("books")
          .loadMany(shelf.getBookIdsList()));
}
 

@Override
protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws IOException {

  DataLoaderRegistry dataLoaderRegistry = registryProvider.get();
  Instrumentation instrumentation =
      new ChainedInstrumentation(
          Arrays.asList(
              GuavaListenableFutureSupport.listenableFutureInstrumentation(),
              new TracingInstrumentation()));
  GraphQL graphql = GraphQL.newGraphQL(schema).instrumentation(instrumentation).build();

  Map<String, Object> json = readJson(req);
  String query = (String) json.get("query");
  if (query == null) {
    resp.setStatus(400);
    return;
  }
  String operationName = (String) json.get("operationName");
  Map<String, Object> variables = getVariables(json.get("variables"));

  ExecutionInput executionInput =
      ExecutionInput.newExecutionInput()
          .query(query)
          .operationName(operationName)
          .variables(variables)
          .dataLoaderRegistry(dataLoaderRegistry)
          .context(dataLoaderRegistry)
          .build();
  ExecutionResult executionResult = graphql.execute(executionInput);
  resp.setContentType("application/json");
  resp.setStatus(HttpServletResponse.SC_OK);
  GSON.toJson(executionResult.toSpecification(), resp.getWriter());
  logger.info("stats: " + dataLoaderRegistry.getStatistics());
}
 

@Override
public DataLoaderRegistry getDataLoaderRegistry() {
    return delegate.getDataLoaderRegistry();
}
 

Context(DataLoaderRegistry dataLoaderRegistry) {
    this.dataLoaderRegistry = dataLoaderRegistry;
}
 

@Autowired
public ContextProvider(DataLoaderRegistry dataLoaderRegistry) {
    this.dataLoaderRegistry = dataLoaderRegistry;
}
 

private void start(ApolloWSMessage message) {
  String opId = message.content().getString("id");

  // Unsubscribe if it's subscribed
  if (subscriptions.containsKey(opId)) {
    stop(opId);
  }

  GraphQLQuery payload = new GraphQLQuery(message.content().getJsonObject("payload"));
  ExecutionInput.Builder builder = ExecutionInput.newExecutionInput();
  builder.query(payload.getQuery());

  builder.context(apolloWSHandler.getQueryContext().apply(message));

  DataLoaderRegistry registry = apolloWSHandler.getDataLoaderRegistry().apply(message);
  if (registry != null) {
    builder.dataLoaderRegistry(registry);
  }

  Locale locale = apolloWSHandler.getLocale().apply(message);
  if (locale != null) {
    builder.locale(locale);
  }

  String operationName = payload.getOperationName();
  if (operationName != null) {
    builder.operationName(operationName);
  }
  Map<String, Object> variables = payload.getVariables();
  if (variables != null) {
    builder.variables(variables);
  }

  apolloWSHandler.getGraphQL().executeAsync(builder).whenCompleteAsync((executionResult, throwable) -> {
    if (throwable == null) {
      if (executionResult.getData() instanceof Publisher) {
        subscribe(opId, executionResult);
      } else {
        sendMessage(opId, DATA, new JsonObject(executionResult.toSpecification()));
        sendMessage(opId, COMPLETE, null);
      }
    } else {
      if (log.isDebugEnabled()) {
        log.debug("Failed to execute GraphQL query, opId=" + opId, throwable);
      }
      sendMessage(opId, ERROR, toJsonObject(throwable));
    }
  }, context);
}
 

synchronized Function<ApolloWSMessage, DataLoaderRegistry> getDataLoaderRegistry() {
  return dataLoaderRegistryFactory;
}
 

@Override
public synchronized ApolloWSHandler dataLoaderRegistry(Function<ApolloWSMessage, DataLoaderRegistry> factory) {
  dataLoaderRegistryFactory = factory != null ? factory : DEFAULT_DATA_LOADER_REGISTRY_FACTORY;
  return this;
}
 

@Autowired
public GraphQLProvider(DataLoaderRegistry dataLoaderRegistry, StarWarsWiring starWarsWiring) {
    this.dataLoaderRegistry = dataLoaderRegistry;
    this.starWarsWiring = starWarsWiring;
}
 

@Override
public synchronized GraphQLHandler dataLoaderRegistry(Function<RoutingContext, DataLoaderRegistry> factory) {
  dataLoaderRegistryFactory = factory != null ? factory : DEFAULT_DATA_LOADER_REGISTRY_FACTORY;
  return this;
}
 

public DataLoaderRegistry getDataLoaderRegistry() {
    return dataLoaderRegistry;
}
 

public Context() {
    this.dataLoaderRegistry = new DataLoaderRegistry();
    dataLoaderRegistry.register("characters", newCharacterDataLoader());
}
 

private void handleStarWars(HttpServletRequest httpRequest, HttpServletResponse httpResponse) throws IOException {
    //
    // this builds out the parameters we need like the graphql query from the http request
    QueryParameters parameters = QueryParameters.from(httpRequest);
    if (parameters.getQuery() == null) {
        //
        // how to handle nonsensical requests is up to your application
        httpResponse.setStatus(400);
        return;
    }

    ExecutionInput.Builder executionInput = newExecutionInput()
            .query(parameters.getQuery())
            .operationName(parameters.getOperationName())
            .variables(parameters.getVariables());

    //
    // the context object is something that means something to down stream code.  It is instructions
    // from yourself to your other code such as DataFetchers.  The engine passes this on unchanged and
    // makes it available to inner code
    //
    // the graphql guidance says  :
    //
    //  - GraphQL should be placed after all authentication middleware, so that you
    //  - have access to the same session and user information you would in your
    //  - HTTP endpoint handlers.
    //
    StarWarsWiring.Context context = new StarWarsWiring.Context();
    executionInput.context(context);

    //
    // you need a schema in order to execute queries
    GraphQLSchema schema = buildStarWarsSchema();

    //
    // This example uses the DataLoader technique to ensure that the most efficient
    // loading of data (in this case StarWars characters) happens.  We pass that to data
    // fetchers via the graphql context object.
    //
    DataLoaderRegistry dataLoaderRegistry = context.getDataLoaderRegistry();


    DataLoaderDispatcherInstrumentation dlInstrumentation =
            new DataLoaderDispatcherInstrumentation(dataLoaderRegistry, newOptions().includeStatistics(true));

    Instrumentation instrumentation = new ChainedInstrumentation(
            asList(new TracingInstrumentation(), dlInstrumentation)
    );

    // finally you build a runtime graphql object and execute the query
    GraphQL graphQL = GraphQL
            .newGraphQL(schema)
            // instrumentation is pluggable
            .instrumentation(instrumentation)
            .build();
    ExecutionResult executionResult = graphQL.execute(executionInput.build());

    returnAsJson(httpResponse, executionResult);
}
 

public StarWarsWiring() {
    this.dataLoaderRegistry = new DataLoaderRegistry();
    dataLoaderRegistry.register("characters", newCharacterDataLoader());
}
 

@Bean
public DataLoaderRegistry getDataLoaderRegistry() {
    return dataLoaderRegistry;
}
 

GraphQLServerRequestStream(
    DataLoaderRegistry registry, Instrumentation instrumentation, GraphQLSchema graphQLSchema) {
  this.registry = registry;
  this.instrumentation = instrumentation;
  this.graphQLSchema = graphQLSchema;
}
 

GraphQLServerRequestResponse(
    DataLoaderRegistry registry, Instrumentation instrumentation, GraphQLSchema graphQLSchema) {
  this.registry = registry;
  this.instrumentation = instrumentation;
  this.graphQLSchema = graphQLSchema;
}
 

@Override
public I dataLoadRegister(DataLoaderRegistry dataLoadRegistry) {
  this.dataLoadRegistry = Objects.requireNonNull(dataLoadRegistry);
  return this;
}
 

public void dataLoaderRegistry(GraphQL graphQL, BatchLoaderWithContext<String, Link> linksBatchLoader) {
  GraphQLHandler handler = GraphQLHandler.create(graphQL).dataLoaderRegistry(rc -> {

    DataLoader<String, Link> linkDataLoader = DataLoader.newDataLoader(linksBatchLoader);

    return new DataLoaderRegistry().register("link", linkDataLoader);

  });
}
 

/**
 * Customize the {@link DataLoaderRegistry}.
 * The provided {@code factory} method will be invoked for each incoming GraphQL request.
 *
 * @return a reference to this, so the API can be used fluently
 */
@Fluent
@GenIgnore(GenIgnore.PERMITTED_TYPE)
GraphQLHandler dataLoaderRegistry(Function<RoutingContext, DataLoaderRegistry> factory);
 

/**
 * Customize the {@link DataLoaderRegistry}.
 * The provided {@code factory} method will be invoked for each incoming GraphQL request.
 *
 * @return a reference to this, so the API can be used fluently
 */
@Fluent
@GenIgnore(GenIgnore.PERMITTED_TYPE)
ApolloWSHandler dataLoaderRegistry(Function<ApolloWSMessage, DataLoaderRegistry> factory);
 

I<O> dataLoadRegister(DataLoaderRegistry dataLoadRegistry);