org.dataloader.DataLoader Java Examples

private void keyContextExample() {
    DataLoaderOptions options = DataLoaderOptions.newOptions()
            .setBatchLoaderContextProvider(() -> SecurityCtx.getCallingUserCtx());

    BatchLoaderWithContext<String, String> batchLoader = new BatchLoaderWithContext<String, String>() {
        @Override
        public CompletionStage<List<String>> load(List<String> keys, BatchLoaderEnvironment environment) {
            SecurityCtx callCtx = environment.getContext();
            //
            // this is the load context objects in map form by key
            // in this case [ keyA : contextForA, keyB : contextForB ]
            //
            Map<Object, Object> keyContexts = environment.getKeyContexts();
            //
            // this is load context in list form
            //
            // in this case [ contextForA, contextForB ]
            return callDatabaseForResults(callCtx, keys);
        }
    };

    DataLoader<String, String> loader = DataLoader.newDataLoader(batchLoader, options);
    loader.load("keyA", "contextForA");
    loader.load("keyB", "contextForB");
}
 

@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);
  });
}
 

@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;
}
 

@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);
  });
}
 

private void callContextExample() {
    DataLoaderOptions options = DataLoaderOptions.newOptions()
            .setBatchLoaderContextProvider(() -> SecurityCtx.getCallingUserCtx());

    BatchLoaderWithContext<String, String> batchLoader = new BatchLoaderWithContext<String, String>() {
        @Override
        public CompletionStage<List<String>> load(List<String> keys, BatchLoaderEnvironment environment) {
            SecurityCtx callCtx = environment.getContext();
            return callDatabaseForResults(callCtx, keys);
        }
    };

    DataLoader<String, String> loader = DataLoader.newDataLoader(batchLoader, options);
}
 

private void tryBatcLoader() {
    DataLoader<String, User> dataLoader = DataLoader.newDataLoaderWithTry(new BatchLoader<String, Try<User>>() {
        @Override
        public CompletionStage<List<Try<User>>> load(List<String> keys) {
            return CompletableFuture.supplyAsync(() -> {
                List<Try<User>> users = new ArrayList<>();
                for (String key : keys) {
                    Try<User> userTry = loadUser(key);
                    users.add(userTry);
                }
                return users;
            });
        }
    });
}
 

public DataLoader<String, Object> getCharacterDataLoader() {
    return dataLoaderRegistry.getDataLoader("characters");
}
 

@Override
public <K, V> DataLoader<K, V> getDataLoader(String dataLoaderName) {
  throw new UnsupportedOperationException("Not yet implemented");//FIXME: implement
}
 

private Object getLinkPostedBy(DataFetchingEnvironment env) {
  Link link = env.getSource();
  DataLoader<String, User> user = env.getDataLoader("user");
  return user.load(link.getUserId());
}
 

private Object getLinkPostedBy(DataFetchingEnvironment env) {
  Link link = env.getSource();
  DataLoader<String, User> user = env.getDataLoader("user");
  return user.load(link.getUserId());
}
 

@Override
public <K, V> DataLoader<K, V> getDataLoader(String dataLoaderName) {
    return delegate.getDataLoader(dataLoaderName);
}
 

private static DataLoader<String, Object> newCharacterDataLoader() {
    return new DataLoader<>(characterBatchLoader);
}
 

public DataLoader<String, Object> getCharacterDataLoader() {
    return dataLoaderRegistry.getDataLoader("characters");
}
 

private void statsConfigExample() {

        DataLoaderOptions options = DataLoaderOptions.newOptions().setStatisticsCollector(() -> new ThreadLocalStatisticsCollector());
        DataLoader<String, User> userDataLoader = DataLoader.newDataLoader(userBatchLoader, options);
    }
 

private DataLoader<String, Object> newCharacterDataLoader() {
    return new DataLoader<>(characterBatchLoader);
}
 

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);

  });
}
 

private void disableCache() {
    DataLoader.newDataLoader(userBatchLoader, DataLoaderOptions.newOptions().setCachingEnabled(false));


    userDataLoader.load("A");
    userDataLoader.load("B");
    userDataLoader.load("A");

    userDataLoader.dispatch();

    // will result in keys to the batch loader with [ "A", "B", "A" ]
}