com.amazonaws.athena.connectors.jdbc.manager.FederationExpressionParser Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of athena-jdbc Show documentation
The newest version!
/*-
 * #%L
 * Amazon Athena Query Federation SDK
 * %%
 * Copyright (C) 2019 - 2022 Amazon Web Services
 * %%
 * 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.
 * #L%
 */
package com.amazonaws.athena.connectors.jdbc.manager;

import com.amazonaws.athena.connector.lambda.data.Block;
import com.amazonaws.athena.connector.lambda.domain.predicate.Constraints;
import com.amazonaws.athena.connector.lambda.domain.predicate.expression.ConstantExpression;
import com.amazonaws.athena.connector.lambda.domain.predicate.expression.FederationExpression;
import com.amazonaws.athena.connector.lambda.domain.predicate.expression.FunctionCallExpression;
import com.amazonaws.athena.connector.lambda.domain.predicate.expression.VariableExpression;
import com.amazonaws.athena.connector.lambda.domain.predicate.functions.FunctionName;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Joiner;
import com.google.common.collect.ImmutableList;
import org.apache.arrow.vector.complex.reader.FieldReader;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

import static com.amazonaws.athena.connector.lambda.domain.predicate.expression.ConstantExpression.DEFAULT_CONSTANT_EXPRESSION_BLOCK_NAME;

public abstract class FederationExpressionParser 
{
    static final Logger LOGGER = LoggerFactory.getLogger(FederationExpressionParser.class);

    private static final String quoteCharacter = "'"; // AFAIK this is going to be valid for all sources when quoting a constant.

    /**
     * Each datasource has different syntax for various operations, quotes, etc. This is the only method a subclass to implement, and otherwise will just
     * invoke parseComplexExpressions.
     */
    public abstract String mapFunctionToDataSourceSyntax(FunctionName functionName, ArrowType type, List arguments);

    public List parseComplexExpressions(List columns, Constraints constraints, List accumulator)
    {
        if (constraints.getExpression() == null || constraints.getExpression().isEmpty()) {
            return ImmutableList.of();
        }

        List federationExpressions = constraints.getExpression();
        return federationExpressions.stream()
                    .map(federationExpression -> parseFunctionCallExpression((FunctionCallExpression) federationExpression, accumulator))
                    .collect(Collectors.toList());
    }

    /**
     * This is a recursive function, as function calls can have arguments which, themselves, are function calls.
     * @param functionCallExpression
     * @return
     */
    public String parseFunctionCallExpression(FunctionCallExpression functionCallExpression, List accumulator)
    {
        FunctionName functionName = functionCallExpression.getFunctionName();
        List functionArguments = functionCallExpression.getArguments();

        List arguments = functionArguments.stream()
            .map(argument -> {
                // base cases
                if (argument instanceof ConstantExpression) {
                    return parseConstantExpression((ConstantExpression) argument, accumulator);
                }
                else if (argument instanceof VariableExpression) {
                    return parseVariableExpression((VariableExpression) argument);
                }
                // recursive case
                else if (argument instanceof FunctionCallExpression) {
                    return parseFunctionCallExpression((FunctionCallExpression) argument, accumulator);
                }
                throw new RuntimeException("Should not reach this case - a new subclass was introduced and is not handled.");
            }).collect(Collectors.toList());

        return mapFunctionToDataSourceSyntax(functionName, functionCallExpression.getType(), arguments);
    }
    
    // basing this off of the toString() impl in Block.java
    @VisibleForTesting
    public String parseConstantExpression(ConstantExpression constantExpression, List accumulator)
    {
        Block values = constantExpression.getValues();
        FieldReader fieldReader = values.getFieldReader(DEFAULT_CONSTANT_EXPRESSION_BLOCK_NAME);
        
        for (int i = 0; i < values.getRowCount(); i++) {
            fieldReader.setPosition(i);
            accumulator.add(new TypeAndValue(constantExpression.getType(), fieldReader.readObject()));
        }

        return Joiner.on(",").join(Collections.nCopies(values.getRowCount(), "?"));
    }

    /**
     * Various connectors have different standards for wrapping column names in some specific quote character.
     * This default implementation just returns the column name for connectors that don't have this specification.
     */
    @VisibleForTesting
    public String parseVariableExpression(VariableExpression variableExpression) 
    {
        return variableExpression.getColumnName();
    }
}