QueryExecutionAuthorizerTest.java

package com.airbnb.airpal.core.execution;

import com.airbnb.airpal.presto.Table;
import com.google.common.collect.ImmutableSet;
import org.junit.Test;

import java.util.Set;

import static com.airbnb.airpal.core.execution.QueryExecutionAuthorizer.tablesUsedByQuery;
import static org.junit.Assert.assertEquals;

public class QueryExecutionAuthorizerTest
{
    static String defaultConnector = "hive";
    static String defaultSchema = "default";

    static String TEST_CREATE_TABLE = "CREATE TABLE the_gibson.users_pii AS SELECT * FROM users;";
    @Test
    public void testTableReferencesCreateTable()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_CREATE_TABLE, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table(defaultConnector, defaultSchema, "users"),
                new Table(defaultConnector, "the_gibson", "users_pii")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    static String TEST_CREATE_VIEW = "CREATE VIEW the_gibson.users_pii AS SELECT pii_col FROM users;";
    @Test
    public void testTableReferencesCreateView()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_CREATE_VIEW, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table(defaultConnector, defaultSchema, "users"),
                new Table(defaultConnector, "the_gibson", "users_pii")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    static String TEST_REPLACE_VIEW = "CREATE OR REPLACE VIEW the_gibson.users AS SELECT pii_col FROM users;";
    @Test
    public void testTableReferencesReplaceView()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_REPLACE_VIEW, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table(defaultConnector, defaultSchema, "users"),
                new Table(defaultConnector, "the_gibson", "users")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    static String TEST_ALTER_TABLE = "ALTER TABLE default.users RENAME TO the_gibson.users_pii; SELECT pii_col FROM the_gibson.users_pii;";
    @Test
    public void testTableReferencesRenameTable()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_ALTER_TABLE, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table(defaultConnector, "default", "users"),
                new Table(defaultConnector, "the_gibson", "users_pii")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    static String TEST_ALTER_TABLE2 = "USE SCHEMA the_gibson; ALTER TABLE default.users RENAME TO users_pii; SELECT pii_col FROM users_pii;";
    @Test
    public void testTableReferencesRenameTable2()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_ALTER_TABLE2, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table(defaultConnector, "default", "users"),
                new Table(defaultConnector, "the_gibson", "users_pii")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    static String TEST_ALTER_TABLE3 = "USE CATALOG cassandra; USE SCHEMA the_gibson; ALTER TABLE hive.default.users RENAME TO users_pii; SELECT pii_col FROM users_pii;";
    @Test
    public void testTableReferencesRenameTable3()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_ALTER_TABLE3, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table(defaultConnector, "default", "users"),
                new Table("cassandra", "the_gibson", "users_pii")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    static String TEST_ACCESS_PII_UNION_VIEW = "SELECT str_col FROM the_gibson.users UNION ALL SELECT pii_str_col FROM users;";
    @Test
    public void testTableReferencesSelectUnion()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_ACCESS_PII_UNION_VIEW, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table(defaultConnector, defaultSchema, "users"),
                new Table(defaultConnector, "the_gibson", "users")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    static String TEST_ACCESS_WITH = "WITH a AS (SELECT * FROM users) SELECT * FROM the_gibson.users UNION ALL SELECT * FROM a;";
    @Test
    public void testTableReferencesSelectWith()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_ACCESS_WITH, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table(defaultConnector, defaultSchema, "users"),
                new Table(defaultConnector, "the_gibson", "users")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    static String TEST_BASIC_SELECT_COUNT = "SELECT COUNT(*) FROM users;";
    @Test
    public void testTableReferencesSelectCount()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_BASIC_SELECT_COUNT, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table(defaultConnector, defaultSchema, "users")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    static String TEST_SELECT_ALL = "SELECT * FROM users;";
    @Test
    public void testTableReferencesSelectStar()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_SELECT_ALL, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table(defaultConnector, defaultSchema, "users")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    static String TEST_SELECT_SUBQUERY = "SELECT * FROM (SELECT pii_str_col FROM users) UNION ALL SELECT * FROM the_gibson.users;";
    @Test
    public void testTableReferencesSelectStarSubquery()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_SELECT_SUBQUERY, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table(defaultConnector, defaultSchema, "users"),
                new Table(defaultConnector, "the_gibson", "users")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    static String TEST_DROP_VIEW = "DROP VIEW my_view";
    @Test
    public void testTableReferencesDropView()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_DROP_VIEW, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table(defaultConnector, defaultSchema, "my_view")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    static String TEST_DROP_TABLE = "DROP TABLE my_table";
    @Test
    public void testTableReferencesDropTable()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_DROP_TABLE, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table(defaultConnector, defaultSchema, "my_table")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    static String TEST_SELECT_ALIAS = "SELECT pii FROM users u JOIN users_pii p ON u.id = p.id;";
    @Test
    public void testTableReferencesJoinAlias()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_SELECT_ALIAS, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table(defaultConnector, defaultSchema, "users"),
                new Table(defaultConnector, defaultSchema, "users_pii")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    static String TEST_SELECT_CONNECTOR = "SELECT pii FROM cassandra.pii.users u JOIN users_pii p ON u.id = p.id;";
    @Test
    public void testTableReferencesSelectConnector()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_SELECT_CONNECTOR, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table("cassandra", "pii", "users"),
                new Table(defaultConnector, defaultSchema, "users_pii")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    static String TEST_CROSS_JOIN_SUBSELECT = "SELECT * " +
            "FROM cassandra.pii.users u " +
            "CROSS JOIN (SELECT * FROM hive.pii.users) u2 " +
            "WHERE NOT u2.id IS NULL;";
    @Test
    public void testTableReferencesCrossJoinSubselect()
            throws Exception
    {
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_CROSS_JOIN_SUBSELECT, defaultConnector, defaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table("cassandra", "pii", "users"),
                new Table("hive", "pii", "users")
        );

        assertEquals(tablesExpected, tablesUsed);
    }

    @Test
    public void testAlternateDefaultConnectorSchemaReferences()
            throws Exception
    {
        String alternateDefaultConnector = "cassandra";
        String alternateDefaultSchema = "default2";
        Set<Table> tablesUsed = tablesUsedByQuery(TEST_SELECT_ALIAS, alternateDefaultConnector, alternateDefaultSchema);
        Set<Table> tablesExpected = ImmutableSet.of(
                new Table(alternateDefaultConnector, alternateDefaultSchema, "users"),
                new Table(alternateDefaultConnector, alternateDefaultSchema, "users_pii")
        );

        assertEquals(tablesExpected, tablesUsed);
    }
}