integration-tests/src/test/java/org/apache/druid/tests/query/ITWikipediaQueryTest.java - druid - Git at Google

 /*
  * 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.druid.tests.query;

 import com.fasterxml.jackson.jaxrs.smile.SmileMediaTypes;
 import com.google.common.collect.ImmutableMap;
 import com.google.inject.Inject;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.java.util.http.client.response.StatusResponseHolder;
 import org.apache.druid.query.Druids;
 import org.apache.druid.query.QueryCapacityExceededException;
 import org.apache.druid.query.aggregation.CountAggregatorFactory;
 import org.apache.druid.testing.IntegrationTestingConfig;
 import org.apache.druid.testing.clients.CoordinatorResourceTestClient;
 import org.apache.druid.testing.clients.QueryResourceTestClient;
 import org.apache.druid.testing.guice.DruidTestModuleFactory;
 import org.apache.druid.testing.utils.ITRetryUtil;
 import org.apache.druid.testing.utils.TestQueryHelper;
 import org.apache.druid.tests.TestNGGroup;
 import org.jboss.netty.handler.codec.http.HttpResponseStatus;
 import org.testng.Assert;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Guice;
 import org.testng.annotations.Test;

 import javax.ws.rs.core.MediaType;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.UUID;
 import java.util.concurrent.Future;

 @Test(groups = TestNGGroup.QUERY)
 @Guice(moduleFactory = DruidTestModuleFactory.class)
 public class ITWikipediaQueryTest
 {
   private static final Logger LOG = new Logger(ITWikipediaQueryTest.class);

   public static final String WIKIPEDIA_DATA_SOURCE = "wikipedia_editstream";
   private static final String WIKI_LOOKUP = "wiki-simple";
   private static final String WIKIPEDIA_QUERIES_RESOURCE = "/queries/wikipedia_editstream_queries.json";
   private static final String WIKIPEDIA_LOOKUP_RESOURCE = "/queries/wiki-lookup-config.json";

   @Inject
   private CoordinatorResourceTestClient coordinatorClient;
   @Inject
   private TestQueryHelper queryHelper;
   @Inject
   private QueryResourceTestClient queryClient;
   @Inject
   private IntegrationTestingConfig config;

   @BeforeMethod
   public void before() throws Exception
   {
     // ensure that wikipedia segments are loaded completely
     ITRetryUtil.retryUntilTrue(
         () -> coordinatorClient.areSegmentsLoaded(WIKIPEDIA_DATA_SOURCE), "wikipedia segment load"
     );
     if (!coordinatorClient.areLookupsLoaded(WIKI_LOOKUP)) {
       coordinatorClient.initializeLookups(WIKIPEDIA_LOOKUP_RESOURCE);
       ITRetryUtil.retryUntilTrue(
           () -> coordinatorClient.areLookupsLoaded(WIKI_LOOKUP), "wikipedia lookup load"
       );
     }
   }

   /**
    * A combination of request Content-Type and Accept HTTP header
    * The first is Content-Type which can not be null while the 2nd is Accept which could be null
    * <p>
    * When Accept is null, its value defaults to value of Content-Type
    */
   @DataProvider
   public static Object[][] encodingCombination()
   {
     return new Object[][]{
         {MediaType.APPLICATION_JSON, null},
         {MediaType.APPLICATION_JSON, MediaType.APPLICATION_JSON},
         {MediaType.APPLICATION_JSON, SmileMediaTypes.APPLICATION_JACKSON_SMILE},
         {SmileMediaTypes.APPLICATION_JACKSON_SMILE, null},
         {SmileMediaTypes.APPLICATION_JACKSON_SMILE, MediaType.APPLICATION_JSON},
         {SmileMediaTypes.APPLICATION_JACKSON_SMILE, SmileMediaTypes.APPLICATION_JACKSON_SMILE},
         };
   }

   @Test(dataProvider = "encodingCombination")
   public void testWikipediaQueriesFromFile(String contentType, String accept)
       throws Exception
   {
     // run tests on a new query helper
     TestQueryHelper queryHelper = this.queryHelper.withEncoding(contentType, accept);

     queryHelper.testQueriesFromFile(WIKIPEDIA_QUERIES_RESOURCE);
   }

   @Test
   public void testQueryLaningLaneIsLimited() throws Exception
   {
     ITRetryUtil.retryUntil(
         () -> {
           // the broker is configured with a manually defined query lane, 'one' with limit 1
           //  -Ddruid.query.scheduler.laning.type=manual
           //  -Ddruid.query.scheduler.laning.lanes.one=1
           // by issuing 50 queries, at least 1 of them will succeed on 'one', and at least 1 of them will overlap enough to
           // get limited.
           // It's possible but unlikely that these queries execute in a way that none of them overlap, so we
           // retry this test a few times to compensate for this.
           final int numQueries = 50;
           List<Future<StatusResponseHolder>> futures = new ArrayList<>(numQueries);
           for (int i = 0; i < numQueries; i++) {
             futures.add(
                 queryClient.queryAsync(
                     queryHelper.getQueryURL(config.getBrokerUrl()),
                     getQueryBuilder().build()
                 )
             );
           }

           int success = 0;
           int limited = 0;

           for (Future<StatusResponseHolder> future : futures) {
             StatusResponseHolder status = future.get();
             if (status.getStatus().getCode() == QueryCapacityExceededException.STATUS_CODE) {
               limited++;
               Assert.assertTrue(status.getContent().contains(QueryCapacityExceededException.makeLaneErrorMessage("one", 1)));
             } else if (status.getStatus().getCode() == HttpResponseStatus.OK.getCode()) {
               success++;
             }
           }

           try {
             Assert.assertTrue(success > 0);
             Assert.assertTrue(limited > 0);
             return true;
           }
           catch (AssertionError ae) {
             LOG.error(ae, "Got assertion error in testQueryLaningLaneIsLimited");
             return false;
           }
         },
         true,
         5000,
         3,
         "testQueryLaningLaneIsLimited"
     );

     // test another to make sure we can still issue one query at a time
     StatusResponseHolder followUp = queryClient.queryAsync(
         queryHelper.getQueryURL(config.getBrokerUrl()),
         getQueryBuilder().build()
     ).get();

     Assert.assertEquals(followUp.getStatus().getCode(), HttpResponseStatus.OK.getCode());

     StatusResponseHolder andAnother = queryClient.queryAsync(
         queryHelper.getQueryURL(config.getBrokerUrl()),
         getQueryBuilder().build()
     ).get();

     Assert.assertEquals(andAnother.getStatus().getCode(), HttpResponseStatus.OK.getCode());
   }

   @Test
   public void testQueryLaningWithNoLane() throws Exception
   {
     // the broker is configured with a manually defined query lane, 'one' with limit 1
     //  -Ddruid.query.scheduler.laning.type=manual
     //  -Ddruid.query.scheduler.laning.lanes.one=1
     // these queries will not belong to the lane so none of them should be limited
     final int numQueries = 50;
     List<Future<StatusResponseHolder>> futures = new ArrayList<>(numQueries);
     for (int i = 0; i < numQueries; i++) {
       futures.add(
           queryClient.queryAsync(
               queryHelper.getQueryURL(config.getBrokerUrl()),
               getQueryBuilder().context(ImmutableMap.of("queryId", UUID.randomUUID().toString())).build()
           )
       );
     }

     int success = 0;
     int limited = 0;

     for (Future<StatusResponseHolder> future : futures) {
       StatusResponseHolder status = future.get();
       if (status.getStatus().getCode() == QueryCapacityExceededException.STATUS_CODE) {
         limited++;
       } else if (status.getStatus().getCode() == HttpResponseStatus.OK.getCode()) {
         success++;
       }
     }

     Assert.assertTrue(success > 0);
     Assert.assertEquals(limited, 0);

   }

   private Druids.TimeseriesQueryBuilder getQueryBuilder()
   {
     return Druids.newTimeseriesQueryBuilder()
                  .dataSource("wikipedia_editstream")
                  .aggregators(new CountAggregatorFactory("chocula"))
                  .intervals("2013-01-01T00:00:00.000/2013-01-08T00:00:00.000")
                  .context(ImmutableMap.of("lane", "one", "queryId", UUID.randomUUID().toString()));
   }
 }
	/*
	* 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.druid.tests.query;

	import com.fasterxml.jackson.jaxrs.smile.SmileMediaTypes;
	import com.google.common.collect.ImmutableMap;
	import com.google.inject.Inject;
	import org.apache.druid.java.util.common.logger.Logger;
	import org.apache.druid.java.util.http.client.response.StatusResponseHolder;
	import org.apache.druid.query.Druids;
	import org.apache.druid.query.QueryCapacityExceededException;
	import org.apache.druid.query.aggregation.CountAggregatorFactory;
	import org.apache.druid.testing.IntegrationTestingConfig;
	import org.apache.druid.testing.clients.CoordinatorResourceTestClient;
	import org.apache.druid.testing.clients.QueryResourceTestClient;
	import org.apache.druid.testing.guice.DruidTestModuleFactory;
	import org.apache.druid.testing.utils.ITRetryUtil;
	import org.apache.druid.testing.utils.TestQueryHelper;
	import org.apache.druid.tests.TestNGGroup;
	import org.jboss.netty.handler.codec.http.HttpResponseStatus;
	import org.testng.Assert;
	import org.testng.annotations.BeforeMethod;
	import org.testng.annotations.DataProvider;
	import org.testng.annotations.Guice;
	import org.testng.annotations.Test;

	import javax.ws.rs.core.MediaType;
	import java.util.ArrayList;
	import java.util.List;
	import java.util.UUID;
	import java.util.concurrent.Future;

	@Test(groups = TestNGGroup.QUERY)
	@Guice(moduleFactory = DruidTestModuleFactory.class)
	public class ITWikipediaQueryTest
	{
	private static final Logger LOG = new Logger(ITWikipediaQueryTest.class);

	public static final String WIKIPEDIA_DATA_SOURCE = "wikipedia_editstream";
	private static final String WIKI_LOOKUP = "wiki-simple";
	private static final String WIKIPEDIA_QUERIES_RESOURCE = "/queries/wikipedia_editstream_queries.json";
	private static final String WIKIPEDIA_LOOKUP_RESOURCE = "/queries/wiki-lookup-config.json";

	@Inject
	private CoordinatorResourceTestClient coordinatorClient;
	@Inject
	private TestQueryHelper queryHelper;
	@Inject
	private QueryResourceTestClient queryClient;
	@Inject
	private IntegrationTestingConfig config;

	@BeforeMethod
	public void before() throws Exception
	{
	// ensure that wikipedia segments are loaded completely
	ITRetryUtil.retryUntilTrue(
	() -> coordinatorClient.areSegmentsLoaded(WIKIPEDIA_DATA_SOURCE), "wikipedia segment load"
	);
	if (!coordinatorClient.areLookupsLoaded(WIKI_LOOKUP)) {
	coordinatorClient.initializeLookups(WIKIPEDIA_LOOKUP_RESOURCE);
	ITRetryUtil.retryUntilTrue(
	() -> coordinatorClient.areLookupsLoaded(WIKI_LOOKUP), "wikipedia lookup load"
	);
	}
	}

	/**
	* A combination of request Content-Type and Accept HTTP header
	* The first is Content-Type which can not be null while the 2nd is Accept which could be null
	* <p>
	* When Accept is null, its value defaults to value of Content-Type
	*/
	@DataProvider
	public static Object[][] encodingCombination()
	{
	return new Object[][]{
	{MediaType.APPLICATION_JSON, null},
	{MediaType.APPLICATION_JSON, MediaType.APPLICATION_JSON},
	{MediaType.APPLICATION_JSON, SmileMediaTypes.APPLICATION_JACKSON_SMILE},
	{SmileMediaTypes.APPLICATION_JACKSON_SMILE, null},
	{SmileMediaTypes.APPLICATION_JACKSON_SMILE, MediaType.APPLICATION_JSON},
	{SmileMediaTypes.APPLICATION_JACKSON_SMILE, SmileMediaTypes.APPLICATION_JACKSON_SMILE},
	};
	}

	@Test(dataProvider = "encodingCombination")
	public void testWikipediaQueriesFromFile(String contentType, String accept)
	throws Exception
	{
	// run tests on a new query helper
	TestQueryHelper queryHelper = this.queryHelper.withEncoding(contentType, accept);

	queryHelper.testQueriesFromFile(WIKIPEDIA_QUERIES_RESOURCE);
	}

	@Test
	public void testQueryLaningLaneIsLimited() throws Exception
	{
	ITRetryUtil.retryUntil(
	() -> {
	// the broker is configured with a manually defined query lane, 'one' with limit 1
	// -Ddruid.query.scheduler.laning.type=manual
	// -Ddruid.query.scheduler.laning.lanes.one=1
	// by issuing 50 queries, at least 1 of them will succeed on 'one', and at least 1 of them will overlap enough to
	// get limited.
	// It's possible but unlikely that these queries execute in a way that none of them overlap, so we
	// retry this test a few times to compensate for this.
	final int numQueries = 50;
	List<Future<StatusResponseHolder>> futures = new ArrayList<>(numQueries);
	for (int i = 0; i < numQueries; i++) {
	futures.add(
	queryClient.queryAsync(
	queryHelper.getQueryURL(config.getBrokerUrl()),
	getQueryBuilder().build()
	)
	);
	}

	int success = 0;
	int limited = 0;

	for (Future<StatusResponseHolder> future : futures) {
	StatusResponseHolder status = future.get();
	if (status.getStatus().getCode() == QueryCapacityExceededException.STATUS_CODE) {
	limited++;
	Assert.assertTrue(status.getContent().contains(QueryCapacityExceededException.makeLaneErrorMessage("one", 1)));
	} else if (status.getStatus().getCode() == HttpResponseStatus.OK.getCode()) {
	success++;
	}
	}

	try {
	Assert.assertTrue(success > 0);
	Assert.assertTrue(limited > 0);
	return true;
	}
	catch (AssertionError ae) {
	LOG.error(ae, "Got assertion error in testQueryLaningLaneIsLimited");
	return false;
	}
	},
	true,
	5000,
	3,
	"testQueryLaningLaneIsLimited"
	);

	// test another to make sure we can still issue one query at a time
	StatusResponseHolder followUp = queryClient.queryAsync(
	queryHelper.getQueryURL(config.getBrokerUrl()),
	getQueryBuilder().build()
	).get();

	Assert.assertEquals(followUp.getStatus().getCode(), HttpResponseStatus.OK.getCode());

	StatusResponseHolder andAnother = queryClient.queryAsync(
	queryHelper.getQueryURL(config.getBrokerUrl()),
	getQueryBuilder().build()
	).get();

	Assert.assertEquals(andAnother.getStatus().getCode(), HttpResponseStatus.OK.getCode());
	}

	@Test
	public void testQueryLaningWithNoLane() throws Exception
	{
	// the broker is configured with a manually defined query lane, 'one' with limit 1
	// -Ddruid.query.scheduler.laning.type=manual
	// -Ddruid.query.scheduler.laning.lanes.one=1
	// these queries will not belong to the lane so none of them should be limited
	final int numQueries = 50;
	List<Future<StatusResponseHolder>> futures = new ArrayList<>(numQueries);
	for (int i = 0; i < numQueries; i++) {
	futures.add(
	queryClient.queryAsync(
	queryHelper.getQueryURL(config.getBrokerUrl()),
	getQueryBuilder().context(ImmutableMap.of("queryId", UUID.randomUUID().toString())).build()
	)
	);
	}

	int success = 0;
	int limited = 0;

	for (Future<StatusResponseHolder> future : futures) {
	StatusResponseHolder status = future.get();
	if (status.getStatus().getCode() == QueryCapacityExceededException.STATUS_CODE) {
	limited++;
	} else if (status.getStatus().getCode() == HttpResponseStatus.OK.getCode()) {
	success++;
	}
	}

	Assert.assertTrue(success > 0);
	Assert.assertEquals(limited, 0);

	}

	private Druids.TimeseriesQueryBuilder getQueryBuilder()
	{
	return Druids.newTimeseriesQueryBuilder()
	.dataSource("wikipedia_editstream")
	.aggregators(new CountAggregatorFactory("chocula"))
	.intervals("2013-01-01T00:00:00.000/2013-01-08T00:00:00.000")
	.context(ImmutableMap.of("lane", "one", "queryId", UUID.randomUUID().toString()));
	}
	}