use of org.apache.druid.java.util.common.guava.LazySequence in project druid by druid-io.
the class SqlResourceTest method setUp.
@Before
public void setUp() throws Exception {
final QueryScheduler scheduler = new QueryScheduler(5, ManualQueryPrioritizationStrategy.INSTANCE, new HiLoQueryLaningStrategy(40), new ServerConfig()) {
@Override
public <T> Sequence<T> run(Query<?> query, Sequence<T> resultSequence) {
return super.run(query, new LazySequence<T>(() -> {
if (sleep) {
try {
// pretend to be a query that is waiting on results
Thread.sleep(500);
} catch (InterruptedException ignored) {
}
}
return resultSequence;
}));
}
};
executorService = MoreExecutors.listeningDecorator(Execs.multiThreaded(8, "test_sql_resource_%s"));
walker = CalciteTests.createMockWalker(conglomerate, temporaryFolder.newFolder(), scheduler);
final PlannerConfig plannerConfig = new PlannerConfig() {
@Override
public boolean shouldSerializeComplexValues() {
return false;
}
};
final DruidSchemaCatalog rootSchema = CalciteTests.createMockRootSchema(conglomerate, walker, plannerConfig, CalciteTests.TEST_AUTHORIZER_MAPPER);
final DruidOperatorTable operatorTable = CalciteTests.createOperatorTable();
final ExprMacroTable macroTable = CalciteTests.createExprMacroTable();
req = EasyMock.createStrictMock(HttpServletRequest.class);
EasyMock.expect(req.getRemoteAddr()).andReturn(null).once();
EasyMock.expect(req.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)).andReturn(CalciteTests.REGULAR_USER_AUTH_RESULT).anyTimes();
EasyMock.expect(req.getAttribute(AuthConfig.DRUID_ALLOW_UNSECURED_PATH)).andReturn(null).anyTimes();
EasyMock.expect(req.getAttribute(AuthConfig.DRUID_AUTHORIZATION_CHECKED)).andReturn(null).anyTimes();
EasyMock.expect(req.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)).andReturn(CalciteTests.REGULAR_USER_AUTH_RESULT).anyTimes();
req.setAttribute(AuthConfig.DRUID_AUTHORIZATION_CHECKED, true);
EasyMock.expectLastCall().anyTimes();
EasyMock.expect(req.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)).andReturn(CalciteTests.REGULAR_USER_AUTH_RESULT).anyTimes();
EasyMock.replay(req);
testRequestLogger = new TestRequestLogger();
final PlannerFactory plannerFactory = new PlannerFactory(rootSchema, CalciteTests.createMockQueryMakerFactory(walker, conglomerate), operatorTable, macroTable, plannerConfig, CalciteTests.TEST_AUTHORIZER_MAPPER, CalciteTests.getJsonMapper(), CalciteTests.DRUID_SCHEMA_NAME);
lifecycleManager = new SqlLifecycleManager() {
@Override
public void add(String sqlQueryId, SqlLifecycle lifecycle) {
super.add(sqlQueryId, lifecycle);
if (lifecycleAddLatch != null) {
lifecycleAddLatch.countDown();
}
}
};
final ServiceEmitter emitter = new NoopServiceEmitter();
sqlLifecycleFactory = new SqlLifecycleFactory(plannerFactory, emitter, testRequestLogger, scheduler) {
@Override
public SqlLifecycle factorize() {
return new TestSqlLifecycle(plannerFactory, emitter, testRequestLogger, scheduler, System.currentTimeMillis(), System.nanoTime(), validateAndAuthorizeLatchSupplier, planLatchSupplier, executeLatchSupplier, sequenceMapFnSupplier);
}
};
resource = new SqlResource(JSON_MAPPER, CalciteTests.TEST_AUTHORIZER_MAPPER, sqlLifecycleFactory, lifecycleManager, new ServerConfig());
}
use of org.apache.druid.java.util.common.guava.LazySequence in project druid by druid-io.
the class QueryResourceTest method createScheduledQueryResource.
private void createScheduledQueryResource(QueryScheduler scheduler, Collection<CountDownLatch> beforeScheduler, Collection<CountDownLatch> inScheduler) {
QuerySegmentWalker texasRanger = new QuerySegmentWalker() {
@Override
public <T> QueryRunner<T> getQueryRunnerForIntervals(Query<T> query, Iterable<Interval> intervals) {
return (queryPlus, responseContext) -> {
beforeScheduler.forEach(CountDownLatch::countDown);
return scheduler.run(scheduler.prioritizeAndLaneQuery(queryPlus, ImmutableSet.of()), new LazySequence<T>(() -> {
inScheduler.forEach(CountDownLatch::countDown);
try {
// pretend to be a query that is waiting on results
Thread.sleep(500);
} catch (InterruptedException ignored) {
}
// all that waiting for nothing :(
return Sequences.empty();
}));
};
}
@Override
public <T> QueryRunner<T> getQueryRunnerForSegments(Query<T> query, Iterable<SegmentDescriptor> specs) {
return getQueryRunnerForIntervals(null, null);
}
};
queryResource = new QueryResource(new QueryLifecycleFactory(WAREHOUSE, texasRanger, new DefaultGenericQueryMetricsFactory(), new NoopServiceEmitter(), testRequestLogger, new AuthConfig(), AuthTestUtils.TEST_AUTHORIZER_MAPPER, Suppliers.ofInstance(new DefaultQueryConfig(ImmutableMap.of()))), jsonMapper, smileMapper, scheduler, new AuthConfig(), null, ResponseContextConfig.newConfig(true), DRUID_NODE);
}
use of org.apache.druid.java.util.common.guava.LazySequence in project druid by druid-io.
the class TestClusterQuerySegmentWalker method getQueryRunnerForSegments.
@Override
public <T> QueryRunner<T> getQueryRunnerForSegments(final Query<T> query, final Iterable<SegmentDescriptor> specs) {
final QueryRunnerFactory<T, Query<T>> factory = conglomerate.findFactory(query);
if (factory == null) {
throw new ISE("Unknown query type[%s].", query.getClass());
}
final DataSourceAnalysis analysis = DataSourceAnalysis.forDataSource(query.getDataSource());
if (!analysis.isConcreteTableBased()) {
throw new ISE("Cannot handle datasource: %s", query.getDataSource());
}
final String dataSourceName = ((TableDataSource) analysis.getBaseDataSource()).getName();
final QueryToolChest<T, Query<T>> toolChest = factory.getToolchest();
// Make sure this query type can handle the subquery, if present.
if (analysis.isQuery() && !toolChest.canPerformSubquery(((QueryDataSource) analysis.getDataSource()).getQuery())) {
throw new ISE("Cannot handle subquery: %s", analysis.getDataSource());
}
final Function<SegmentReference, SegmentReference> segmentMapFn = joinableFactoryWrapper.createSegmentMapFn(analysis.getJoinBaseTableFilter().map(Filters::toFilter).orElse(null), analysis.getPreJoinableClauses(), new AtomicLong(), analysis.getBaseQuery().orElse(query));
final QueryRunner<T> baseRunner = new FinalizeResultsQueryRunner<>(toolChest.postMergeQueryDecoration(toolChest.mergeResults(toolChest.preMergeQueryDecoration(makeTableRunner(toolChest, factory, getSegmentsForTable(dataSourceName, specs), segmentMapFn)))), toolChest);
// to actually serve the queries
return (theQuery, responseContext) -> {
responseContext.initializeRemainingResponses();
responseContext.addRemainingResponse(theQuery.getQuery().getMostSpecificId(), 0);
if (scheduler != null) {
Set<SegmentServerSelector> segments = new HashSet<>();
specs.forEach(spec -> segments.add(new SegmentServerSelector(spec)));
return scheduler.run(scheduler.prioritizeAndLaneQuery(theQuery, segments), new LazySequence<>(() -> baseRunner.run(theQuery.withQuery(Queries.withSpecificSegments(theQuery.getQuery(), ImmutableList.copyOf(specs))), responseContext)));
} else {
return baseRunner.run(theQuery.withQuery(Queries.withSpecificSegments(theQuery.getQuery(), ImmutableList.copyOf(specs))), responseContext);
}
};
}
Aggregations