content: package com.airbnb.airpal.presto.metadata;
  
  import com.airbnb.airpal.core.BackgroundCacheLoader;
  import com.airbnb.airpal.core.execution.QueryClient;
  import com.airbnb.airpal.presto.QueryRunner;
  import com.airbnb.airpal.presto.Table;
  import com.airbnb.airpal.presto.Util;
  import com.airbnb.airpal.presto.hive.HivePartition;
  import com.facebook.presto.client.QueryResults;
  import com.facebook.presto.client.StatementClient;
  import com.google.common.base.Function;
  import com.google.common.base.Optional;
  import com.google.common.cache.CacheBuilder;
  import com.google.common.cache.LoadingCache;
  import com.google.common.collect.ImmutableList;
  import com.google.common.util.concurrent.ListeningExecutorService;
  import com.google.common.util.concurrent.MoreExecutors;
  import io.airlift.units.Duration;
  import lombok.extern.slf4j.Slf4j;
  
  import javax.annotation.Nullable;
  
  import java.util.List;
  import java.util.Objects;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.ExecutorService;
  
  import static com.google.common.base.Preconditions.checkNotNull;
  import static java.lang.String.format;
  
  @Slf4j
  public class PreviewTableCache
  {
      private final LoadingCache<PartitionedTableWithValue, List<List<Object>>> previewTableCache;
      private final QueryRunner.QueryRunnerFactory queryRunnerFactory;
  
      public PreviewTableCache(final QueryRunner.QueryRunnerFactory queryRunnerFactory,
                               final Duration previewCacheLifetime,
                               final ExecutorService executor,
                               final int previewLimit)
      {
          this.queryRunnerFactory = checkNotNull(queryRunnerFactory, "queryRunnerFactory session was null!");
  
          ListeningExecutorService listeningExecutor = MoreExecutors.listeningDecorator(executor);
  
          BackgroundCacheLoader<PartitionedTableWithValue, List<List<Object>>> tableLoader =
                  new BackgroundCacheLoader<PartitionedTableWithValue, List<List<Object>>>(listeningExecutor)
          {
                      @Override
                      public List<List<Object>> load(PartitionedTableWithValue key)
                              throws Exception
                      {
                          return queryRows(buildQueryWithLimit(key, previewLimit));
                      }
          };
  
          this.previewTableCache = CacheBuilder.newBuilder()
                                              .expireAfterWrite(Math.round(previewCacheLifetime.getValue()),
                                                                previewCacheLifetime.getUnit())
                                              .maximumSize(previewLimit)
                                              .build(tableLoader);
      }
  
      private static String buildQueryWithLimit(PartitionedTableWithValue tableWithValue, int limit)
      {
          Table table = tableWithValue.getTable();
          HivePartition partition = tableWithValue.getPartition().orNull();
          String partitionClause = "";
  
          if (partition != null) {
              String value = tableWithValue.getValue();
              String partitionValue = (Objects.equals(partition.getType(), "varchar")) ?
                      "'" + value + "'" :
                      String.valueOf(value);
  
              partitionClause = format("WHERE %s = %s", partition.getName(), partitionValue);
          }
  
          return format("SELECT * FROM %s %s LIMIT %d",
                  Util.fqn(table.getConnectorId(), table.getSchema(), table.getTable()),
                  partitionClause,
                  limit);
      }
  
      private List<List<Object>> queryRows(String query)
      {
          final ImmutableList.Builder<List<Object>> cache = ImmutableList.builder();
          QueryRunner queryRunner = queryRunnerFactory.create();
          QueryClient queryClient = new QueryClient(queryRunner, io.dropwizard.util.Duration.seconds(60), query);
  
          try {
              queryClient.executeWith(new Function<StatementClient, Void>() {
                  @Nullable
                  @Override
                  public Void apply(StatementClient client)
                  {
                      QueryResults results = client.current();
                      if (results.getData() != null) {
                          cache.addAll(results.getData());
                      }
  
                      return null;
                  }
              });
          }
          catch (QueryClient.QueryTimeOutException e) {
              log.error("Caught timeout loading columns", e);
          }
  
          return cache.build();
      }
  
      public List<List<Object>> getPreview(final String connectorId,
              final String schema,
              final String table,
              final Optional<HivePartition> partition,
              final String partitionValue)
              throws ExecutionException
      {
          return previewTableCache.get(
                  new PartitionedTableWithValue(
                          new Table(connectorId,
                                  schema,
                                  table),
                          partition,
                          partitionValue
                  ));
      }
  }
  