SchemaCache.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.facebook.presto.client.QueryResults;
import com.facebook.presto.client.StatementClient;
import com.google.common.base.Function;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.LoadingCache;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import lombok.extern.slf4j.Slf4j;

import javax.annotation.Nullable;

import java.io.Closeable;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;

import static com.google.common.base.Preconditions.checkNotNull;
import static java.lang.String.format;

@Slf4j
public class SchemaCache
        implements Closeable
{
    private static final long RELOAD_TIME_MINUTES = 2;
    private static final Set<String> EXCLUDED_SCHEMAS = Sets.newHashSet("sys", "information_schema");

    private final ExecutorService executor;
    private final QueryRunner.QueryRunnerFactory queryRunnerFactory;
    private final LoadingCache<String, Map<String, List<String>>> schemaTableCache;

    public SchemaCache(final QueryRunner.QueryRunnerFactory queryRunnerFactory,
            final ExecutorService executor)
    {
        this.queryRunnerFactory = checkNotNull(queryRunnerFactory, "queryRunnerFactory session was null!");
        this.executor = checkNotNull(executor, "executor was null!");

        ListeningExecutorService listeningExecutor = MoreExecutors.listeningDecorator(executor);
        BackgroundCacheLoader<String, Map<String, List<String>>> loader =
                new BackgroundCacheLoader<String, Map<String, List<String>>>(listeningExecutor)
                {
                    @Override
                    public Map<String, List<String>> load(String catalogName)
                    {
                        return queryMetadata(format(
                                "SELECT table_catalog, table_schema, table_name " +
                                        "FROM information_schema.tables " +
                                        "WHERE table_catalog = '%s'",
                                catalogName));
                    }
                };

        schemaTableCache = CacheBuilder.newBuilder()
                .refreshAfterWrite(RELOAD_TIME_MINUTES, TimeUnit.MINUTES)
                .build(loader);
    }

    private Map<String, List<String>> queryMetadata(String query)
    {
        final Map<String, List<String>> cache = 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) {
                        for (List<Object> row : results.getData()) {
                            String schema = (String) row.get(1);
                            String table = (String) row.get(2);

                            if (EXCLUDED_SCHEMAS.contains(schema)) {
                                continue;
                            }

                            List<String> tables = cache.get(schema);

                            if (tables == null) {
                                tables = Lists.newArrayList();
                                cache.put(schema, tables);
                            }

                            tables.add(table);
                        }
                    }

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

        return ImmutableMap.copyOf(cache);
    }

    public void populateCache(final String catalog)
    {
        checkNotNull(catalog, "schemaName is null");
        executor.execute(new Runnable()
        {
            @Override
            public void run()
            {
                schemaTableCache.refresh(catalog);
            }
        });
    }

    public Set<String> getCatalogs()
    {
        return schemaTableCache.asMap().keySet();
    }

    public Map<String, List<String>> getSchemaMap(final String catalog)
    {
        try {
            return schemaTableCache.get(catalog);
        }
        catch (ExecutionException e) {
            e.printStackTrace();
            return Maps.newHashMap();
        }
    }

    @Override
    public void close()
    {
        executor.shutdownNow();
    }

    public static ThreadFactory daemonThreadsNamed(String nameFormat)
    {
        return new ThreadFactoryBuilder().setNameFormat(nameFormat).setDaemon(true).build();
    }
}