PreviewTableCache.java

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