org.apache.lens.driver.hive.TestHiveDriver.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.lens.driver.hive.TestHiveDriver.java

Source

/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.lens.driver.hive;

import static org.testng.Assert.*;

import java.io.*;
import java.text.ParseException;
import java.util.*;

import org.apache.lens.api.LensConf;
import org.apache.lens.api.Priority;
import org.apache.lens.api.query.QueryHandle;
import org.apache.lens.cube.metadata.FactPartition;
import org.apache.lens.cube.metadata.UpdatePeriod;
import org.apache.lens.server.api.LensConfConstants;
import org.apache.lens.server.api.driver.*;
import org.apache.lens.server.api.driver.DriverQueryStatus.DriverQueryState;
import org.apache.lens.server.api.driver.hooks.DriverQueryHook;
import org.apache.lens.server.api.error.LensException;
import org.apache.lens.server.api.query.ExplainQueryContext;
import org.apache.lens.server.api.query.PreparedQueryContext;
import org.apache.lens.server.api.query.QueryContext;
import org.apache.lens.server.api.query.cost.QueryCost;
import org.apache.lens.server.api.query.priority.CostRangePriorityDecider;
import org.apache.lens.server.api.query.priority.CostToPriorityRangeConf;
import org.apache.lens.server.api.user.MockDriverQueryHook;
import org.apache.lens.server.api.util.LensUtil;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.*;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.hive.conf.HiveConf;
import org.apache.hadoop.hive.metastore.api.Database;
import org.apache.hadoop.hive.ql.HiveDriverRunHook;
import org.apache.hadoop.hive.ql.HiveDriverRunHookContext;
import org.apache.hadoop.hive.ql.metadata.Hive;
import org.apache.hadoop.hive.ql.session.SessionState;
import org.apache.hive.service.cli.ColumnDescriptor;

import org.testng.annotations.*;

import com.beust.jcommander.internal.Maps;
import com.google.common.collect.Lists;

/**
 * The Class TestHiveDriver.
 */
public class TestHiveDriver {

    /** The Constant TEST_DATA_FILE. */
    public static final String TEST_DATA_FILE = "testdata/testdata1.data";

    /** The test output dir. */
    private final String testOutputDir = "target/" + this.getClass().getSimpleName() + "/test-output";

    /** The conf. */
    protected Configuration driverConf = new Configuration();
    protected HiveConf hiveConf = new HiveConf();
    protected Configuration queryConf = new Configuration();

    /** The driver. */
    protected HiveDriver driver;

    /** Driver list * */
    protected Collection<LensDriver> drivers;

    /** The data base. */
    String dataBase = this.getClass().getSimpleName().toLowerCase();

    protected String sessionid;
    protected SessionState ss;
    private CostRangePriorityDecider alwaysNormalPriorityDecider = new CostRangePriorityDecider(
            new CostToPriorityRangeConf(""));

    /**
     * Before test.
     *
     * @throws Exception the exception
     */
    @BeforeTest
    public void beforeTest() throws Exception {
        // Check if hadoop property set
        System.out.println("###HADOOP_PATH " + System.getProperty("hadoop.bin.path"));
        assertNotNull(System.getProperty("hadoop.bin.path"));
        createDriver();
        ss = new SessionState(hiveConf, "testuser");
        SessionState.start(ss);
        Hive client = Hive.get(hiveConf);
        Database database = new Database();
        database.setName(dataBase);
        client.createDatabase(database, true);
        SessionState.get().setCurrentDatabase(dataBase);
        sessionid = SessionState.get().getSessionId();
        driverConf.setBoolean(LensConfConstants.QUERY_ADD_INSERT_OVEWRITE, false);
        QueryContext context = createContext("USE " + dataBase, this.queryConf);
        driver.execute(context);
        driverConf.setBoolean(LensConfConstants.QUERY_ADD_INSERT_OVEWRITE, true);
        driverConf.setBoolean(LensConfConstants.QUERY_PERSISTENT_RESULT_INDRIVER, true);
    }

    protected void createDriver() throws LensException {
        driverConf.addResource("drivers/hive/hive1/hivedriver-site.xml");
        driverConf.setClass(HiveDriver.HIVE_CONNECTION_CLASS, EmbeddedThriftConnection.class,
                ThriftConnection.class);
        driverConf.setClass(LensConfConstants.DRIVER_HOOK_CLASSES_SFX, MockDriverQueryHook.class,
                DriverQueryHook.class);
        driverConf.set("hive.lock.manager", "org.apache.hadoop.hive.ql.lockmgr.EmbeddedLockManager");
        driverConf.setBoolean(HiveDriver.HS2_CALCULATE_PRIORITY, true);
        driver = new HiveDriver();
        driver.configure(driverConf, "hive", "hive1");
        drivers = Lists.<LensDriver>newArrayList(driver);
        System.out.println("TestHiveDriver created");
    }

    @BeforeMethod
    public void setDB() {
        SessionState.setCurrentSessionState(ss);
    }

    protected QueryContext createContext(final String query, Configuration conf) throws LensException {
        QueryContext context = new QueryContext(query, "testuser", new LensConf(), conf, drivers);
        // session id has to be set before calling setDriverQueriesAndPlans
        context.setLensSessionIdentifier(sessionid);
        return context;
    }

    protected QueryContext createContext(final String query, Configuration conf, LensDriver driver)
            throws LensException {
        QueryContext context = new QueryContext(query, "testuser", new LensConf(), conf, Arrays.asList(driver));
        // session id has to be set before calling setDriverQueriesAndPlans
        context.setLensSessionIdentifier(sessionid);
        return context;
    }

    protected QueryContext createContext(PreparedQueryContext query, Configuration conf) {
        QueryContext context = new QueryContext(query, "testuser", new LensConf(), conf);
        context.setLensSessionIdentifier(sessionid);
        return context;
    }

    protected ExplainQueryContext createExplainContext(final String query, Configuration conf) {
        ExplainQueryContext ectx = new ExplainQueryContext(UUID.randomUUID().toString(), query, "testuser", null,
                conf, drivers);
        ectx.setLensSessionIdentifier(sessionid);
        return ectx;
    }

    /**
     * After test.
     *
     * @throws Exception the exception
     */
    @AfterTest
    public void afterTest() throws Exception {
        verifyThriftLogs();
        driver.close();
        Hive.get(hiveConf).dropDatabase(dataBase, true, true, true);
    }

    private void verifyThriftLogs() throws IOException {
        BufferedReader br = new BufferedReader(new FileReader(new File("target/test.log")));
        for (String line = br.readLine(); line != null; line = br.readLine()) {
            if (line.contains("Update from hive")) {
                return;
            }
        }
        fail("No updates from hive found in the logs");
    }

    /**
     * Creates the test table.
     *
     * @param tableName the table name
     * @throws Exception the exception
     */
    protected void createTestTable(String tableName) throws Exception {
        int handleSize = getHandleSize();
        System.out.println("Hadoop Location: " + System.getProperty("hadoop.bin.path"));
        String createTable = "CREATE TABLE IF NOT EXISTS " + tableName + "(ID STRING)" + " TBLPROPERTIES ('"
                + LensConfConstants.STORAGE_COST + "'='500')";
        String dataLoad = "LOAD DATA LOCAL INPATH '" + TEST_DATA_FILE + "' OVERWRITE INTO TABLE " + tableName;
        // Create test table
        QueryContext context = createContext(createTable, queryConf);
        LensResultSet resultSet = driver.execute(context);
        assertNull(resultSet);
        // Load some data into the table
        context = createContext(dataLoad, queryConf);
        resultSet = driver.execute(context);
        assertNull(resultSet);
        assertHandleSize(handleSize);
    }

    /**
     * Creates the test table.
     *
     * @param tableName the table name
     * @throws Exception the exception
     */
    protected void createPartitionedTable(String tableName) throws Exception {
        int handleSize = getHandleSize();
        System.out.println("Hadoop Location: " + System.getProperty("hadoop.bin.path"));
        String createTable = "CREATE TABLE IF NOT EXISTS " + tableName + "(ID STRING)"
                + " PARTITIONED BY (dt string) TBLPROPERTIES ('" + LensConfConstants.STORAGE_COST + "'='500')";
        queryConf.setBoolean(LensConfConstants.QUERY_PERSISTENT_RESULT_INDRIVER, false);
        // Craete again
        QueryContext context = createContext(createTable, queryConf);
        LensResultSet resultSet = driver.execute(context);
        assertNull(resultSet);

        // Load some data into the table
        String dataLoad = "LOAD DATA LOCAL INPATH '" + TEST_DATA_FILE + "' OVERWRITE INTO TABLE " + tableName
                + " partition (dt='today')";
        context = createContext(dataLoad, queryConf);
        resultSet = driver.execute(context);
        assertNull(resultSet);
        assertHandleSize(handleSize);
    }

    // Tests

    /**
     * Test insert overwrite conf.
     *
     * @throws Exception the exception
     */
    @Test
    public void testInsertOverwriteConf() throws Exception {
        createTestTable("test_insert_overwrite");
        queryConf.setBoolean(LensConfConstants.QUERY_ADD_INSERT_OVEWRITE, false);
        String query = "SELECT ID FROM test_insert_overwrite";
        QueryContext context = createContext(query, queryConf);
        driver.addPersistentPath(context);
        assertEquals(context.getUserQuery(), query);
        assertNotNull(context.getDriverContext().getDriverQuery(driver));
        assertEquals(context.getDriverContext().getDriverQuery(driver), context.getUserQuery());
    }

    /**
     * Test temptable.
     *
     * @throws Exception the exception
     */
    @Test
    public void testTemptable() throws Exception {
        int handleSize = getHandleSize();
        createTestTable("test_temp");
        queryConf.setBoolean(LensConfConstants.QUERY_PERSISTENT_RESULT_INDRIVER, false);
        Hive.get(hiveConf).dropTable("test_temp_output");
        String query = "CREATE TABLE test_temp_output AS SELECT ID FROM test_temp";
        QueryContext context = createContext(query, queryConf);
        LensResultSet resultSet = driver.execute(context);
        assertNull(resultSet);
        assertHandleSize(handleSize);

        // fetch results from temp table
        String select = "SELECT * FROM test_temp_output";
        context = createContext(select, queryConf);
        resultSet = driver.execute(context);
        assertHandleSize(handleSize);
        validateInMemoryResult(resultSet, "test_temp_output");
        assertHandleSize(handleSize);
    }

    /**
     * Test execute query.
     *
     * @throws Exception the exception
     */
    @Test
    public void testExecuteQuery() throws Exception {
        int handleSize = getHandleSize();
        createTestTable("test_execute");
        LensResultSet resultSet = null;
        // Execute a select query
        queryConf.setBoolean(LensConfConstants.QUERY_PERSISTENT_RESULT_INDRIVER, false);
        String select = "SELECT ID FROM test_execute";
        QueryContext context = createContext(select, queryConf);
        resultSet = driver.execute(context);
        assertNotNull(context.getDriverConf(driver).get("mapred.job.name"));
        validateInMemoryResult(resultSet);
        queryConf.setBoolean(LensConfConstants.QUERY_PERSISTENT_RESULT_INDRIVER, true);
        context = createContext(select, queryConf);
        resultSet = driver.execute(context);
        validatePersistentResult(resultSet, TEST_DATA_FILE, context.getHDFSResultDir(), false);
        queryConf.set(LensConfConstants.QUERY_OUTPUT_DIRECTORY_FORMAT,
                "ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe'"
                        + " WITH SERDEPROPERTIES ('serialization.null.format'='-NA-',"
                        + " 'field.delim'=','  ) STORED AS TEXTFILE ");
        select = "SELECT ID, null, ID FROM test_execute";
        context = createContext(select, queryConf);
        resultSet = driver.execute(context);
        validatePersistentResult(resultSet, TEST_DATA_FILE, context.getHDFSResultDir(), true);
        assertHandleSize(handleSize);
    }

    /**
     * Validate in memory result.
     *
     * @param resultSet the result set
     * @throws LensException the lens exception
     * @throws IOException   Signals that an I/O exception has occurred.
     */
    private void validateInMemoryResult(LensResultSet resultSet) throws LensException, IOException {
        validateInMemoryResult(resultSet, null);
    }

    /**
     * Validate in memory result.
     *
     * @param resultSet   the result set
     * @param outputTable the output table
     * @throws LensException the lens exception
     * @throws IOException   Signals that an I/O exception has occurred.
     */
    private void validateInMemoryResult(LensResultSet resultSet, String outputTable)
            throws LensException, IOException {
        assertNotNull(resultSet);
        assertTrue(resultSet instanceof HiveInMemoryResultSet);
        HiveInMemoryResultSet inmemrs = (HiveInMemoryResultSet) resultSet;

        // check metadata
        LensResultSetMetadata rsMeta = inmemrs.getMetadata();
        List<ColumnDescriptor> columns = rsMeta.getColumns();
        assertNotNull(columns);
        assertEquals(columns.size(), 1);
        String expectedCol = "";
        if (outputTable != null) {
            expectedCol += outputTable + ".";
        }
        expectedCol += "ID";
        assertTrue(columns.get(0).getName().toLowerCase().equals(expectedCol.toLowerCase())
                || columns.get(0).getName().toLowerCase().equals("ID".toLowerCase()));
        assertEquals(columns.get(0).getTypeName().toLowerCase(), "STRING".toLowerCase());

        List<String> expectedRows = new ArrayList<String>();
        // Read data from the test file into expectedRows
        BufferedReader br = new BufferedReader(new InputStreamReader(new FileInputStream(TEST_DATA_FILE)));
        String line = "";
        while ((line = br.readLine()) != null) {
            expectedRows.add(line.trim());
        }
        br.close();

        List<String> actualRows = new ArrayList<String>();
        while (inmemrs.hasNext()) {
            List<Object> row = inmemrs.next().getValues();
            actualRows.add((String) row.get(0));
        }
        assertEquals(actualRows, expectedRows);
    }

    /**
     * The Class FailHook.
     */
    public static class FailHook implements HiveDriverRunHook {

        /*
         * (non-Javadoc)
         *
         * @see
         * org.apache.hadoop.hive.ql.HiveDriverRunHook#postDriverRun(org.apache.hadoop.hive.ql.HiveDriverRunHookContext)
         */
        @Override
        public void postDriverRun(HiveDriverRunHookContext arg0) throws Exception {
            // TODO Auto-generated method stub

        }

        /*
         * (non-Javadoc)
         *
         * @see org.apache.hadoop.hive.ql.HiveDriverRunHook#preDriverRun(org.apache.hadoop.hive.ql.HiveDriverRunHookContext)
         */
        @Override
        public void preDriverRun(HiveDriverRunHookContext arg0) throws Exception {
            throw new LensException("Failing this run");
        }

    }

    // executeAsync

    /**
     * Test execute query async.
     *
     * @throws Exception the exception
     */
    @Test
    public void testExecuteQueryAsync() throws Exception {
        int handleSize = getHandleSize();
        createTestTable("test_execute_sync");

        // Now run a command that would fail
        String expectFail = "SELECT ID FROM test_execute_sync";
        queryConf.setBoolean(LensConfConstants.QUERY_PERSISTENT_RESULT_INDRIVER, true);
        Configuration failConf = new Configuration(queryConf);
        failConf.set("hive.exec.driver.run.hooks", FailHook.class.getName());
        QueryContext context = createContext(expectFail, failConf);
        driver.executeAsync(context);
        assertHandleSize(handleSize + 1);
        validateExecuteAsync(context, DriverQueryState.FAILED, true, false);
        assertHandleSize(handleSize + 1);
        driver.closeQuery(context.getQueryHandle());
        assertHandleSize(handleSize);
        // Async select query
        String select = "SELECT ID FROM test_execute_sync";
        queryConf.setBoolean(LensConfConstants.QUERY_PERSISTENT_RESULT_INDRIVER, false);
        context = createContext(select, queryConf);
        driver.executeAsync(context);
        assertNotNull(context.getDriverConf(driver).get("mapred.job.name"));
        assertNotNull(context.getDriverConf(driver).get("mapred.job.priority"));
        assertHandleSize(handleSize + 1);
        validateExecuteAsync(context, DriverQueryState.SUCCESSFUL, false, false);
        driver.closeQuery(context.getQueryHandle());
        assertHandleSize(handleSize);

        queryConf.setBoolean(LensConfConstants.QUERY_PERSISTENT_RESULT_INDRIVER, true);
        context = createContext(select, queryConf);
        driver.executeAsync(context);
        assertHandleSize(handleSize + 1);
        validateExecuteAsync(context, DriverQueryState.SUCCESSFUL, true, false);
        driver.closeQuery(context.getQueryHandle());
        assertHandleSize(handleSize);

        queryConf.set(LensConfConstants.QUERY_OUTPUT_DIRECTORY_FORMAT,
                "ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe'"
                        + " WITH SERDEPROPERTIES ('serialization.null.format'='-NA-',"
                        + " 'field.delim'=','  ) STORED AS TEXTFILE ");
        select = "SELECT ID, null, ID FROM test_execute_sync";
        context = createContext(select, queryConf);
        driver.executeAsync(context);
        assertHandleSize(handleSize + 1);
        validateExecuteAsync(context, DriverQueryState.SUCCESSFUL, true, true);
        driver.closeQuery(context.getQueryHandle());
        assertHandleSize(handleSize);
    }

    /**
     * Validate execute async.
     *
     * @param ctx          the ctx
     * @param finalState   the final state
     * @param isPersistent the is persistent
     * @param formatNulls  the format nulls
     * @param driver       the driver
     * @throws Exception the exception
     */
    protected void validateExecuteAsync(QueryContext ctx, DriverQueryState finalState, boolean isPersistent,
            boolean formatNulls, HiveDriver driver) throws Exception {
        waitForAsyncQuery(ctx, driver);
        driver.updateStatus(ctx);
        assertEquals(ctx.getDriverStatus().getState(), finalState, "Expected query to finish with" + finalState);
        assertTrue(ctx.getDriverStatus().getDriverFinishTime() > 0);
        if (finalState.equals(DriverQueryState.SUCCESSFUL)) {
            System.out.println("Progress:" + ctx.getDriverStatus().getProgressMessage());
            assertNotNull(ctx.getDriverStatus().getProgressMessage());
            if (!isPersistent) {
                validateInMemoryResult(driver.fetchResultSet(ctx));
            } else {
                validatePersistentResult(driver.fetchResultSet(ctx), TEST_DATA_FILE, ctx.getHDFSResultDir(),
                        formatNulls);
            }
        } else if (finalState.equals(DriverQueryState.FAILED)) {
            System.out.println("Error:" + ctx.getDriverStatus().getErrorMessage());
            System.out.println("Status:" + ctx.getDriverStatus().getStatusMessage());
            assertNotNull(ctx.getDriverStatus().getErrorMessage());
        }
    }

    /**
     * Validate execute async.
     *
     * @param ctx          the ctx
     * @param finalState   the final state
     * @param isPersistent the is persistent
     * @param formatNulls  the format nulls
     * @throws Exception the exception
     */
    protected void validateExecuteAsync(QueryContext ctx, DriverQueryState finalState, boolean isPersistent,
            boolean formatNulls) throws Exception {
        validateExecuteAsync(ctx, finalState, isPersistent, formatNulls, driver);
    }

    /**
     * Test cancel async query.
     *
     * @throws Exception the exception
     */
    @Test
    public void testCancelAsyncQuery() throws Exception {
        int handleSize = getHandleSize();
        createTestTable("test_cancel_async");
        queryConf.setBoolean(LensConfConstants.QUERY_PERSISTENT_RESULT_INDRIVER, false);
        QueryContext context = createContext("select a.id aid, b.id bid from "
                + "((SELECT ID FROM test_cancel_async) a full outer join (select id from test_cancel_async) b)",
                queryConf);
        driver.executeAsync(context);
        driver.cancelQuery(context.getQueryHandle());
        driver.updateStatus(context);
        assertEquals(context.getDriverStatus().getState(), DriverQueryState.CANCELED,
                "Expecting query to be cancelled");
        driver.closeQuery(context.getQueryHandle());
        assertHandleSize(handleSize);

        try {
            driver.cancelQuery(context.getQueryHandle());
            fail("Cancel on closed query should throw error");
        } catch (LensException exc) {
            assertTrue(exc.getMessage().startsWith("Query not found"));
        }
    }

    /**
     * Validate persistent result.
     *
     * @param resultSet   the result set
     * @param dataFile    the data file
     * @param outptuDir   the outptu dir
     * @param formatNulls the format nulls
     * @throws Exception the exception
     */
    private void validatePersistentResult(LensResultSet resultSet, String dataFile, Path outptuDir,
            boolean formatNulls) throws Exception {
        assertTrue(resultSet instanceof HivePersistentResultSet,
                "resultset class: " + resultSet.getClass().getName());
        HivePersistentResultSet persistentResultSet = (HivePersistentResultSet) resultSet;
        String path = persistentResultSet.getOutputPath();

        Path actualPath = new Path(path);
        FileSystem fs = actualPath.getFileSystem(driverConf);
        assertEquals(actualPath, fs.makeQualified(outptuDir));
        List<String> actualRows = new ArrayList<String>();
        for (FileStatus stat : fs.listStatus(actualPath, new PathFilter() {
            @Override
            public boolean accept(Path path) {
                return !new File(path.toUri()).isDirectory();
            }
        })) {
            FSDataInputStream in = fs.open(stat.getPath());
            BufferedReader br = null;
            try {
                br = new BufferedReader(new InputStreamReader(in));
                String line = "";

                while ((line = br.readLine()) != null) {
                    System.out.println("Actual:" + line);
                    actualRows.add(line.trim());
                }
            } finally {
                if (br != null) {
                    br.close();
                }
            }
        }

        BufferedReader br = null;
        List<String> expectedRows = new ArrayList<String>();

        try {
            br = new BufferedReader(new FileReader(new File(dataFile)));
            String line = "";
            while ((line = br.readLine()) != null) {
                String row = line.trim();
                if (formatNulls) {
                    row += ",-NA-,";
                    row += line.trim();
                }
                expectedRows.add(row);
            }
        } finally {
            if (br != null) {
                br.close();
            }
        }
        assertEquals(actualRows, expectedRows);
    }

    /**
     * Test persistent result set.
     *
     * @throws Exception the exception
     */
    @Test
    public void testPersistentResultSet() throws Exception {
        int handleSize = getHandleSize();
        createTestTable("test_persistent_result_set");
        queryConf.setBoolean(LensConfConstants.QUERY_PERSISTENT_RESULT_INDRIVER, true);
        queryConf.setBoolean(LensConfConstants.QUERY_ADD_INSERT_OVEWRITE, true);
        queryConf.set(LensConfConstants.RESULT_SET_PARENT_DIR, testOutputDir);
        QueryContext ctx = createContext("SELECT ID FROM test_persistent_result_set", queryConf);
        LensResultSet resultSet = driver.execute(ctx);
        validatePersistentResult(resultSet, TEST_DATA_FILE, ctx.getHDFSResultDir(), false);
        assertHandleSize(handleSize);

        ctx = createContext("SELECT ID FROM test_persistent_result_set", queryConf);
        driver.executeAsync(ctx);
        assertHandleSize(handleSize + 1);
        validateExecuteAsync(ctx, DriverQueryState.SUCCESSFUL, true, false);
        driver.closeQuery(ctx.getQueryHandle());
        assertHandleSize(handleSize);

        queryConf.set(LensConfConstants.QUERY_OUTPUT_DIRECTORY_FORMAT,
                "ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe'"
                        + " WITH SERDEPROPERTIES ('serialization.null.format'='-NA-',"
                        + " 'field.delim'=','  ) STORED AS TEXTFILE ");
        ctx = createContext("SELECT ID, null, ID FROM test_persistent_result_set", queryConf);
        resultSet = driver.execute(ctx);
        assertHandleSize(handleSize);
        validatePersistentResult(resultSet, TEST_DATA_FILE, ctx.getHDFSResultDir(), true);
        driver.closeQuery(ctx.getQueryHandle());
        assertHandleSize(handleSize);

        ctx = createContext("SELECT ID, null, ID FROM test_persistent_result_set", queryConf);
        driver.executeAsync(ctx);
        assertHandleSize(handleSize + 1);
        validateExecuteAsync(ctx, DriverQueryState.SUCCESSFUL, true, true);
        driver.closeQuery(ctx.getQueryHandle());
        assertHandleSize(handleSize);
    }

    /**
     * Wait for async query.
     *
     * @param ctx    the ctx
     * @param driver the driver
     * @throws Exception the exception
     */
    private void waitForAsyncQuery(QueryContext ctx, HiveDriver driver) throws Exception {
        while (true) {
            driver.updateStatus(ctx);
            System.out.println("#W Waiting for query " + ctx.getQueryHandle() + " status: "
                    + ctx.getDriverStatus().getState());
            assertNotNull(ctx.getDriverStatus());
            if (ctx.getDriverStatus().isFinished()) {
                assertTrue(ctx.getDriverStatus().getDriverFinishTime() > 0);
                break;
            }
            System.out.println("Progress:" + ctx.getDriverStatus().getProgressMessage());
            Thread.sleep(1000);
            assertTrue(ctx.getDriverStatus().getDriverStartTime() > 0);
        }
    }

    @Test(expectedExceptions = { UnsupportedOperationException.class })
    public void testEstimateNativeQuery() throws Exception {
        createTestTable("test_estimate");
        SessionState.setCurrentSessionState(ss);
        QueryCost cost = driver.estimate(createExplainContext("SELECT ID FROM test_estimate", queryConf));
        assertEquals(cost.getEstimatedResourceUsage(), Double.MAX_VALUE);
        cost.getEstimatedExecTimeMillis();

    }

    @Test(expectedExceptions = { UnsupportedOperationException.class })
    public void testEstimateOlapQuery() throws Exception {
        SessionState.setCurrentSessionState(ss);
        ExplainQueryContext ctx = createExplainContext("cube SELECT ID FROM test_cube", queryConf);
        ctx.setOlapQuery(true);
        ctx.getDriverContext().setDriverRewriterPlan(driver, new DriverQueryPlan() {
            @Override
            public String getPlan() {
                return null;
            }

            @Override
            public QueryCost getCost() {
                return null;
            }

            @Override
            public Map<String, Set<?>> getPartitions() {
                return Maps.newHashMap();
            }
        });
        QueryCost cost = driver.estimate(ctx);
        assertEquals(cost.getEstimatedResourceUsage(), 0.0);
        cost.getEstimatedExecTimeMillis();
    }

    @Test
    public void testExplainNativeFailingQuery() throws Exception {
        SessionState.setCurrentSessionState(ss);
        try {
            driver.estimate(createExplainContext("SELECT ID FROM nonexist", queryConf));
            fail("Should not reach here");
        } catch (LensException e) {
            assertTrue(LensUtil.getCauseMessage(e).contains("Line 1:32 Table not found 'nonexist'"));
        }
    }

    // explain

    /**
     * Test explain.
     *
     * @throws Exception the exception
     */
    @Test
    public void testExplain() throws Exception {
        int handleSize = getHandleSize();
        SessionState.setCurrentSessionState(ss);
        SessionState.get().setCurrentDatabase(dataBase);
        createTestTable("test_explain");
        DriverQueryPlan plan = driver.explain(createExplainContext("SELECT ID FROM test_explain", queryConf));
        assertTrue(plan instanceof HiveQueryPlan);
        assertEquals(plan.getTableWeight(dataBase + ".test_explain"), 500.0);
        assertHandleSize(handleSize);

        // test execute prepare
        PreparedQueryContext pctx = new PreparedQueryContext("SELECT ID FROM test_explain", null, queryConf,
                drivers);
        pctx.setSelectedDriver(driver);
        pctx.setLensSessionIdentifier(sessionid);

        SessionState.setCurrentSessionState(ss);
        Configuration inConf = new Configuration(queryConf);
        inConf.setBoolean(LensConfConstants.QUERY_PERSISTENT_RESULT_INDRIVER, false);
        plan = driver.explainAndPrepare(pctx);
        QueryContext qctx = createContext(pctx, inConf);
        LensResultSet result = driver.execute(qctx);
        assertHandleSize(handleSize);
        validateInMemoryResult(result);

        // test execute prepare async
        queryConf.setBoolean(LensConfConstants.QUERY_PERSISTENT_RESULT_INDRIVER, true);
        qctx = createContext(pctx, queryConf);
        driver.executeAsync(qctx);
        assertNotNull(qctx.getDriverOpHandle());
        validateExecuteAsync(qctx, DriverQueryState.SUCCESSFUL, true, false);
        assertHandleSize(handleSize + 1);

        driver.closeQuery(qctx.getQueryHandle());
        assertHandleSize(handleSize);

        // for backward compatibility
        qctx = createContext(pctx, inConf);
        qctx.setQueryHandle(new QueryHandle(pctx.getPrepareHandle().getPrepareHandleId()));
        result = driver.execute(qctx);
        assertNotNull(qctx.getDriverOpHandle());
        assertHandleSize(handleSize);
        validateInMemoryResult(result);
        // test execute prepare async
        qctx = createContext(pctx, queryConf);
        qctx.setQueryHandle(new QueryHandle(pctx.getPrepareHandle().getPrepareHandleId()));
        driver.executeAsync(qctx);
        assertHandleSize(handleSize + 1);
        validateExecuteAsync(qctx, DriverQueryState.SUCCESSFUL, true, false);

        driver.closeQuery(qctx.getQueryHandle());
        driver.closePreparedQuery(pctx.getPrepareHandle());
        assertHandleSize(handleSize);
    }

    /**
     * Test explain partitioned table
     *
     * @throws Exception the exception
     */
    @Test
    public void testExplainPartitionedTable() throws Exception {
        int handleSize = getHandleSize();
        createPartitionedTable("test_part_table");
        // acquire
        SessionState.setCurrentSessionState(ss);
        DriverQueryPlan plan = driver.explain(createExplainContext("SELECT ID FROM test_part_table", queryConf));
        assertHandleSize(handleSize);
        assertTrue(plan instanceof HiveQueryPlan);
        assertNotNull(plan.getTablesQueried());
        assertEquals(plan.getTablesQueried().size(), 1);
        System.out.println("Tables:" + plan.getTablesQueried());
        assertEquals(plan.getTableWeight(dataBase + ".test_part_table"), 500.0);
        System.out.println("Parts:" + plan.getPartitions());
        assertFalse(plan.getPartitions().isEmpty());
        assertEquals(plan.getPartitions().size(), 1);
        assertTrue(((String) plan.getPartitions().get(dataBase + ".test_part_table").iterator().next())
                .contains("today"));
        assertTrue(((String) plan.getPartitions().get(dataBase + ".test_part_table").iterator().next())
                .contains("dt"));
    }

    /**
     * Test explain output.
     *
     * @throws Exception the exception
     */
    @Test
    public void testExplainOutput() throws Exception {
        int handleSize = getHandleSize();
        createTestTable("explain_test_1");
        createTestTable("explain_test_2");
        SessionState.setCurrentSessionState(ss);
        DriverQueryPlan plan = driver.explain(createExplainContext("SELECT explain_test_1.ID, count(1) FROM "
                + " explain_test_1  join explain_test_2 on explain_test_1.ID = explain_test_2.ID"
                + " WHERE explain_test_1.ID = 'foo' or explain_test_2.ID = 'bar'" + " GROUP BY explain_test_1.ID",
                queryConf));

        assertHandleSize(handleSize);
        assertTrue(plan instanceof HiveQueryPlan);
        assertNotNull(plan.getTablesQueried());
        assertEquals(plan.getTablesQueried().size(), 2);
        assertNotNull(plan.getTableWeights());
        assertTrue(plan.getTableWeights().containsKey(dataBase + ".explain_test_1"));
        assertTrue(plan.getTableWeights().containsKey(dataBase + ".explain_test_2"));
        assertTrue(plan.getPlan() != null && !plan.getPlan().isEmpty());
        driver.closeQuery(plan.getHandle());
    }

    /**
     * Test explain output persistent.
     *
     * @throws Exception the exception
     */
    @Test
    public void testExplainOutputPersistent() throws Exception {
        int handleSize = getHandleSize();
        createTestTable("explain_test_1");
        queryConf.setBoolean(LensConfConstants.QUERY_PERSISTENT_RESULT_INDRIVER, true);
        SessionState.setCurrentSessionState(ss);
        String query2 = "SELECT DISTINCT ID FROM explain_test_1";
        PreparedQueryContext pctx = createPreparedQueryContext(query2);
        pctx.setSelectedDriver(driver);
        pctx.setLensSessionIdentifier(sessionid);
        DriverQueryPlan plan2 = driver.explainAndPrepare(pctx);
        // assertNotNull(plan2.getResultDestination());
        assertHandleSize(handleSize);
        assertNotNull(plan2.getTablesQueried());
        assertEquals(plan2.getTablesQueried().size(), 1);
        assertTrue(plan2.getTableWeights().containsKey(dataBase + ".explain_test_1"));
        QueryContext ctx = createContext(pctx, queryConf);
        LensResultSet resultSet = driver.execute(ctx);
        assertHandleSize(handleSize);
        HivePersistentResultSet persistentResultSet = (HivePersistentResultSet) resultSet;
        String path = persistentResultSet.getOutputPath();
        assertEquals(ctx.getDriverResultPath(), path);
        driver.closeQuery(plan2.getHandle());
    }

    private PreparedQueryContext createPreparedQueryContext(String query2) {
        PreparedQueryContext pctx = new PreparedQueryContext(query2, null, queryConf, drivers);
        pctx.setSelectedDriver(driver);
        pctx.setLensSessionIdentifier(sessionid);
        return pctx;
    }

    @DataProvider
    public Object[][] priorityDataProvider() throws IOException, ParseException {
        BufferedReader br = new BufferedReader(
                new InputStreamReader(TestHiveDriver.class.getResourceAsStream("/priority_tests.data")));
        String line;
        int numTests = Integer.parseInt(br.readLine());
        Object[][] data = new Object[numTests][2];
        for (int i = 0; i < numTests; i++) {
            String[] kv = br.readLine().split("\\s*:\\s*");
            final Set<FactPartition> partitions = getFactParts(Arrays.asList(kv[0].trim().split("\\s*,\\s*")));
            final Priority expected = Priority.valueOf(kv[1]);
            data[i] = new Object[] { partitions, expected };
        }
        return data;
    }

    /**
     * Testing Duration Based Priority Logic by mocking everything except partitions.
     *
     * @throws IOException
     * @throws LensException
     * @throws ParseException
     */
    @Test(dataProvider = "priorityDataProvider")
    public void testPriority(final Set<FactPartition> partitions, Priority expected) throws Exception {
        Configuration conf = new Configuration();
        QueryContext ctx = createContext("test priority query", conf);
        ctx.getDriverContext().setDriverRewriterPlan(driver, new DriverQueryPlan() {

            @Override
            public String getPlan() {
                return null;
            }

            @Override
            public QueryCost getCost() {
                return null;
            }
        });

        ctx.getDriverContext().getDriverRewriterPlan(driver).getPartitions()
                .putAll(new HashMap<String, Set<FactPartition>>() {
                    {
                        put("table1", partitions);
                    }
                });
        // table weights only for first calculation
        ctx.getDriverContext().getDriverRewriterPlan(driver).getTableWeights()
                .putAll(new HashMap<String, Double>() {
                    {
                        put("table1", 1.0);
                    }
                });
        ctx.setOlapQuery(true);
        Priority priority = driver.decidePriority(ctx);
        assertEquals(priority, expected, "cost: " + ctx.getDriverQueryCost(driver) + "priority: " + priority);
        assertEquals(ctx.getConf().get("mapred.job.priority"), priority.toString());
        assertEquals(driver.decidePriority(ctx, alwaysNormalPriorityDecider), Priority.NORMAL);
    }

    @Test
    public void testPriorityWithoutFactPartitions() throws LensException {
        // test priority without fact partitions
        QueryContext ctx = createContext("test priority query", queryConf);
        ctx.getDriverContext().setDriverRewriterPlan(driver, new DriverQueryPlan() {

            @Override
            public String getPlan() {
                return null;
            }

            @Override
            public QueryCost getCost() {
                return null;
            }
        });

        ctx.getDriverContext().getDriverRewriterPlan(driver).getPartitions()
                .putAll(new HashMap<String, Set<String>>() {
                    {
                        put("table1", new HashSet<String>());
                    }
                });
        ctx.getDriverContext().getDriverRewriterPlan(driver).getTableWeights()
                .putAll(new HashMap<String, Double>() {
                    {
                        put("table1", 1.0);
                    }
                });
        ctx.setDriverCost(driver, driver.queryCostCalculator.calculateCost(ctx, driver));
        assertEquals(driver.decidePriority(ctx, driver.queryPriorityDecider), Priority.VERY_HIGH);
        assertEquals(driver.decidePriority(ctx, alwaysNormalPriorityDecider), Priority.NORMAL);

        // test priority without rewriter plan
        ctx = createContext("test priority query", queryConf);
        ctx.getDriverContext().setDriverRewriterPlan(driver, new DriverQueryPlan() {
            @Override
            public String getPlan() {
                return null;
            }

            @Override
            public QueryCost getCost() {
                return null;
            }
        });
        ctx.setDriverCost(driver, driver.queryCostCalculator.calculateCost(ctx, driver));
        assertEquals(driver.decidePriority(ctx), Priority.VERY_HIGH);
        assertEquals(alwaysNormalPriorityDecider.decidePriority(ctx.getDriverQueryCost(driver)), Priority.NORMAL);
    }

    private Set<FactPartition> getFactParts(List<String> partStrings) throws ParseException {
        Set<FactPartition> factParts = new HashSet<FactPartition>();
        for (String partStr : partStrings) {
            String[] partEls = partStr.split(" ");
            UpdatePeriod p = null;
            String partSpec = partEls[1];
            switch (partSpec.length()) {
            case 7: //monthly
                p = UpdatePeriod.MONTHLY;
                break;
            case 10: // daily
                p = UpdatePeriod.DAILY;
                break;
            case 13: // hourly
                p = UpdatePeriod.HOURLY;
                break;
            }
            FactPartition part = new FactPartition(partEls[0], p.parse(partSpec), p, null, p.format(),
                    Collections.singleton("table1"));
            factParts.add(part);
        }
        return factParts;
    }

    private int getHandleSize() {
        return driver.getHiveHandleSize();
    }

    private void assertHandleSize(int handleSize) {
        assertEquals(getHandleSize(), handleSize,
                "Unexpected handle size, all handles: " + driver.getHiveHandles());
    }
}