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