InputReferenceExtractor.java

package com.airbnb.airpal.core.execution;

import com.airbnb.airpal.presto.Table;
import com.facebook.presto.sql.tree.CreateTable;
import com.facebook.presto.sql.tree.CreateView;
import com.facebook.presto.sql.tree.DefaultTraversalVisitor;
import com.facebook.presto.sql.tree.DropTable;
import com.facebook.presto.sql.tree.DropView;
import com.facebook.presto.sql.tree.Join;
import com.facebook.presto.sql.tree.JoinOn;
import com.facebook.presto.sql.tree.Node;
import com.facebook.presto.sql.tree.QualifiedName;
import com.facebook.presto.sql.tree.RenameTable;
import com.facebook.presto.sql.tree.Use;
import com.facebook.presto.sql.tree.WithQuery;
import com.google.common.collect.Sets;
import lombok.EqualsAndHashCode;
import lombok.Value;

import java.util.HashSet;
import java.util.List;
import java.util.Set;

@Value
@EqualsAndHashCode(callSuper = false)
public class InputReferenceExtractor
        extends DefaultTraversalVisitor<InputReferenceExtractor.CatalogSchemaContext, InputReferenceExtractor.CatalogSchemaContext>
{
    private final Set<Table> references = new HashSet<>();
    private final Set<Table> aliases = new HashSet<>();

    public Set<Table> getReferences()
    {
        return Sets.difference(references, aliases);
    }

    private Table qualifiedNameToTable(QualifiedName name, CatalogSchemaContext context)
    {
        List<String> nameParts = name.getParts();

        String connectorId = context.getCatalog();
        String schema = context.getSchema();
        String table = null;

        if (nameParts.size() == 3) {
            connectorId = nameParts.get(0);
            schema = nameParts.get(1);
            table = nameParts.get(2);
        } else if (nameParts.size() == 2) {
            schema = nameParts.get(0);
            table = nameParts.get(1);
        } else if (nameParts.size() == 1) {
            table = nameParts.get(0);
        }

        return new Table(connectorId, schema, table);
    }

    @Override
    protected CatalogSchemaContext visitCreateView(CreateView node, CatalogSchemaContext context)
    {
        references.add(qualifiedNameToTable(node.getName(), context));
        visitQuery(node.getQuery(), context);

        return context;
    }

    @Override
    protected CatalogSchemaContext visitCreateTable(CreateTable node, CatalogSchemaContext context)
    {
        references.add(qualifiedNameToTable(node.getName(), context));
        visitCreateTable(node, context);

        return context;
    }

    @Override
    protected CatalogSchemaContext visitDropTable(DropTable node, CatalogSchemaContext context)
    {
        references.add(qualifiedNameToTable(node.getTableName(), context));
        return context;
    }

    @Override
    protected CatalogSchemaContext visitDropView(DropView node, CatalogSchemaContext context)
    {
        references.add(qualifiedNameToTable(node.getName(), context));
        return context;
    }

    @Override
    protected CatalogSchemaContext visitTable(com.facebook.presto.sql.tree.Table node, CatalogSchemaContext context)
    {
        references.add(qualifiedNameToTable(node.getName(), context));
        return context;
    }

    @Override
    protected CatalogSchemaContext visitRenameTable(RenameTable node, CatalogSchemaContext context)
    {
        references.add(qualifiedNameToTable(node.getSource(), context));
        return context;
    }

    @Override
    protected CatalogSchemaContext visitWithQuery(WithQuery node, CatalogSchemaContext context)
    {
        aliases.add(new Table(context.getCatalog(), context.getSchema(), node.getName()));
        return super.visitWithQuery(node, context);
    }

    @Override
    protected CatalogSchemaContext visitUse(Use node, CatalogSchemaContext context)
    {
        return new CatalogSchemaContext(node.getCatalog().orElse(context.getCatalog()), node.getSchema());
    }

    @Override
    protected CatalogSchemaContext visitNode(Node node, CatalogSchemaContext context)
    {
        return context;
    }

    @Override
    protected CatalogSchemaContext visitJoin(Join node, CatalogSchemaContext context)
    {
        process(node.getLeft(), context);
        process(node.getRight(), context);

        if (node.getCriteria().isPresent()) {
            if (node.getCriteria().get() instanceof JoinOn) {
                process(((JoinOn) node.getCriteria().get()).getExpression(), context);
            }
        }

        return context;
    }

    @Value
    public static class CatalogSchemaContext
    {
        private final String catalog;
        private final String schema;
    }
}