use of org.apache.druid.query.DefaultQueryConfig in project druid by druid-io.
the class QueryResourceTest method testSecuredCancelQuery.
@Test(timeout = 60_000L)
public void testSecuredCancelQuery() throws Exception {
final CountDownLatch waitForCancellationLatch = new CountDownLatch(1);
final CountDownLatch waitFinishLatch = new CountDownLatch(2);
final CountDownLatch startAwaitLatch = new CountDownLatch(1);
final CountDownLatch cancelledCountDownLatch = new CountDownLatch(1);
EasyMock.expect(testServletRequest.getAttribute(AuthConfig.DRUID_AUTHORIZATION_CHECKED)).andReturn(null).anyTimes();
EasyMock.expect(testServletRequest.getAttribute(AuthConfig.DRUID_ALLOW_UNSECURED_PATH)).andReturn(null).anyTimes();
EasyMock.expect(testServletRequest.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)).andReturn(AUTHENTICATION_RESULT).anyTimes();
testServletRequest.setAttribute(AuthConfig.DRUID_AUTHORIZATION_CHECKED, true);
EasyMock.expectLastCall().times(1);
EasyMock.replay(testServletRequest);
AuthorizerMapper authMapper = new AuthorizerMapper(null) {
@Override
public Authorizer getAuthorizer(String name) {
return new Authorizer() {
@Override
public Access authorize(AuthenticationResult authenticationResult, Resource resource, Action action) {
// WRITE corresponds to cancellation of query
if (action.equals(Action.READ)) {
try {
// Countdown startAwaitLatch as we want query cancellation to happen
// after we enter isAuthorized method so that we can handle the
// InterruptedException here because of query cancellation
startAwaitLatch.countDown();
waitForCancellationLatch.await();
} catch (InterruptedException e) {
// When the query is cancelled the control will reach here,
// countdown the latch and rethrow the exception so that error response is returned for the query
cancelledCountDownLatch.countDown();
throw new RuntimeException(e);
}
return new Access(true);
} else {
return new Access(true);
}
}
};
}
};
queryResource = new QueryResource(new QueryLifecycleFactory(WAREHOUSE, TEST_SEGMENT_WALKER, new DefaultGenericQueryMetricsFactory(), new NoopServiceEmitter(), testRequestLogger, new AuthConfig(), authMapper, Suppliers.ofInstance(new DefaultQueryConfig(ImmutableMap.of()))), jsonMapper, smileMapper, queryScheduler, new AuthConfig(), authMapper, ResponseContextConfig.newConfig(true), DRUID_NODE);
final String queryString = "{\"queryType\":\"timeBoundary\", \"dataSource\":\"allow\"," + "\"context\":{\"queryId\":\"id_1\"}}";
ObjectMapper mapper = new DefaultObjectMapper();
Query<?> query = mapper.readValue(queryString, Query.class);
ListenableFuture<?> future = MoreExecutors.listeningDecorator(Execs.singleThreaded("test_query_resource_%s")).submit(new Runnable() {
@Override
public void run() {
try {
Response response = queryResource.doPost(new ByteArrayInputStream(queryString.getBytes(StandardCharsets.UTF_8)), null, testServletRequest);
Assert.assertEquals(Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(), response.getStatus());
} catch (IOException e) {
throw new RuntimeException(e);
}
waitFinishLatch.countDown();
}
});
queryScheduler.registerQueryFuture(query, future);
startAwaitLatch.await();
Executors.newSingleThreadExecutor().submit(new Runnable() {
@Override
public void run() {
Response response = queryResource.cancelQuery("id_1", testServletRequest);
Assert.assertEquals(Response.Status.ACCEPTED.getStatusCode(), response.getStatus());
waitForCancellationLatch.countDown();
waitFinishLatch.countDown();
}
});
waitFinishLatch.await();
cancelledCountDownLatch.await();
}
use of org.apache.druid.query.DefaultQueryConfig in project druid by druid-io.
the class QueryResourceTest method testDenySecuredCancelQuery.
@Test(timeout = 60_000L)
public void testDenySecuredCancelQuery() throws Exception {
final CountDownLatch waitForCancellationLatch = new CountDownLatch(1);
final CountDownLatch waitFinishLatch = new CountDownLatch(2);
final CountDownLatch startAwaitLatch = new CountDownLatch(1);
EasyMock.expect(testServletRequest.getAttribute(AuthConfig.DRUID_AUTHORIZATION_CHECKED)).andReturn(null).anyTimes();
EasyMock.expect(testServletRequest.getAttribute(AuthConfig.DRUID_ALLOW_UNSECURED_PATH)).andReturn(null).anyTimes();
EasyMock.expect(testServletRequest.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)).andReturn(AUTHENTICATION_RESULT).anyTimes();
testServletRequest.setAttribute(AuthConfig.DRUID_AUTHORIZATION_CHECKED, true);
EasyMock.expectLastCall().times(1);
testServletRequest.setAttribute(AuthConfig.DRUID_AUTHORIZATION_CHECKED, false);
EasyMock.expectLastCall().times(1);
EasyMock.replay(testServletRequest);
AuthorizerMapper authMapper = new AuthorizerMapper(null) {
@Override
public Authorizer getAuthorizer(String name) {
return new Authorizer() {
@Override
public Access authorize(AuthenticationResult authenticationResult, Resource resource, Action action) {
// WRITE corresponds to cancellation of query
if (action.equals(Action.READ)) {
try {
waitForCancellationLatch.await();
} catch (InterruptedException e) {
throw new RuntimeException(e);
}
return new Access(true);
} else {
// Deny access to cancel the query
return new Access(false);
}
}
};
}
};
queryResource = new QueryResource(new QueryLifecycleFactory(WAREHOUSE, TEST_SEGMENT_WALKER, new DefaultGenericQueryMetricsFactory(), new NoopServiceEmitter(), testRequestLogger, new AuthConfig(), authMapper, Suppliers.ofInstance(new DefaultQueryConfig(ImmutableMap.of()))), jsonMapper, smileMapper, queryScheduler, new AuthConfig(), authMapper, ResponseContextConfig.newConfig(true), DRUID_NODE);
final String queryString = "{\"queryType\":\"timeBoundary\", \"dataSource\":\"allow\"," + "\"context\":{\"queryId\":\"id_1\"}}";
ObjectMapper mapper = new DefaultObjectMapper();
Query<?> query = mapper.readValue(queryString, Query.class);
ListenableFuture<?> future = MoreExecutors.listeningDecorator(Execs.singleThreaded("test_query_resource_%s")).submit(new Runnable() {
@Override
public void run() {
try {
startAwaitLatch.countDown();
Response response = queryResource.doPost(new ByteArrayInputStream(queryString.getBytes(StandardCharsets.UTF_8)), null, testServletRequest);
Assert.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus());
} catch (IOException e) {
throw new RuntimeException(e);
}
waitFinishLatch.countDown();
}
});
queryScheduler.registerQueryFuture(query, future);
startAwaitLatch.await();
Executors.newSingleThreadExecutor().submit(new Runnable() {
@Override
public void run() {
try {
queryResource.cancelQuery("id_1", testServletRequest);
} catch (ForbiddenException e) {
waitForCancellationLatch.countDown();
waitFinishLatch.countDown();
}
}
});
waitFinishLatch.await();
}
use of org.apache.druid.query.DefaultQueryConfig in project druid by druid-io.
the class QueryResourceTest method testQueryTimeoutException.
@Test
public void testQueryTimeoutException() throws Exception {
final QuerySegmentWalker timeoutSegmentWalker = new QuerySegmentWalker() {
@Override
public <T> QueryRunner<T> getQueryRunnerForIntervals(Query<T> query, Iterable<Interval> intervals) {
throw new QueryTimeoutException();
}
@Override
public <T> QueryRunner<T> getQueryRunnerForSegments(Query<T> query, Iterable<SegmentDescriptor> specs) {
return getQueryRunnerForIntervals(null, null);
}
};
final QueryResource timeoutQueryResource = new QueryResource(new QueryLifecycleFactory(WAREHOUSE, timeoutSegmentWalker, new DefaultGenericQueryMetricsFactory(), new NoopServiceEmitter(), testRequestLogger, new AuthConfig(), AuthTestUtils.TEST_AUTHORIZER_MAPPER, Suppliers.ofInstance(new DefaultQueryConfig(ImmutableMap.of()))), jsonMapper, jsonMapper, queryScheduler, new AuthConfig(), null, ResponseContextConfig.newConfig(true), DRUID_NODE);
expectPermissiveHappyPathAuth();
Response response = timeoutQueryResource.doPost(new ByteArrayInputStream(SIMPLE_TIMESERIES_QUERY.getBytes(StandardCharsets.UTF_8)), null, /*pretty*/
testServletRequest);
Assert.assertNotNull(response);
Assert.assertEquals(QueryTimeoutException.STATUS_CODE, response.getStatus());
QueryTimeoutException ex;
try {
ex = jsonMapper.readValue((byte[]) response.getEntity(), QueryTimeoutException.class);
} catch (IOException e) {
throw new RuntimeException(e);
}
Assert.assertEquals("Query Timed Out!", ex.getMessage());
Assert.assertEquals(QueryTimeoutException.ERROR_CODE, ex.getErrorCode());
Assert.assertEquals(1, timeoutQueryResource.getTimedOutQueryCount());
}
use of org.apache.druid.query.DefaultQueryConfig in project druid by druid-io.
the class QueryResourceTest method testGoodQueryWithQueryConfigDoesNotOverrideQueryContext.
@Test
public void testGoodQueryWithQueryConfigDoesNotOverrideQueryContext() throws IOException {
String overrideConfigKey = "priority";
String overrideConfigValue = "678";
DefaultQueryConfig overrideConfig = new DefaultQueryConfig(ImmutableMap.of(overrideConfigKey, overrideConfigValue));
queryResource = new QueryResource(new QueryLifecycleFactory(WAREHOUSE, TEST_SEGMENT_WALKER, new DefaultGenericQueryMetricsFactory(), new NoopServiceEmitter(), testRequestLogger, new AuthConfig(), AuthTestUtils.TEST_AUTHORIZER_MAPPER, Suppliers.ofInstance(overrideConfig)), jsonMapper, smileMapper, queryScheduler, new AuthConfig(), null, ResponseContextConfig.newConfig(true), DRUID_NODE);
expectPermissiveHappyPathAuth();
Response response = queryResource.doPost(// SIMPLE_TIMESERIES_QUERY_LOW_PRIORITY context has overrideConfigKey with value of -1
new ByteArrayInputStream(SIMPLE_TIMESERIES_QUERY_LOW_PRIORITY.getBytes(StandardCharsets.UTF_8)), null, /*pretty*/
testServletRequest);
Assert.assertNotNull(response);
final ByteArrayOutputStream baos = new ByteArrayOutputStream();
((StreamingOutput) response.getEntity()).write(baos);
final List<Result<TimeBoundaryResultValue>> responses = jsonMapper.readValue(baos.toByteArray(), new TypeReference<List<Result<TimeBoundaryResultValue>>>() {
});
Assert.assertNotNull(response);
Assert.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus());
Assert.assertEquals(0, responses.size());
Assert.assertEquals(1, testRequestLogger.getNativeQuerylogs().size());
Assert.assertNotNull(testRequestLogger.getNativeQuerylogs().get(0).getQuery());
Assert.assertNotNull(testRequestLogger.getNativeQuerylogs().get(0).getQuery().getContext());
Assert.assertTrue(testRequestLogger.getNativeQuerylogs().get(0).getQuery().getContext().containsKey(overrideConfigKey));
Assert.assertEquals(-1, testRequestLogger.getNativeQuerylogs().get(0).getQuery().getContext().get(overrideConfigKey));
}
use of org.apache.druid.query.DefaultQueryConfig in project druid by druid-io.
the class QueryResourceTest method testGoodQueryWithQueryConfigOverrideDefault.
@Test
public void testGoodQueryWithQueryConfigOverrideDefault() throws IOException {
String overrideConfigKey = "priority";
String overrideConfigValue = "678";
DefaultQueryConfig overrideConfig = new DefaultQueryConfig(ImmutableMap.of(overrideConfigKey, overrideConfigValue));
queryResource = new QueryResource(new QueryLifecycleFactory(WAREHOUSE, TEST_SEGMENT_WALKER, new DefaultGenericQueryMetricsFactory(), new NoopServiceEmitter(), testRequestLogger, new AuthConfig(), AuthTestUtils.TEST_AUTHORIZER_MAPPER, Suppliers.ofInstance(overrideConfig)), jsonMapper, smileMapper, queryScheduler, new AuthConfig(), null, ResponseContextConfig.newConfig(true), DRUID_NODE);
expectPermissiveHappyPathAuth();
Response response = queryResource.doPost(new ByteArrayInputStream(SIMPLE_TIMESERIES_QUERY.getBytes(StandardCharsets.UTF_8)), null, /*pretty*/
testServletRequest);
Assert.assertNotNull(response);
final ByteArrayOutputStream baos = new ByteArrayOutputStream();
((StreamingOutput) response.getEntity()).write(baos);
final List<Result<TimeBoundaryResultValue>> responses = jsonMapper.readValue(baos.toByteArray(), new TypeReference<List<Result<TimeBoundaryResultValue>>>() {
});
Assert.assertNotNull(response);
Assert.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus());
Assert.assertEquals(0, responses.size());
Assert.assertEquals(1, testRequestLogger.getNativeQuerylogs().size());
Assert.assertNotNull(testRequestLogger.getNativeQuerylogs().get(0).getQuery());
Assert.assertNotNull(testRequestLogger.getNativeQuerylogs().get(0).getQuery().getContext());
Assert.assertTrue(testRequestLogger.getNativeQuerylogs().get(0).getQuery().getContext().containsKey(overrideConfigKey));
Assert.assertEquals(overrideConfigValue, testRequestLogger.getNativeQuerylogs().get(0).getQuery().getContext().get(overrideConfigKey));
}
Aggregations