make engine load via injector for msqdrill

This commit is contained in:
Zoltan Haindrich 2024-07-17 16:04:14 +00:00
parent 0eaf4c61b9
commit 51d465df6d
2 changed files with 69 additions and 5 deletions

View File

@ -21,16 +21,21 @@ package org.apache.druid.msq.exec;
import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Supplier; import com.google.common.base.Supplier;
import com.google.common.base.Suppliers;
import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.inject.Binder;
import com.google.inject.Injector; import com.google.inject.Injector;
import com.google.inject.Module; import com.google.inject.Module;
import com.google.inject.Provides; import com.google.inject.Provides;
import com.google.inject.TypeLiteral;
import org.apache.druid.guice.DruidInjectorBuilder; import org.apache.druid.guice.DruidInjectorBuilder;
import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.LazySingleton;
import org.apache.druid.guice.annotations.NativeQuery;
import org.apache.druid.initialization.ServerInjectorBuilderTest.TestDruidModule; import org.apache.druid.initialization.ServerInjectorBuilderTest.TestDruidModule;
import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.java.util.emitter.service.ServiceEmitter;
import org.apache.druid.msq.exec.MSQDrillWindowQueryTest.DrillWindowQueryMSQComponentSupplier; import org.apache.druid.msq.exec.MSQDrillWindowQueryTest.DrillWindowQueryMSQComponentSupplier;
import org.apache.druid.msq.guice.MultiStageQuery;
import org.apache.druid.msq.sql.MSQTaskSqlEngine; import org.apache.druid.msq.sql.MSQTaskSqlEngine;
import org.apache.druid.msq.test.CalciteMSQTestsHelper; import org.apache.druid.msq.test.CalciteMSQTestsHelper;
import org.apache.druid.msq.test.ExtractResultsFactory; import org.apache.druid.msq.test.ExtractResultsFactory;
@ -41,7 +46,13 @@ import org.apache.druid.query.DefaultQueryConfig;
import org.apache.druid.query.groupby.TestGroupByBuffers; import org.apache.druid.query.groupby.TestGroupByBuffers;
import org.apache.druid.server.QueryLifecycleFactory; import org.apache.druid.server.QueryLifecycleFactory;
import org.apache.druid.server.QueryScheduler; import org.apache.druid.server.QueryScheduler;
import org.apache.druid.server.QuerySchedulerProvider;
import org.apache.druid.server.log.RequestLogger; import org.apache.druid.server.log.RequestLogger;
import org.apache.druid.server.log.TestRequestLogger;
import org.apache.druid.server.metrics.NoopServiceEmitter;
import org.apache.druid.server.security.AuthenticatorMapper;
import org.apache.druid.server.security.AuthorizerMapper;
import org.apache.druid.server.security.Escalator;
import org.apache.druid.sql.SqlLifecycleManager; import org.apache.druid.sql.SqlLifecycleManager;
import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.SqlStatementFactory;
import org.apache.druid.sql.SqlToolbox; import org.apache.druid.sql.SqlToolbox;
@ -49,8 +60,13 @@ import org.apache.druid.sql.calcite.DrillWindowQueryTest;
import org.apache.druid.sql.calcite.QueryTestBuilder; import org.apache.druid.sql.calcite.QueryTestBuilder;
import org.apache.druid.sql.calcite.SqlTestFrameworkConfig; import org.apache.druid.sql.calcite.SqlTestFrameworkConfig;
import org.apache.druid.sql.calcite.TempDirProducer; import org.apache.druid.sql.calcite.TempDirProducer;
import org.apache.druid.sql.calcite.planner.CalciteRulesManager;
import org.apache.druid.sql.calcite.planner.CatalogResolver;
import org.apache.druid.sql.calcite.planner.PlannerFactory; import org.apache.druid.sql.calcite.planner.PlannerFactory;
import org.apache.druid.sql.calcite.run.SqlEngine; import org.apache.druid.sql.calcite.run.SqlEngine;
import org.apache.druid.sql.calcite.schema.DruidSchemaName;
import org.apache.druid.sql.calcite.util.CalciteTests;
import org.apache.druid.sql.guice.SqlModule;
@SqlTestFrameworkConfig.ComponentSupplier(DrillWindowQueryMSQComponentSupplier.class) @SqlTestFrameworkConfig.ComponentSupplier(DrillWindowQueryMSQComponentSupplier.class)
public class MSQDrillWindowQueryTest extends DrillWindowQueryTest public class MSQDrillWindowQueryTest extends DrillWindowQueryTest
@ -67,12 +83,39 @@ public class MSQDrillWindowQueryTest extends DrillWindowQueryTest
{ {
super.configureGuice(builder); super.configureGuice(builder);
builder.addModules(CalciteMSQTestsHelper.fetchModules(tempDirProducer::newTempFolder, TestGroupByBuffers.createDefault()).toArray(new Module[0])); builder.addModules(CalciteMSQTestsHelper.fetchModules(tempDirProducer::newTempFolder, TestGroupByBuffers.createDefault()).toArray(new Module[0]));
builder.addModule(new TestSqlModule());
builder.addModule(new TestMSQSqlModule()); builder.addModule(new TestMSQSqlModule());
} }
static class TestSqlModule extends TestDruidModule
{
@Override
public void configure(Binder binder)
{
binder.install(new SqlModule.SqlStatementFactoryModule());
binder.bind(String.class)
.annotatedWith(DruidSchemaName.class)
.toInstance(CalciteTests.DRUID_SCHEMA_NAME);
binder.bind(new TypeLiteral<Supplier<DefaultQueryConfig>>()
{
}).toInstance(Suppliers.ofInstance(new DefaultQueryConfig(ImmutableMap.of())));
binder.bind(CalciteRulesManager.class).toInstance(new CalciteRulesManager(ImmutableSet.of()));
TestRequestLogger testRequestLogger = new TestRequestLogger();
binder.bind(RequestLogger.class).toInstance(testRequestLogger);
binder.bind(CatalogResolver.class).toInstance(CatalogResolver.NULL_RESOLVER);
binder.bind(ServiceEmitter.class).to(NoopServiceEmitter.class);
binder.bind(QueryScheduler.class)
.toProvider(QuerySchedulerProvider.class)
.in(LazySingleton.class);
binder.bind(AuthenticatorMapper.class).toInstance(CalciteTests.TEST_AUTHENTICATOR_MAPPER);
binder.bind(AuthorizerMapper.class).toInstance(CalciteTests.TEST_AUTHORIZER_MAPPER);
binder.bind(Escalator.class).toInstance(CalciteTests.TEST_AUTHENTICATOR_ESCALATOR);
}
}
static class TestMSQSqlModule extends TestDruidModule { static class TestMSQSqlModule extends TestDruidModule {
@Provides // @Provides
@LazySingleton @LazySingleton
public SqlToolbox makeSqlToolbox( public SqlToolbox makeSqlToolbox(
final PlannerFactory plannerFactory, final PlannerFactory plannerFactory,
@ -95,7 +138,7 @@ public class MSQDrillWindowQueryTest extends DrillWindowQueryTest
} }
@Provides @Provides
@NativeQuery @MultiStageQuery
@LazySingleton @LazySingleton
public SqlStatementFactory makeNativeSqlStatementFactory( public SqlStatementFactory makeNativeSqlStatementFactory(
final MSQTaskSqlEngine sqlEngine, final MSQTaskSqlEngine sqlEngine,
@ -140,7 +183,27 @@ public class MSQDrillWindowQueryTest extends DrillWindowQueryTest
Injector injector Injector injector
) )
{ {
if(true) {
return injector.getInstance(MSQTaskSqlEngine.class); return injector.getInstance(MSQTaskSqlEngine.class);
} else {
final WorkerMemoryParameters workerMemoryParameters =
WorkerMemoryParameters.createInstance(
WorkerMemoryParameters.PROCESSING_MINIMUM_BYTES * 50,
2,
10,
2,
0,
0
);
final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient(
queryJsonMapper,
injector,
new MSQTestTaskActionClient(queryJsonMapper, injector),
workerMemoryParameters,
ImmutableList.of()
);
return new MSQTaskSqlEngine(indexingServiceClient, queryJsonMapper);
}
} }
} }

View File

@ -69,6 +69,7 @@ import org.apache.druid.sql.calcite.util.SqlTestFramework;
import org.apache.druid.sql.calcite.util.SqlTestFramework.Builder; import org.apache.druid.sql.calcite.util.SqlTestFramework.Builder;
import org.apache.druid.sql.calcite.util.SqlTestFramework.PlannerComponentSupplier; import org.apache.druid.sql.calcite.util.SqlTestFramework.PlannerComponentSupplier;
import org.apache.druid.sql.calcite.util.SqlTestFramework.QueryComponentSupplier; import org.apache.druid.sql.calcite.util.SqlTestFramework.QueryComponentSupplier;
import org.apache.druid.sql.guice.SqlModule;
import org.apache.http.client.utils.URIBuilder; import org.apache.http.client.utils.URIBuilder;
import org.eclipse.jetty.server.Server; import org.eclipse.jetty.server.Server;
@ -258,7 +259,7 @@ public class DruidAvaticaTestDriver implements Driver
binder.bind(QueryScheduler.class) binder.bind(QueryScheduler.class)
.toProvider(QuerySchedulerProvider.class) .toProvider(QuerySchedulerProvider.class)
.in(LazySingleton.class); .in(LazySingleton.class);
// binder.install(new SqlModule.SqlStatementFactoryModule()); binder.install(new SqlModule.SqlStatementFactoryModule());
binder.bind(new TypeLiteral<Supplier<DefaultQueryConfig>>() binder.bind(new TypeLiteral<Supplier<DefaultQueryConfig>>()
{ {
}).toInstance(Suppliers.ofInstance(new DefaultQueryConfig(ImmutableMap.of()))); }).toInstance(Suppliers.ofInstance(new DefaultQueryConfig(ImmutableMap.of())));