All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.facebook.presto.cli.TableNameCompleter Maven / Gradle / Ivy

Go to download

There is a newer version: 0.289
Show newest version
/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.facebook.presto.cli;

import com.facebook.presto.client.QueryData;
import com.facebook.presto.client.StatementClient;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.collect.ImmutableList;
import jline.console.completer.Completer;

import java.io.Closeable;
import java.util.List;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;

import static com.facebook.airlift.concurrent.Threads.daemonThreadsNamed;
import static com.google.common.cache.CacheLoader.asyncReloading;
import static java.lang.String.format;
import static java.util.Objects.requireNonNull;
import static java.util.concurrent.Executors.newCachedThreadPool;

public class TableNameCompleter
        implements Completer, Closeable
{
    private static final long RELOAD_TIME_MINUTES = 2;

    private final ExecutorService executor = newCachedThreadPool(daemonThreadsNamed("completer-%s"));
    private final QueryRunner queryRunner;
    private final LoadingCache> tableCache;
    private final LoadingCache> functionCache;

    public TableNameCompleter(QueryRunner queryRunner)
    {
        this.queryRunner = requireNonNull(queryRunner, "queryRunner session was null!");

        tableCache = CacheBuilder.newBuilder()
                .refreshAfterWrite(RELOAD_TIME_MINUTES, TimeUnit.MINUTES)
                .build(asyncReloading(CacheLoader.from(this::listTables), executor));

        functionCache = CacheBuilder.newBuilder()
                .build(asyncReloading(CacheLoader.from(this::listFunctions), executor));
    }

    private List listTables(String schemaName)
    {
        return queryMetadata(format("SELECT table_name FROM information_schema.tables WHERE table_schema = '%s'", schemaName));
    }

    @SuppressWarnings("unused")
    private List listFunctions(String schemaName)
    {
        return queryMetadata("SHOW FUNCTIONS");
    }

    private List queryMetadata(String query)
    {
        ImmutableList.Builder cache = ImmutableList.builder();
        try (StatementClient client = queryRunner.startInternalQuery(query)) {
            while (client.isRunning() && !Thread.currentThread().isInterrupted()) {
                QueryData results = client.currentData();
                if (results.getData() != null) {
                    for (List row : results.getData()) {
                        cache.add((String) row.get(0));
                    }
                }
                client.advance();
            }
        }
        return cache.build();
    }

    public void populateCache()
    {
        String schemaName = queryRunner.getSession().getSchema();
        if (schemaName != null) {
            executor.execute(() -> {
                functionCache.refresh(schemaName);
                tableCache.refresh(schemaName);
            });
        }
    }

    @Override
    public int complete(String buffer, int cursor, List candidates)
    {
        if (cursor <= 0) {
            return cursor;
        }
        int blankPos = findLastBlank(buffer.substring(0, cursor));
        String prefix = buffer.substring(blankPos + 1, cursor);
        String schemaName = queryRunner.getSession().getSchema();

        if (schemaName != null) {
            List functionNames = functionCache.getIfPresent(schemaName);
            List tableNames = tableCache.getIfPresent(schemaName);

            SortedSet sortedCandidates = new TreeSet<>();
            if (functionNames != null) {
                sortedCandidates.addAll(filterResults(functionNames, prefix));
            }
            if (tableNames != null) {
                sortedCandidates.addAll(filterResults(tableNames, prefix));
            }

            candidates.addAll(sortedCandidates);
        }

        return blankPos + 1;
    }

    private static int findLastBlank(String buffer)
    {
        for (int i = buffer.length() - 1; i >= 0; i--) {
            if (Character.isWhitespace(buffer.charAt(i))) {
                return i;
            }
        }
        return -1;
    }

    private static List filterResults(List values, String prefix)
    {
        ImmutableList.Builder builder = ImmutableList.builder();
        for (String value : values) {
            if (value.startsWith(prefix)) {
                builder.add(value);
            }
        }
        return builder.build();
    }

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

    

    
            
    
            







    
© 2015 - 2024 Weber Informatics LLC | Privacy Policy