ColumnCache.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.Util;
import com.airbnb.airpal.presto.hive.HiveColumn;
import com.airbnb.airpal.presto.hive.HivePartition;
import com.facebook.presto.client.ClientTypeSignature;
import com.facebook.presto.client.Column;
import com.facebook.presto.client.QueryResults;
import com.facebook.presto.client.StatementClient;
import com.facebook.presto.spi.type.TypeSignature;
import com.google.common.base.Function;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.LoadingCache;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
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.Map;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;

import static com.airbnb.airpal.presto.QueryRunner.QueryRunnerFactory;
import static com.google.common.base.Preconditions.checkNotNull;
import static java.lang.String.format;

@Slf4j
public class ColumnCache
{
    private final LoadingCache<String, List<HiveColumn>> schemaTableCache;
    private final LoadingCache<String, List<HivePartition>> partitionCache;
    private final QueryRunnerFactory queryRunnerFactory;

    public ColumnCache(final QueryRunnerFactory queryRunnerFactory,
                       final Duration columnCacheLifetime,
                       final Duration partitionCacheLifetime,
                       final ExecutorService executor)
    {
        this.queryRunnerFactory = checkNotNull(queryRunnerFactory, "queryRunnerFactory session was null!");
        checkNotNull(executor, "executor was null!");

        ListeningExecutorService listeningExecutor = MoreExecutors.listeningDecorator(executor);

        BackgroundCacheLoader<String, List<HiveColumn>> columnLoader = new BackgroundCacheLoader<String,
                List<HiveColumn>>(listeningExecutor)
        {
            @Override
            public List<HiveColumn> load(String fqTableName)
            {
                return queryColumns(format("SHOW COLUMNS FROM %s", fqTableName));
            }
        };

        BackgroundCacheLoader<String, List<HivePartition>> partitionLoader = new BackgroundCacheLoader<String,
                List<HivePartition>>(listeningExecutor)
        {
            @Override
            public List<HivePartition> load(String fqTableName) throws Exception
            {
                return queryPartitions(format("SHOW PARTITIONS FROM %s", fqTableName));
            }
        };

        this.schemaTableCache = CacheBuilder.newBuilder()
                                            .expireAfterWrite(Math.round(columnCacheLifetime.getValue()),
                                                              columnCacheLifetime.getUnit())
                                            .build(columnLoader);
        this.partitionCache = CacheBuilder.newBuilder()
                                          .expireAfterWrite(Math.round(partitionCacheLifetime.getValue()),
                                                            partitionCacheLifetime.getUnit())
                                          .build(partitionLoader);
    }

    private List<HivePartition> queryPartitions(String query)
    {
        final ImmutableList.Builder<HivePartition> cache = ImmutableList.builder();
        final Map<Column, List<Object>> objects = Maps.newHashMap();
        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 && results.getColumns() != null) {
                        final List<Column> columns = results.getColumns();

                        for (Column column : columns) {
                            objects.put(column, Lists.newArrayList());
                        }

                        for (List<Object> row : results.getData()) {
                            for (int i = 0; i < row.size(); i++) {
                                Column column = columns.get(i);
                                objects.get(column).add(row.get(i));
                            }
                        }
                    }

                    return null;
                }
            });
        }
        catch (QueryClient.QueryTimeOutException e) {
            log.error("Caught timeout loading columns", e);
        }

        for (Map.Entry<Column, List<Object>> entry : objects.entrySet()) {
            cache.add(HivePartition.fromColumn(entry.getKey(), entry.getValue()));
        }

        return cache.build();
    }

    private List<HiveColumn> queryColumns(String query)
    {
        final ImmutableList.Builder<HiveColumn> 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) {
                        for (List<Object> row : results.getData()) {
                            Column column = new Column((String) row.get(0), (String) row.get(1), new ClientTypeSignature(TypeSignature.parseTypeSignature((String)row.get(1))));
                            boolean isNullable = (Boolean) row.get(2);
                            boolean isPartition = (Boolean) row.get(3);

                            cache.add(HiveColumn.fromColumn(column, isNullable, isPartition));
                        }
                    }

                    return null;
                }
            });
        }
        catch (QueryClient.QueryTimeOutException e) {
            log.error("Caught timeout loading columns", e);
        }

        return cache.build();
    }

    public List<HiveColumn> getColumns(String databaseName, String tableName) throws ExecutionException
    {
        return schemaTableCache.get(Util.fqn(databaseName, tableName));
    }

    public List<HivePartition> getPartitions(String databaseName, String tableName) throws ExecutionException
    {
        return partitionCache.get(Util.fqn(databaseName, tableName));
    }

}