QueryCache.java

package com.linkedin.thirdeye.client.cache;

import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

import com.linkedin.thirdeye.client.ThirdEyeClient;
import com.linkedin.thirdeye.client.ThirdEyeRequest;
import com.linkedin.thirdeye.client.ThirdEyeResponse;

public class QueryCache {
  private final ExecutorService executorService;
  private final ThirdEyeClient client;

  public QueryCache(ThirdEyeClient clientMap, ExecutorService executorService) {
    this.executorService = executorService;
    this.client = clientMap;
  }

  public ThirdEyeClient getClient() {
    return client;
  }

  public ThirdEyeResponse getQueryResult(ThirdEyeRequest request) throws Exception {
    ThirdEyeResponse response = client.execute(request);
    return response;
  }

  public Future<ThirdEyeResponse> getQueryResultAsync(final ThirdEyeRequest request)
      throws Exception {
    return executorService.submit(new Callable<ThirdEyeResponse>() {
      @Override
      public ThirdEyeResponse call() throws Exception {
        return getQueryResult(request);
      }
    });
  }

  public Map<ThirdEyeRequest, Future<ThirdEyeResponse>> getQueryResultsAsync(
      final List<ThirdEyeRequest> requests) throws Exception {
    Map<ThirdEyeRequest, Future<ThirdEyeResponse>> responseFuturesMap = new LinkedHashMap<>();
    for (final ThirdEyeRequest request : requests) {
      Future<ThirdEyeResponse> responseFuture =
          executorService.submit(new Callable<ThirdEyeResponse>() {
            @Override
            public ThirdEyeResponse call() throws Exception {
              return getQueryResult(request);
            }
          });
      responseFuturesMap.put(request, responseFuture);
    }
    return responseFuturesMap;
  }

  public Map<ThirdEyeRequest, ThirdEyeResponse> getQueryResultsAsyncAndWait(
      final List<ThirdEyeRequest> requests) throws Exception {
    return getQueryResultsAsyncAndWait(requests, 60);
  }

  public Map<ThirdEyeRequest, ThirdEyeResponse> getQueryResultsAsyncAndWait(
      final List<ThirdEyeRequest> requests, int timeoutSeconds) throws Exception {
    Map<ThirdEyeRequest, ThirdEyeResponse> responseMap = new LinkedHashMap<>();
    for (final ThirdEyeRequest request : requests) {
      Future<ThirdEyeResponse> responseFuture =
          executorService.submit(new Callable<ThirdEyeResponse>() {
            @Override
            public ThirdEyeResponse call() throws Exception {
              return getQueryResult(request);
            }
          });
      responseMap.put(request, responseFuture.get(timeoutSeconds, TimeUnit.SECONDS));
    }
    return responseMap;
  }

  public void clear() throws Exception {
    client.clear();
  }

}