com.inmobi.grill.server.query.TestQueryService.java Source code

Java tutorial

Introduction

Here is the source code for com.inmobi.grill.server.query.TestQueryService.java

Source

package com.inmobi.grill.server.query;

/*
 * #%L
 * Grill Server
 * %%
 * Copyright (C) 2014 Inmobi
 * %%
 * 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%
 */

import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertTrue;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

import javax.ws.rs.BadRequestException;
import javax.ws.rs.client.Entity;
import javax.ws.rs.core.Application;
import javax.ws.rs.core.GenericType;
import javax.ws.rs.core.Response;
import javax.ws.rs.client.WebTarget;
import javax.ws.rs.core.MediaType;

import com.inmobi.grill.api.APIResult;
import com.inmobi.grill.api.GrillConf;
import com.inmobi.grill.api.GrillException;
import com.inmobi.grill.api.GrillSessionHandle;
import com.inmobi.grill.api.query.GrillPreparedQuery;
import com.inmobi.grill.api.query.GrillQuery;
import com.inmobi.grill.api.query.InMemoryQueryResult;
import com.inmobi.grill.api.query.PersistentQueryResult;
import com.inmobi.grill.api.query.QueryHandle;
import com.inmobi.grill.api.query.QueryHandleWithResultSet;
import com.inmobi.grill.api.query.QueryPlan;
import com.inmobi.grill.api.query.QueryPrepareHandle;
import com.inmobi.grill.api.query.QueryResultSetMetadata;
import com.inmobi.grill.api.query.QueryStatus;
import com.inmobi.grill.api.query.QueryStatus.Status;
import com.inmobi.grill.driver.hive.TestHiveDriver.FailHook;
import com.inmobi.grill.server.GrillJerseyTest;
import com.inmobi.grill.server.GrillServices;
import com.inmobi.grill.server.api.GrillConfConstants;
import com.inmobi.grill.server.api.metrics.MetricsService;
import com.inmobi.grill.server.query.QueryApp;
import com.inmobi.grill.server.query.QueryExecutionServiceImpl;

import org.glassfish.jersey.client.ClientConfig;
import org.glassfish.jersey.media.multipart.FormDataBodyPart;
import org.glassfish.jersey.media.multipart.FormDataContentDisposition;
import org.glassfish.jersey.media.multipart.FormDataMultiPart;
import org.glassfish.jersey.media.multipart.MultiPartFeature;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.testng.Assert;
import org.testng.annotations.AfterClass;
import org.testng.annotations.AfterTest;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.BeforeTest;
import org.testng.annotations.Test;

@Test(groups = "unit-test")
public class TestQueryService extends GrillJerseyTest {

    QueryExecutionServiceImpl queryService;
    MetricsService metricsSvc;
    GrillSessionHandle grillSessionId;

    @BeforeTest
    public void setUp() throws Exception {
        super.setUp();
        queryService = (QueryExecutionServiceImpl) GrillServices.get().getService("query");
        metricsSvc = (MetricsService) GrillServices.get().getService(MetricsService.NAME);
        grillSessionId = queryService.openSession("foo", "bar", new HashMap<String, String>());
    }

    @AfterTest
    public void tearDown() throws Exception {
        queryService.closeSession(grillSessionId);
        super.tearDown();
    }

    @BeforeClass
    public void createTables() throws InterruptedException {
        createTable(testTable);
        loadData(testTable);
    }

    @AfterClass
    public void dropTables() throws InterruptedException {
        dropTable(testTable);
    }

    @Override
    protected Application configure() {
        return new QueryApp();
    }

    @Override
    protected void configureClient(ClientConfig config) {
        config.register(MultiPartFeature.class);
    }

    private static String testTable = "TEST_TABLE";
    public static final String TEST_DATA_FILE = "../grill-driver-hive/testdata/testdata2.txt";

    private void createTable(String tblName) throws InterruptedException {
        GrillConf conf = new GrillConf();
        conf.addProperty(GrillConfConstants.GRILL_PERSISTENT_RESULT_SET, "false");
        final WebTarget target = target().path("queryapi/queries");

        final FormDataMultiPart mp = new FormDataMultiPart();
        String createTable = "CREATE TABLE IF NOT EXISTS " + tblName + "(ID INT, IDSTR STRING)";

        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(), createTable));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "execute"));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(), conf,
                MediaType.APPLICATION_XML_TYPE));

        final QueryHandle handle = target.request().post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA_TYPE),
                QueryHandle.class);

        // wait till the query finishes
        GrillQuery ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                .get(GrillQuery.class);
        QueryStatus stat = ctx.getStatus();
        while (!stat.isFinished()) {
            ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                    .get(GrillQuery.class);
            stat = ctx.getStatus();
            Thread.sleep(1000);
        }
        assertTrue(ctx.getSubmissionTime() > 0);
        assertTrue(ctx.getLaunchTime() > 0);
        assertTrue(ctx.getDriverStartTime() > 0);
        assertTrue(ctx.getDriverFinishTime() > 0);
        assertTrue(ctx.getFinishTime() > 0);
        Assert.assertEquals(ctx.getStatus().getStatus(), QueryStatus.Status.SUCCESSFUL);
    }

    private void loadData(String tblName) throws InterruptedException {
        GrillConf conf = new GrillConf();
        conf.addProperty(GrillConfConstants.GRILL_PERSISTENT_RESULT_SET, "false");
        final WebTarget target = target().path("queryapi/queries");

        final FormDataMultiPart mp = new FormDataMultiPart();
        String dataLoad = "LOAD DATA LOCAL INPATH '" + TEST_DATA_FILE + "' OVERWRITE INTO TABLE " + tblName;

        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(), dataLoad));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "execute"));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(), conf,
                MediaType.APPLICATION_XML_TYPE));

        final QueryHandle handle = target.request().post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA_TYPE),
                QueryHandle.class);

        // wait till the query finishes
        GrillQuery ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                .get(GrillQuery.class);
        QueryStatus stat = ctx.getStatus();
        while (!stat.isFinished()) {
            ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                    .get(GrillQuery.class);
            stat = ctx.getStatus();
            Thread.sleep(1000);
        }
        Assert.assertEquals(ctx.getStatus().getStatus(), QueryStatus.Status.SUCCESSFUL);

    }

    private void dropTable(String tblName) throws InterruptedException {
        GrillConf conf = new GrillConf();
        conf.addProperty(GrillConfConstants.GRILL_PERSISTENT_RESULT_SET, "false");
        final WebTarget target = target().path("queryapi/queries");

        final FormDataMultiPart mp = new FormDataMultiPart();
        String createTable = "DROP TABLE IF EXISTS " + tblName;

        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(), createTable));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "execute"));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(), conf,
                MediaType.APPLICATION_XML_TYPE));

        final QueryHandle handle = target.request().post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA_TYPE),
                QueryHandle.class);

        // wait till the query finishes
        GrillQuery ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                .get(GrillQuery.class);
        QueryStatus stat = ctx.getStatus();
        while (!stat.isFinished()) {
            ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                    .get(GrillQuery.class);
            stat = ctx.getStatus();
            Thread.sleep(1000);
        }
        Assert.assertEquals(ctx.getStatus().getStatus(), QueryStatus.Status.SUCCESSFUL);
    }

    // test get a random query, should return 400
    @Test
    public void testGetRandomQuery() {
        final WebTarget target = target().path("queryapi/queries");

        Response rs = target.path("random").queryParam("sessionid", grillSessionId).request().get();
        Assert.assertEquals(rs.getStatus(), 400);
    }

    @Test
    public void testLaunchFail() throws InterruptedException {
        final WebTarget target = target().path("queryapi/queries");
        long failedQueries = metricsSvc.getTotalFailedQueries();
        System.out.println("%% " + failedQueries);
        GrillConf conf = new GrillConf();
        final FormDataMultiPart mp = new FormDataMultiPart();
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(),
                "select ID from non_exist_table"));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "execute"));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(), conf,
                MediaType.APPLICATION_XML_TYPE));
        final QueryHandle handle = target.request().post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA_TYPE),
                QueryHandle.class);

        Assert.assertNotNull(handle);

        GrillQuery ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                .get(GrillQuery.class);
        QueryStatus stat = ctx.getStatus();
        while (!stat.isFinished()) {
            ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                    .get(GrillQuery.class);
            stat = ctx.getStatus();
            System.out.println("%% query " + ctx.getQueryHandle() + " status:" + stat);
            Thread.sleep(1000);
        }

        assertTrue(ctx.getSubmissionTime() > 0);
        assertEquals(ctx.getLaunchTime(), 0);
        assertEquals(ctx.getDriverStartTime(), 0);
        assertEquals(ctx.getDriverFinishTime(), 0);
        assertTrue(ctx.getFinishTime() > 0);
        Assert.assertEquals(ctx.getStatus().getStatus(), QueryStatus.Status.FAILED);
        System.out.println("%% " + metricsSvc.getTotalFailedQueries());
        Assert.assertEquals(metricsSvc.getTotalFailedQueries(), failedQueries + 1);
    }

    // test with execute async post, get all queries, get query context,
    // get wrong uuid query
    @Test
    public void testQueriesAPI() throws InterruptedException {
        // test post execute op
        final WebTarget target = target().path("queryapi/queries");
        GrillConf conf = new GrillConf();
        conf.addProperty("hive.exec.driver.run.hooks", FailHook.class.getCanonicalName());
        final FormDataMultiPart mp = new FormDataMultiPart();
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(),
                "select ID from " + testTable));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "execute"));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(), conf,
                MediaType.APPLICATION_XML_TYPE));

        long queuedQueries = metricsSvc.getQueuedQueries();
        long runningQueries = metricsSvc.getRunningQueries();
        long finishedQueries = metricsSvc.getFinishedQueries();

        final QueryHandle handle = target.request().post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA_TYPE),
                QueryHandle.class);

        Assert.assertNotNull(handle);

        // Get all queries
        // XML 
        List<QueryHandle> allQueriesXML = target.queryParam("sessionid", grillSessionId)
                .request(MediaType.APPLICATION_XML).get(new GenericType<List<QueryHandle>>() {
                });
        Assert.assertTrue(allQueriesXML.size() >= 1);

        //JSON
        //  List<QueryHandle> allQueriesJSON = target.request(
        //      MediaType.APPLICATION_JSON).get(new GenericType<List<QueryHandle>>() {
        //  });
        //  Assert.assertEquals(allQueriesJSON.size(), 1);
        //JAXB
        List<QueryHandle> allQueries = (List<QueryHandle>) target.queryParam("sessionid", grillSessionId).request()
                .get(new GenericType<List<QueryHandle>>() {
                });
        Assert.assertTrue(allQueries.size() >= 1);
        Assert.assertTrue(allQueries.contains(handle));

        // Get query
        // Invocation.Builder builderjson = target.path(handle.toString()).request(MediaType.APPLICATION_JSON);
        // String responseJSON = builderjson.get(String.class);
        // System.out.println("query JSON:" + responseJSON);
        String queryXML = target.path(handle.toString()).queryParam("sessionid", grillSessionId)
                .request(MediaType.APPLICATION_XML).get(String.class);
        System.out.println("query XML:" + queryXML);

        Response response = target.path(handle.toString() + "001").queryParam("sessionid", grillSessionId).request()
                .get();
        Assert.assertEquals(response.getStatus(), 404);

        GrillQuery ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                .get(GrillQuery.class);
        // Assert.assertEquals(ctx.getStatus().getStatus(), QueryStatus.Status.QUEUED);

        // wait till the query finishes
        QueryStatus stat = ctx.getStatus();
        while (!stat.isFinished()) {
            ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                    .get(GrillQuery.class);
            stat = ctx.getStatus();
            switch (stat.getStatus()) {
            case RUNNING:
                assertEquals(metricsSvc.getRunningQueries(), runningQueries + 1);
                break;
            case QUEUED:
                assertEquals(metricsSvc.getQueuedQueries(), queuedQueries + 1);
                break;
            default: // nothing
            }
            Thread.sleep(1000);
        }
        assertTrue(ctx.getSubmissionTime() > 0);
        assertTrue(ctx.getLaunchTime() > 0);
        assertTrue(ctx.getDriverStartTime() > 0);
        assertTrue(ctx.getDriverFinishTime() > 0);
        assertTrue(ctx.getFinishTime() > 0);
        Assert.assertEquals(ctx.getStatus().getStatus(), QueryStatus.Status.FAILED);
        assertEquals(metricsSvc.getFinishedQueries(), finishedQueries + 1);

        // Update conf for query
        final FormDataMultiPart confpart = new FormDataMultiPart();
        conf = new GrillConf();
        conf.addProperty("my.property", "myvalue");
        confpart.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        confpart.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(),
                conf, MediaType.APPLICATION_XML_TYPE));
        APIResult updateConf = target.path(handle.toString()).request()
                .put(Entity.entity(confpart, MediaType.MULTIPART_FORM_DATA_TYPE), APIResult.class);
        Assert.assertEquals(updateConf.getStatus(), APIResult.Status.FAILED);
    }

    @Test
    public void testExecuteWithoutSessionId() throws Exception {
        // test post execute op
        final WebTarget target = target().path("queryapi/queries");
        GrillConf conf = new GrillConf();
        conf.addProperty("hive.exec.driver.run.hooks", FailHook.class.getCanonicalName());
        final FormDataMultiPart mp = new FormDataMultiPart();

        /**
         * We are not passing session id in this test
         */
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(),
                "select ID from " + testTable));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "execute"));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(), conf,
                MediaType.APPLICATION_XML_TYPE));

        try {
            final QueryHandle handle = target.request().post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA_TYPE),
                    QueryHandle.class);
            Assert.fail("Should have thrown bad request error");
        } catch (BadRequestException badReqeust) {
            // pass
        }
    }

    // Test explain query
    @Test
    public void testExplainQuery() throws InterruptedException {
        final WebTarget target = target().path("queryapi/queries");

        final FormDataMultiPart mp = new FormDataMultiPart();
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(),
                "select ID from " + testTable));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "explain"));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(),
                new GrillConf(), MediaType.APPLICATION_XML_TYPE));

        final QueryPlan plan = target.request().post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA_TYPE),
                QueryPlan.class);
        Assert.assertEquals(plan.getNumSels(), 1);
        Assert.assertEquals(plan.getTablesQueried().size(), 1);
        Assert.assertTrue(plan.getTablesQueried().get(0).equalsIgnoreCase(testTable));
        Assert.assertNull(plan.getPrepareHandle());
    }

    // post to preparedqueries
    // get all prepared queries
    // get a prepared query
    // update a prepared query
    // post to prepared query multiple times
    // delete a prepared query
    @Test
    public void testPrepareQuery() throws InterruptedException {
        final WebTarget target = target().path("queryapi/preparedqueries");

        final FormDataMultiPart mp = new FormDataMultiPart();
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(),
                "select ID from " + testTable));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "prepare"));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(),
                new GrillConf(), MediaType.APPLICATION_XML_TYPE));

        final QueryPrepareHandle pHandle = target.request()
                .post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA_TYPE), QueryPrepareHandle.class);

        // Get all prepared queries
        List<QueryPrepareHandle> allQueries = (List<QueryPrepareHandle>) target
                .queryParam("sessionid", grillSessionId).request().get(new GenericType<List<QueryPrepareHandle>>() {
                });
        Assert.assertTrue(allQueries.size() >= 1);
        Assert.assertTrue(allQueries.contains(pHandle));

        GrillPreparedQuery ctx = target.path(pHandle.toString()).queryParam("sessionid", grillSessionId).request()
                .get(GrillPreparedQuery.class);
        Assert.assertTrue(ctx.getUserQuery().equalsIgnoreCase("select ID from " + testTable));
        Assert.assertTrue(ctx.getDriverQuery().equalsIgnoreCase("select ID from " + testTable));
        Assert.assertEquals(ctx.getSelectedDriverClassName(),
                com.inmobi.grill.driver.hive.HiveDriver.class.getCanonicalName());
        Assert.assertNull(ctx.getConf().getProperties().get("my.property"));

        // Update conf for prepared query
        final FormDataMultiPart confpart = new FormDataMultiPart();
        GrillConf conf = new GrillConf();
        conf.addProperty("my.property", "myvalue");
        confpart.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        confpart.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(),
                conf, MediaType.APPLICATION_XML_TYPE));
        APIResult updateConf = target.path(pHandle.toString()).request()
                .put(Entity.entity(confpart, MediaType.MULTIPART_FORM_DATA_TYPE), APIResult.class);
        Assert.assertEquals(updateConf.getStatus(), APIResult.Status.SUCCEEDED);

        ctx = target.path(pHandle.toString()).queryParam("sessionid", grillSessionId).request()
                .get(GrillPreparedQuery.class);
        Assert.assertEquals(ctx.getConf().getProperties().get("my.property"), "myvalue");

        QueryHandle handle1 = target.path(pHandle.toString()).request()
                .post(Entity.entity(confpart, MediaType.MULTIPART_FORM_DATA_TYPE), QueryHandle.class);

        // do post once again
        QueryHandle handle2 = target.path(pHandle.toString()).request()
                .post(Entity.entity(confpart, MediaType.MULTIPART_FORM_DATA_TYPE), QueryHandle.class);
        Assert.assertNotEquals(handle1, handle2);

        GrillQuery ctx1 = target().path("queryapi/queries").path(handle1.toString())
                .queryParam("sessionid", grillSessionId).request().get(GrillQuery.class);
        // wait till the query finishes
        QueryStatus stat = ctx1.getStatus();
        while (!stat.isFinished()) {
            ctx1 = target().path("queryapi/queries").path(handle1.toString())
                    .queryParam("sessionid", grillSessionId).request().get(GrillQuery.class);
            stat = ctx1.getStatus();
            Thread.sleep(1000);
        }
        Assert.assertEquals(ctx1.getStatus().getStatus(), QueryStatus.Status.SUCCESSFUL);

        GrillQuery ctx2 = target().path("queryapi/queries").path(handle2.toString())
                .queryParam("sessionid", grillSessionId).request().get(GrillQuery.class);
        // wait till the query finishes
        stat = ctx2.getStatus();
        while (!stat.isFinished()) {
            ctx2 = target().path("queryapi/queries").path(handle1.toString())
                    .queryParam("sessionid", grillSessionId).request().get(GrillQuery.class);
            stat = ctx2.getStatus();
            Thread.sleep(1000);
        }
        Assert.assertEquals(ctx1.getStatus().getStatus(), QueryStatus.Status.SUCCESSFUL);

        // destroy prepared
        APIResult result = target.path(pHandle.toString()).queryParam("sessionid", grillSessionId).request()
                .delete(APIResult.class);
        Assert.assertEquals(result.getStatus(), APIResult.Status.SUCCEEDED);

        // Post on destroyed query
        Response response = target.path(pHandle.toString()).request()
                .post(Entity.entity(confpart, MediaType.MULTIPART_FORM_DATA_TYPE), Response.class);
        Assert.assertEquals(response.getStatus(), 404);
    }

    @Test
    public void testExplainAndPrepareQuery() throws InterruptedException {
        final WebTarget target = target().path("queryapi/preparedqueries");

        final FormDataMultiPart mp = new FormDataMultiPart();
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(),
                "select ID from " + testTable));
        mp.bodyPart(
                new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "explain_and_prepare"));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(),
                new GrillConf(), MediaType.APPLICATION_XML_TYPE));

        final QueryPlan plan = target.request().post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA_TYPE),
                QueryPlan.class);
        Assert.assertEquals(plan.getNumSels(), 1);
        Assert.assertEquals(plan.getTablesQueried().size(), 1);
        Assert.assertTrue(plan.getTablesQueried().get(0).equalsIgnoreCase(testTable));
        Assert.assertNotNull(plan.getPrepareHandle());

        GrillPreparedQuery ctx = target.path(plan.getPrepareHandle().toString())
                .queryParam("sessionid", grillSessionId).request().get(GrillPreparedQuery.class);
        Assert.assertTrue(ctx.getUserQuery().equalsIgnoreCase("select ID from " + testTable));
        Assert.assertTrue(ctx.getDriverQuery().equalsIgnoreCase("select ID from " + testTable));
        Assert.assertEquals(ctx.getSelectedDriverClassName(),
                com.inmobi.grill.driver.hive.HiveDriver.class.getCanonicalName());
        Assert.assertNull(ctx.getConf().getProperties().get("my.property"));

        // Update conf for prepared query
        final FormDataMultiPart confpart = new FormDataMultiPart();
        GrillConf conf = new GrillConf();
        conf.addProperty("my.property", "myvalue");
        confpart.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        confpart.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(),
                conf, MediaType.APPLICATION_XML_TYPE));
        APIResult updateConf = target.path(plan.getPrepareHandle().toString()).request()
                .put(Entity.entity(confpart, MediaType.MULTIPART_FORM_DATA_TYPE), APIResult.class);
        Assert.assertEquals(updateConf.getStatus(), APIResult.Status.SUCCEEDED);

        ctx = target.path(plan.getPrepareHandle().toString()).queryParam("sessionid", grillSessionId).request()
                .get(GrillPreparedQuery.class);
        Assert.assertEquals(ctx.getConf().getProperties().get("my.property"), "myvalue");

        QueryHandle handle1 = target.path(plan.getPrepareHandle().toString()).request()
                .post(Entity.entity(confpart, MediaType.MULTIPART_FORM_DATA_TYPE), QueryHandle.class);

        // do post once again
        QueryHandle handle2 = target.path(plan.getPrepareHandle().toString()).request()
                .post(Entity.entity(confpart, MediaType.MULTIPART_FORM_DATA_TYPE), QueryHandle.class);
        Assert.assertNotEquals(handle1, handle2);

        GrillQuery ctx1 = target().path("queryapi/queries").path(handle1.toString())
                .queryParam("sessionid", grillSessionId).request().get(GrillQuery.class);
        // wait till the query finishes
        QueryStatus stat = ctx1.getStatus();
        while (!stat.isFinished()) {
            ctx1 = target().path("queryapi/queries").path(handle1.toString())
                    .queryParam("sessionid", grillSessionId).request().get(GrillQuery.class);
            stat = ctx1.getStatus();
            Thread.sleep(1000);
        }
        Assert.assertEquals(ctx1.getStatus().getStatus(), QueryStatus.Status.SUCCESSFUL);

        GrillQuery ctx2 = target().path("queryapi/queries").path(handle2.toString())
                .queryParam("sessionid", grillSessionId).request().get(GrillQuery.class);
        // wait till the query finishes
        stat = ctx2.getStatus();
        while (!stat.isFinished()) {
            ctx2 = target().path("queryapi/queries").path(handle1.toString())
                    .queryParam("sessionid", grillSessionId).request().get(GrillQuery.class);
            stat = ctx2.getStatus();
            Thread.sleep(1000);
        }
        Assert.assertEquals(ctx1.getStatus().getStatus(), QueryStatus.Status.SUCCESSFUL);

        // destroy prepared
        APIResult result = target.path(plan.getPrepareHandle().toString()).queryParam("sessionid", grillSessionId)
                .request().delete(APIResult.class);
        Assert.assertEquals(result.getStatus(), APIResult.Status.SUCCEEDED);

        // Post on destroyed query
        Response response = target.path(plan.getPrepareHandle().toString()).request()
                .post(Entity.entity(confpart, MediaType.MULTIPART_FORM_DATA_TYPE), Response.class);
        Assert.assertEquals(response.getStatus(), 404);

    }

    // test with execute async post, get query, get results
    // test cancel query
    @Test
    public void testExecuteAsync() throws InterruptedException, IOException {
        // test post execute op
        final WebTarget target = target().path("queryapi/queries");

        long queuedQueries = metricsSvc.getQueuedQueries();
        long runningQueries = metricsSvc.getRunningQueries();

        final FormDataMultiPart mp = new FormDataMultiPart();
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(),
                "select ID, IDSTR from " + testTable));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "execute"));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(),
                new GrillConf(), MediaType.APPLICATION_XML_TYPE));
        final QueryHandle handle = target.request().post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA_TYPE),
                QueryHandle.class);

        Assert.assertNotNull(handle);

        // Get query
        GrillQuery ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                .get(GrillQuery.class);
        Assert.assertTrue(ctx.getStatus().getStatus().equals(Status.QUEUED)
                || ctx.getStatus().getStatus().equals(Status.LAUNCHED)
                || ctx.getStatus().getStatus().equals(Status.RUNNING));

        // wait till the query finishes
        QueryStatus stat = ctx.getStatus();
        while (!stat.isFinished()) {
            ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                    .get(GrillQuery.class);
            stat = ctx.getStatus();
            switch (stat.getStatus()) {
            case RUNNING:
                assertEquals(metricsSvc.getRunningQueries(), runningQueries + 1,
                        "Asserting queries for " + ctx.getQueryHandle());
                break;
            case QUEUED:
                assertEquals(metricsSvc.getQueuedQueries(), queuedQueries + 1);
                break;
            default: // nothing
            }
            Thread.sleep(1000);
        }
        assertTrue(ctx.getSubmissionTime() > 0);
        assertTrue(ctx.getLaunchTime() > 0);
        assertTrue(ctx.getDriverStartTime() > 0);
        assertTrue(ctx.getDriverFinishTime() > 0);
        assertTrue(ctx.getFinishTime() > 0);
        Assert.assertEquals(ctx.getStatus().getStatus(), QueryStatus.Status.SUCCESSFUL);

        validatePersistedResult(handle);

        // test cancel query
        final QueryHandle handle2 = target.request().post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA_TYPE),
                QueryHandle.class);

        Assert.assertNotNull(handle2);
        APIResult result = target.path(handle2.toString()).queryParam("sessionid", grillSessionId).request()
                .delete(APIResult.class);
        // cancel would fail query is already successful
        Assert.assertTrue(result.getStatus().equals(APIResult.Status.SUCCEEDED)
                || result.getStatus().equals(APIResult.Status.FAILED));

        GrillQuery ctx2 = target.path(handle2.toString()).queryParam("sessionid", grillSessionId).request()
                .get(GrillQuery.class);
        if (result.getStatus().equals(APIResult.Status.FAILED)) {
            Assert.assertTrue(ctx2.getStatus().getStatus() == QueryStatus.Status.SUCCESSFUL);
        } else {
            Assert.assertTrue(ctx2.getStatus().getStatus() == QueryStatus.Status.CANCELED);
        }
    }

    private void validatePersistedResult(QueryHandle handle) throws IOException {
        final WebTarget target = target().path("queryapi/queries");
        // fetch results
        validateResultSetMetadata(handle);

        String presultset = target.path(handle.toString()).path("resultset").queryParam("sessionid", grillSessionId)
                .request().get(String.class);
        System.out.println("PERSISTED RESULT:" + presultset);

        PersistentQueryResult resultset = target.path(handle.toString()).path("resultset")
                .queryParam("sessionid", grillSessionId).request().get(PersistentQueryResult.class);
        validatePersistentResult(resultset, handle);
    }

    private void validatePersistentResult(PersistentQueryResult resultset, QueryHandle handle) throws IOException {
        Assert.assertTrue(resultset.getPersistedURI().endsWith(handle.toString()));
        Path actualPath = new Path(resultset.getPersistedURI());
        FileSystem fs = actualPath.getFileSystem(new Configuration());
        List<String> actualRows = new ArrayList<String>();
        for (FileStatus fstat : fs.listStatus(actualPath)) {
            FSDataInputStream in = fs.open(fstat.getPath());
            BufferedReader br = null;
            try {
                br = new BufferedReader(new InputStreamReader(in));
                String line = "";

                while ((line = br.readLine()) != null) {
                    actualRows.add(line);
                }
            } finally {
                if (br != null) {
                    br.close();
                }
                if (in != null) {
                    in.close();
                }
            }
        }
        Assert.assertEquals(actualRows.get(0), "1one");
        Assert.assertEquals(actualRows.get(1), "\\Ntwo");
        Assert.assertEquals(actualRows.get(2), "3\\N");
        Assert.assertEquals(actualRows.get(3), "\\N\\N");
        Assert.assertEquals(actualRows.get(4), "5");
    }

    // test with execute async post, get query, get results
    // test cancel query
    @Test
    public void testExecuteAsyncInMemoryResult() throws InterruptedException, IOException {
        // test post execute op
        final WebTarget target = target().path("queryapi/queries");

        final FormDataMultiPart mp = new FormDataMultiPart();
        GrillConf conf = new GrillConf();
        conf.addProperty(GrillConfConstants.GRILL_PERSISTENT_RESULT_SET, "false");
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(),
                "select ID, IDSTR from " + testTable));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "execute"));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(), conf,
                MediaType.APPLICATION_XML_TYPE));
        final QueryHandle handle = target.request().post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA_TYPE),
                QueryHandle.class);

        Assert.assertNotNull(handle);

        // Get query
        GrillQuery ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                .get(GrillQuery.class);
        Assert.assertTrue(ctx.getStatus().getStatus().equals(Status.QUEUED)
                || ctx.getStatus().getStatus().equals(Status.LAUNCHED)
                || ctx.getStatus().getStatus().equals(Status.RUNNING));

        // wait till the query finishes
        QueryStatus stat = ctx.getStatus();
        while (!stat.isFinished()) {
            ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                    .get(GrillQuery.class);
            stat = ctx.getStatus();
            Thread.sleep(1000);
        }
        Assert.assertEquals(ctx.getStatus().getStatus(), QueryStatus.Status.SUCCESSFUL);

        // fetch results
        validateResultSetMetadata(handle);

        InMemoryQueryResult resultset = target.path(handle.toString()).path("resultset")
                .queryParam("sessionid", grillSessionId).request().get(InMemoryQueryResult.class);
        validateInmemoryResult(resultset);
    }

    @Test
    public void testExecuteAsyncTempTable() throws InterruptedException, IOException {
        // test post execute op
        final WebTarget target = target().path("queryapi/queries");

        final FormDataMultiPart mp = new FormDataMultiPart();
        GrillConf conf = new GrillConf();
        conf.addProperty(GrillConfConstants.GRILL_PERSISTENT_RESULT_SET, "false");
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(),
                "create table temp_output as select ID, IDSTR from " + testTable));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "execute"));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(), conf,
                MediaType.APPLICATION_XML_TYPE));
        final QueryHandle handle = target.request().post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA_TYPE),
                QueryHandle.class);

        Assert.assertNotNull(handle);

        // Get query
        GrillQuery ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                .get(GrillQuery.class);
        Assert.assertTrue(ctx.getStatus().getStatus().equals(Status.QUEUED)
                || ctx.getStatus().getStatus().equals(Status.LAUNCHED)
                || ctx.getStatus().getStatus().equals(Status.RUNNING));

        // wait till the query finishes
        QueryStatus stat = ctx.getStatus();
        while (!stat.isFinished()) {
            ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                    .get(GrillQuery.class);
            stat = ctx.getStatus();
            Thread.sleep(1000);
        }
        Assert.assertEquals(ctx.getStatus().getStatus(), QueryStatus.Status.SUCCESSFUL);

        String select = "SELECT * FROM temp_output";
        final FormDataMultiPart fetch = new FormDataMultiPart();
        fetch.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        fetch.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(), select));
        fetch.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "execute"));
        fetch.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(), conf,
                MediaType.APPLICATION_XML_TYPE));
        final QueryHandle handle2 = target.request().post(Entity.entity(fetch, MediaType.MULTIPART_FORM_DATA_TYPE),
                QueryHandle.class);

        Assert.assertNotNull(handle2);

        // Get query
        ctx = target.path(handle2.toString()).queryParam("sessionid", grillSessionId).request()
                .get(GrillQuery.class);

        // wait till the query finishes
        stat = ctx.getStatus();
        while (!stat.isFinished()) {
            ctx = target.path(handle2.toString()).queryParam("sessionid", grillSessionId).request()
                    .get(GrillQuery.class);
            stat = ctx.getStatus();
            Thread.sleep(1000);
        }
        Assert.assertEquals(ctx.getStatus().getStatus(), QueryStatus.Status.SUCCESSFUL);

        // fetch results
        validateResultSetMetadata(handle2, "temp_output.");

        InMemoryQueryResult resultset = target.path(handle2.toString()).path("resultset")
                .queryParam("sessionid", grillSessionId).request().get(InMemoryQueryResult.class);
        validateInmemoryResult(resultset);
    }

    private void validateResultSetMetadata(QueryHandle handle) {
        validateResultSetMetadata(handle, "");
    }

    private void validateResultSetMetadata(QueryHandle handle, String outputTablePfx) {
        final WebTarget target = target().path("queryapi/queries");

        QueryResultSetMetadata metadata = target.path(handle.toString()).path("resultsetmetadata")
                .queryParam("sessionid", grillSessionId).request().get(QueryResultSetMetadata.class);
        Assert.assertEquals(metadata.getColumns().size(), 2);
        assertTrue(
                metadata.getColumns().get(0).getName().toLowerCase().equals((outputTablePfx + "ID").toLowerCase())
                        || metadata.getColumns().get(0).getName().toLowerCase().equals("ID".toLowerCase()));
        assertEquals("INT".toLowerCase(), metadata.getColumns().get(0).getType().name().toLowerCase());
        assertTrue(metadata.getColumns().get(1).getName().toLowerCase()
                .equals((outputTablePfx + "IDSTR").toLowerCase())
                || metadata.getColumns().get(0).getName().toLowerCase().equals("IDSTR".toLowerCase()));
        assertEquals("STRING".toLowerCase(), metadata.getColumns().get(1).getType().name().toLowerCase());
    }

    private void validateInmemoryResult(InMemoryQueryResult resultset) {
        Assert.assertEquals(resultset.getRows().size(), 5);
        Assert.assertEquals(resultset.getRows().get(0).getValues().get(0), 1);
        Assert.assertEquals((String) resultset.getRows().get(0).getValues().get(1), "one");

        Assert.assertNull(resultset.getRows().get(1).getValues().get(0));
        Assert.assertEquals((String) resultset.getRows().get(1).getValues().get(1), "two");

        Assert.assertEquals(resultset.getRows().get(2).getValues().get(0), 3);
        Assert.assertNull(resultset.getRows().get(2).getValues().get(1));

        Assert.assertNull(resultset.getRows().get(3).getValues().get(0));
        Assert.assertNull(resultset.getRows().get(3).getValues().get(1));
        Assert.assertEquals(resultset.getRows().get(4).getValues().get(0), 5);
        Assert.assertEquals(resultset.getRows().get(4).getValues().get(1), "");
    }

    // test execute with timeout, fetch results
    // cancel the query with execute_with_timeout
    @Test
    public void testExecuteWithTimeoutQuery() throws IOException {
        final WebTarget target = target().path("queryapi/queries");

        final FormDataMultiPart mp = new FormDataMultiPart();
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(),
                "select ID, IDSTR from " + testTable));
        mp.bodyPart(
                new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "execute_with_timeout"));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(),
                new GrillConf(), MediaType.APPLICATION_XML_TYPE));

        QueryHandleWithResultSet result = target.request()
                .post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA_TYPE), QueryHandleWithResultSet.class);
        Assert.assertNotNull(result.getQueryHandle());
        Assert.assertNotNull(result.getResult());
        validatePersistentResult((PersistentQueryResult) result.getResult(), result.getQueryHandle());

        final FormDataMultiPart mp2 = new FormDataMultiPart();
        GrillConf conf = new GrillConf();
        conf.addProperty(GrillConfConstants.GRILL_PERSISTENT_RESULT_SET, "false");
        mp2.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        mp2.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(),
                "select ID, IDSTR from " + testTable));
        mp2.bodyPart(
                new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "execute_with_timeout"));
        mp2.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(), conf,
                MediaType.APPLICATION_XML_TYPE));

        result = target.request().post(Entity.entity(mp2, MediaType.MULTIPART_FORM_DATA_TYPE),
                QueryHandleWithResultSet.class);
        Assert.assertNotNull(result.getQueryHandle());
        Assert.assertNotNull(result.getResult());
        validateInmemoryResult((InMemoryQueryResult) result.getResult());

    }

    @Test
    public void testServerRestart() throws InterruptedException, IOException, GrillException {
        // test post execute op
        final WebTarget target = target().path("queryapi/queries");

        final FormDataMultiPart mp = new FormDataMultiPart();
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("sessionid").build(), grillSessionId,
                MediaType.APPLICATION_XML_TYPE));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("query").build(),
                "select ID, IDSTR from " + testTable));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("operation").build(), "execute"));
        mp.bodyPart(new FormDataBodyPart(FormDataContentDisposition.name("conf").fileName("conf").build(),
                new GrillConf(), MediaType.APPLICATION_XML_TYPE));
        final QueryHandle handle = target.request().post(Entity.entity(mp, MediaType.MULTIPART_FORM_DATA_TYPE),
                QueryHandle.class);

        Assert.assertNotNull(handle);

        GrillQuery ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                .get(GrillQuery.class);
        // wait till the query is launched
        QueryStatus stat = ctx.getStatus();
        while (stat.getStatus().equals(Status.QUEUED)) {
            ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                    .get(GrillQuery.class);
            stat = ctx.getStatus();
            Thread.sleep(100);
        }
        // Restart the server
        System.out.println("Restarting grill server!");
        restartGrillServer();
        queryService = (QueryExecutionServiceImpl) GrillServices.get().getService("query");
        grillSessionId = queryService.openSession("foo", "bar", new HashMap<String, String>());

        // wait till the query finishes
        ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                .get(GrillQuery.class);
        stat = ctx.getStatus();
        while (!stat.isFinished()) {
            System.out.println("Status:" + stat);
            ctx = target.path(handle.toString()).queryParam("sessionid", grillSessionId).request()
                    .get(GrillQuery.class);
            stat = ctx.getStatus();
            Thread.sleep(1000);
        }
        Assert.assertEquals(ctx.getStatus().getStatus(), QueryStatus.Status.SUCCESSFUL);
        validatePersistedResult(handle);
    }

    @Override
    protected int getTestPort() {
        return 8083;
    }
}