package com.airbnb.airpal.core.execution;

import com.airbnb.airpal.core.AuthorizationUtil;
import com.airbnb.airpal.presto.Table;
import com.facebook.presto.sql.parser.SqlParser;
import com.facebook.presto.sql.tree.Statement;
import com.google.common.base.Splitter;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import org.apache.shiro.subject.Subject;

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

import static com.airbnb.airpal.core.AuthorizationUtil.AuthorizedTablesPredicate;
import static com.airbnb.airpal.core.execution.InputReferenceExtractor.CatalogSchemaContext;
import static com.google.common.base.Preconditions.checkNotNull;

public class QueryExecutionAuthorizer
{
    private static final SqlParser SQL_PARSER = new SqlParser();

    private final Subject subject;
    private final Set<Table> seen = new HashSet<>();
    private final String defaultConnector;
    private final String defaultSchema;
    private boolean lastResult = false;

    public QueryExecutionAuthorizer(Subject subject, String defaultConnector, String defaultSchema)
    {
        this.subject = checkNotNull(subject);
        this.defaultConnector = checkNotNull(defaultConnector);
        this.defaultSchema = checkNotNull(defaultSchema);
    }

    public boolean isAuthorizedWrite(String connectorId, String schema, String table)
    {
        return AuthorizationUtil.isAuthorizedWrite(subject, connectorId, schema, table);
    }

    public boolean isAuthorizedRead(Set<Table> tables)
    {
        if (seen.containsAll(tables)) {
            return lastResult;
        }

        seen.addAll(tables);
        return (lastResult = Iterables.all(tables, new AuthorizedTablesPredicate(subject)));
    }

    private static Splitter STATEMENT_SPLITTER = Splitter.on(";").omitEmptyStrings();

    public static Set<Table> tablesUsedByQuery(String query, String defaultConnector, String defaultSchema)
    {
        List<String> statements = STATEMENT_SPLITTER.splitToList(query);
        ImmutableSet.Builder<Table> tables = ImmutableSet.builder();
        CatalogSchemaContext context = new CatalogSchemaContext(defaultConnector, defaultSchema);

        for (String strStatement : statements) {
            InputReferenceExtractor extractor = new InputReferenceExtractor();
            Statement statement = SQL_PARSER.createStatement(strStatement);
            context = statement.accept(extractor, context);

            tables.addAll(extractor.getReferences());
        }

        return tables.build();
    }

    public Set<Table> tablesUsedByQuery(String query)
    {
        return tablesUsedByQuery(query, defaultConnector, defaultSchema);
    }
}

@@ -397,34 +397,8 @@
 t;%0A%0A
-import java.util.HashSet;%0A
 impo
@@ -416,16 +416,16 @@
 l.List;%0A
+
 import j
@@ -825,61 +825,8 @@
 ct;%0A
-    private final Set%3CTable%3E seen = new HashSet%3C%3E();%0A
     
@@ -907,48 +907,8 @@
 ema;
-%0A    private boolean lastResult = false;
 %0A%0A  
@@ -1447,139 +1447,14 @@
     
-if (seen.containsAll(tables)) %7B%0A            return lastResult;%0A        %7D%0A%0A        seen.addAll(tables);%0A        return (lastResult =
+return
  Ite
@@ -1511,17 +1511,16 @@
 ubject))
-)
 ;%0A    %7D%0A

Remove old caching logic