Apache
/
druid
mirror of https://github.com/apache/druid.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

use druidhookdispatcherr#1

This commit is contained in:
Zoltan Haindrich 2024-07-30 10:33:38 +00:00
parent ce667eeb5e
commit f6cc540368
9 changed files with 31 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java
View File

@ -186,6 +186,7 @@ import org.apache.druid.sql.calcite.util.SqlTestFramework.StandardComponentSuppl
import org.apache.druid.sql.calcite.util.TestDataBuilder;
import org.apache.druid.sql.calcite.view.InProcessViewManager;
import org.apache.druid.sql.guice.SqlBindings;
import org.apache.druid.sql.hook.DruidHookDispatcher;
import org.apache.druid.storage.StorageConfig;
import org.apache.druid.storage.StorageConnector;
import org.apache.druid.storage.StorageConnectorModule;
@ -571,7 +572,8 @@ public class MSQTestBase extends BaseCalciteQueryTest
new CalciteRulesManager(ImmutableSet.of()),
CalciteTests.createJoinableFactoryWrapper(),
catalogResolver,
new AuthConfig()
new AuthConfig(),
new DruidHookDispatcher()
);
sqlStatementFactory = CalciteTests.createSqlStatementFactory(engine, plannerFactory);

8
sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerFactory.java
View File

@ -50,6 +50,7 @@ import org.apache.druid.sql.calcite.run.SqlEngine;
import org.apache.druid.sql.calcite.schema.DruidSchemaCatalog;
import org.apache.druid.sql.calcite.schema.DruidSchemaName;
import org.apache.druid.sql.hook.DruidHook;
import org.apache.druid.sql.hook.DruidHookDispatcher;
import java.util.Map;
import java.util.Properties;
@ -65,6 +66,7 @@ public class PlannerFactory extends PlannerToolbox
.setConformance(DruidConformance.instance())
.setParserFactory(new DruidSqlParserImplFactory()) // Custom SQL parser factory
.build();
private final DruidHookDispatcher hookDispatcher;
@Inject
public PlannerFactory(
@ -78,7 +80,8 @@ public class PlannerFactory extends PlannerToolbox
final CalciteRulesManager calciteRuleManager,
final JoinableFactoryWrapper joinableFactoryWrapper,
final CatalogResolver catalog,
final AuthConfig authConfig
final AuthConfig authConfig,
final DruidHookDispatcher hookDispatcher
)
{
super(
@ -94,6 +97,7 @@ public class PlannerFactory extends PlannerToolbox
authorizerMapper,
authConfig
);
this.hookDispatcher = hookDispatcher;
}
/**
@ -113,7 +117,7 @@ public class PlannerFactory extends PlannerToolbox
queryContext,
hook
);
DruidHook.dispatch(DruidHook.SQL, sql);
hookDispatcher.dispatch(DruidHook.SQL, sql);
return new DruidPlanner(buildFrameworkConfig(context), context, engine, hook);
}

8
sql/src/main/java/org/apache/druid/sql/hook/DruidHookDispatcher.java
View File

@ -37,17 +37,17 @@ public class DruidHookDispatcher
Map<HookKey<?>, List<DruidHook<?>>> GLOBAL = new HashMap<>();
void register(HookKey<?> label, DruidHook<?> hook)
public void register(HookKey<?> label, DruidHook<?> hook)
{
GLOBAL.computeIfAbsent(label, k -> new ArrayList<>()).add(hook);
}
void unregister(HookKey<?> key, DruidHook<?> hook)
public void unregister(HookKey<?> key, DruidHook<?> hook)
{
GLOBAL.get(key).remove(hook);
}
<T> Closeable withHook(HookKey<T> key, DruidHook<T> hook)
public <T> Closeable withHook(HookKey<T> key, DruidHook<T> hook)
{
register(key, hook);
return new Closeable()
@ -61,7 +61,7 @@ public class DruidHookDispatcher
}
@SuppressWarnings({"rawtypes", "unchecked"})
<T> void dispatch(HookKey<T> key, T object)
public <T> void dispatch(HookKey<T> key, T object)
{
List<DruidHook<?>> hooks = GLOBAL.get(key);
if (hooks != null) {

4
sql/src/test/java/org/apache/druid/sql/SqlStatementTest.java
View File

@ -58,6 +58,7 @@ import org.apache.druid.sql.calcite.planner.PlannerFactory;
import org.apache.druid.sql.calcite.planner.PrepareResult;
import org.apache.druid.sql.calcite.schema.DruidSchemaCatalog;
import org.apache.druid.sql.calcite.util.CalciteTests;
import org.apache.druid.sql.hook.DruidHookDispatcher;
import org.apache.druid.sql.http.SqlQuery;
import org.easymock.EasyMock;
import org.hamcrest.MatcherAssert;
@ -158,7 +159,8 @@ public class SqlStatementTest
new CalciteRulesManager(ImmutableSet.of()),
joinableFactoryWrapper,
CatalogResolver.NULL_RESOLVER,
new AuthConfig()
new AuthConfig(),
new DruidHookDispatcher()
);
this.sqlStatementFactory = new SqlStatementFactory(

4
sql/src/test/java/org/apache/druid/sql/avatica/DruidAvaticaHandlerTest.java
View File

@ -89,6 +89,7 @@ import org.apache.druid.sql.calcite.schema.NamedSchema;
import org.apache.druid.sql.calcite.util.CalciteTestBase;
import org.apache.druid.sql.calcite.util.CalciteTests;
import org.apache.druid.sql.guice.SqlModule;
import org.apache.druid.sql.hook.DruidHookDispatcher;
import org.eclipse.jetty.server.Server;
import org.joda.time.DateTime;
import org.joda.time.DateTimeZone;
@ -1048,7 +1049,8 @@ public class DruidAvaticaHandlerTest extends CalciteTestBase
new CalciteRulesManager(ImmutableSet.of()),
CalciteTests.createJoinableFactoryWrapper(),
CatalogResolver.NULL_RESOLVER,
new AuthConfig()
new AuthConfig(),
new DruidHookDispatcher()
)
);
}

4
sql/src/test/java/org/apache/druid/sql/avatica/DruidStatementTest.java
View File

@ -47,6 +47,7 @@ import org.apache.druid.sql.calcite.planner.PlannerFactory;
import org.apache.druid.sql.calcite.schema.DruidSchemaCatalog;
import org.apache.druid.sql.calcite.util.CalciteTestBase;
import org.apache.druid.sql.calcite.util.CalciteTests;
import org.apache.druid.sql.hook.DruidHookDispatcher;
import org.junit.Assert;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.AfterEach;
@ -112,7 +113,8 @@ public class DruidStatementTest extends CalciteTestBase
new CalciteRulesManager(ImmutableSet.of()),
joinableFactoryWrapper,
CatalogResolver.NULL_RESOLVER,
new AuthConfig()
new AuthConfig(),
new DruidHookDispatcher()
);
this.sqlStatementFactory = CalciteTests.createSqlStatementFactory(
CalciteTests.createMockSqlEngine(walker, conglomerate),

4
sql/src/test/java/org/apache/druid/sql/calcite/SqlVectorizedExpressionSanityTest.java
View File

@ -49,6 +49,7 @@ import org.apache.druid.sql.calcite.planner.PlannerResult;
import org.apache.druid.sql.calcite.run.SqlEngine;
import org.apache.druid.sql.calcite.schema.DruidSchemaCatalog;
import org.apache.druid.sql.calcite.util.CalciteTests;
import org.apache.druid.sql.hook.DruidHookDispatcher;
import org.apache.druid.testing.InitializedNullHandlingTest;
import org.apache.druid.timeline.DataSegment;
import org.apache.druid.timeline.partition.LinearShardSpec;
@ -152,7 +153,8 @@ public class SqlVectorizedExpressionSanityTest extends InitializedNullHandlingTe
new CalciteRulesManager(ImmutableSet.of()),
joinableFactoryWrapper,
CatalogResolver.NULL_RESOLVER,
new AuthConfig()
new AuthConfig(),
new DruidHookDispatcher()
);
}

4
sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java
View File

@ -68,6 +68,7 @@ import org.apache.druid.sql.calcite.schema.NoopDruidSchemaManager;
import org.apache.druid.sql.calcite.view.DruidViewMacroFactory;
import org.apache.druid.sql.calcite.view.InProcessViewManager;
import org.apache.druid.sql.calcite.view.ViewManager;
import org.apache.druid.sql.hook.DruidHookDispatcher;
import org.apache.druid.timeline.DataSegment;
import javax.inject.Named;
@ -502,7 +503,8 @@ public class SqlTestFramework
new CalciteRulesManager(componentSupplier.extensionCalciteRules()),
framework.injector.getInstance(JoinableFactoryWrapper.class),
framework.builder.catalogResolver,
authConfig != null ? authConfig : new AuthConfig()
authConfig != null ? authConfig : new AuthConfig(),
new DruidHookDispatcher()
);
componentSupplier.finalizePlanner(this);
this.statementFactory = QueryFrameworkUtils.createSqlStatementFactory(

4
sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java
View File

@ -102,6 +102,7 @@ import org.apache.druid.sql.calcite.run.NativeSqlEngine;
import org.apache.druid.sql.calcite.schema.DruidSchemaCatalog;
import org.apache.druid.sql.calcite.util.CalciteTestBase;
import org.apache.druid.sql.calcite.util.CalciteTests;
import org.apache.druid.sql.hook.DruidHookDispatcher;
import org.hamcrest.CoreMatchers;
import org.junit.Assert;
import org.junit.jupiter.api.AfterAll;
@ -259,7 +260,8 @@ public class SqlResourceTest extends CalciteTestBase
new CalciteRulesManager(ImmutableSet.of()),
CalciteTests.createJoinableFactoryWrapper(),
CatalogResolver.NULL_RESOLVER,
new AuthConfig()
new AuthConfig(),
new DruidHookDispatcher()
);
lifecycleManager = new SqlLifecycleManager()