com.facebook.presto.verifier.framework.VerifierUtil Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of presto-verifier Show documentation
There is a newer version: 0.290
/*
 * 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.verifier.framework;

import com.facebook.presto.common.type.Type;
import com.facebook.presto.common.type.TypeManager;
import com.facebook.presto.common.type.TypeSignature;
import com.facebook.presto.sql.parser.ParsingOptions;
import com.facebook.presto.sql.tree.Identifier;
import com.facebook.presto.verifier.prestoaction.QueryActionStats;

import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.IntStream;

import static com.facebook.presto.sql.parser.ParsingOptions.DecimalLiteralTreatment.AS_DOUBLE;
import static com.google.common.base.Functions.identity;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableMap.toImmutableMap;

public class VerifierUtil
{
    private VerifierUtil() {}

    public static final ParsingOptions PARSING_OPTIONS = ParsingOptions.builder().setDecimalLiteralTreatment(AS_DOUBLE).build();

    public static Identifier delimitedIdentifier(String name)
    {
        return new Identifier(name, true);
    }

    public static void runAndConsume(Callable callable, Consumer queryStatsConsumer)
    {
        runAndConsume(callable, queryStatsConsumer, e -> {});
    }

    public static void runAndConsume(Callable callable, Consumer queryStatsConsumer, Consumer queryExceptionConsumer)
    {
        callAndConsume(callable, identity(), queryStatsConsumer, queryExceptionConsumer);
    }

    public static  QueryResult callAndConsume(Callable> callable, Consumer queryStatsConsumer)
    {
        return callAndConsume(callable, queryStatsConsumer, e -> {});
    }

    public static  QueryResult callAndConsume(Callable> callable, Consumer queryStatsConsumer, Consumer queryExceptionConsumer)
    {
        return callAndConsume(callable, QueryResult::getQueryActionStats, queryStatsConsumer, queryExceptionConsumer);
    }

    private static  V callAndConsume(
            Callable callable,
            Function queryStatsTransformer,
            Consumer queryStatsConsumer,
            Consumer queryExceptionConsumer)
    {
        try {
            V result = callable.call();
            queryStatsConsumer.accept(queryStatsTransformer.apply(result));
            return result;
        }
        catch (PrestoQueryException e) {
            queryStatsConsumer.accept(e.getQueryActionStats());
            queryExceptionConsumer.accept(e);
            throw e;
        }
    }

    public static List getColumnNames(ResultSetMetaData metadata)
    {
        return callUnchecked(() ->
                IntStream.rangeClosed(1, metadata.getColumnCount())
                        .mapToObj(i -> callUnchecked(() -> metadata.getColumnName(i)))
                        .collect(toImmutableList()));
    }

    public static Map getColumnIndices(ResultSetMetaData metadata)
    {
        return callUnchecked(() ->
                IntStream.rangeClosed(1, metadata.getColumnCount())
                        .boxed()
                        .collect(toImmutableMap(i -> callUnchecked(() -> metadata.getColumnName(i)), i -> i - 1)));
    }

    public static List getColumnTypes(TypeManager typeManager, ResultSetMetaData metadata)
    {
        return callUnchecked(() ->
                IntStream.rangeClosed(1, metadata.getColumnCount())
                        .mapToObj(i -> callUnchecked(() -> metadata.getColumnTypeName(i)))
                        .map(TypeSignature::parseTypeSignature)
                        .map(typeManager::getType)
                        .collect(toImmutableList()));
    }

    private static  V callUnchecked(SqlExceptionCallable callable)
    {
        try {
            return callable.call();
        }
        catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }

    public interface Callable
    {
        V call();
    }

    public interface SqlExceptionCallable
    {
        V call()
                throws SQLException;
    }
}