diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java index 5f0bd545b7c..cce6f14c7cc 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java @@ -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); diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerFactory.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerFactory.java index ac369fd3f08..43a67bda9fe 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerFactory.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerFactory.java @@ -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); } diff --git a/sql/src/main/java/org/apache/druid/sql/hook/DruidHookDispatcher.java b/sql/src/main/java/org/apache/druid/sql/hook/DruidHookDispatcher.java index cfc37523a02..688a5e9a3e4 100644 --- a/sql/src/main/java/org/apache/druid/sql/hook/DruidHookDispatcher.java +++ b/sql/src/main/java/org/apache/druid/sql/hook/DruidHookDispatcher.java @@ -37,17 +37,17 @@ public class DruidHookDispatcher Map, List>> 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); } - Closeable withHook(HookKey key, DruidHook hook) + public Closeable withHook(HookKey key, DruidHook hook) { register(key, hook); return new Closeable() @@ -61,7 +61,7 @@ public class DruidHookDispatcher } @SuppressWarnings({"rawtypes", "unchecked"}) - void dispatch(HookKey key, T object) + public void dispatch(HookKey key, T object) { List> hooks = GLOBAL.get(key); if (hooks != null) { diff --git a/sql/src/test/java/org/apache/druid/sql/SqlStatementTest.java b/sql/src/test/java/org/apache/druid/sql/SqlStatementTest.java index b332b6dabfa..1a7b6e3a7c5 100644 --- a/sql/src/test/java/org/apache/druid/sql/SqlStatementTest.java +++ b/sql/src/test/java/org/apache/druid/sql/SqlStatementTest.java @@ -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( diff --git a/sql/src/test/java/org/apache/druid/sql/avatica/DruidAvaticaHandlerTest.java b/sql/src/test/java/org/apache/druid/sql/avatica/DruidAvaticaHandlerTest.java index cccaab75b6f..518e225c491 100644 --- a/sql/src/test/java/org/apache/druid/sql/avatica/DruidAvaticaHandlerTest.java +++ b/sql/src/test/java/org/apache/druid/sql/avatica/DruidAvaticaHandlerTest.java @@ -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() ) ); } diff --git a/sql/src/test/java/org/apache/druid/sql/avatica/DruidStatementTest.java b/sql/src/test/java/org/apache/druid/sql/avatica/DruidStatementTest.java index 81cd1c69231..ed855f9c5f8 100644 --- a/sql/src/test/java/org/apache/druid/sql/avatica/DruidStatementTest.java +++ b/sql/src/test/java/org/apache/druid/sql/avatica/DruidStatementTest.java @@ -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), diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/SqlVectorizedExpressionSanityTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/SqlVectorizedExpressionSanityTest.java index 304400bc3d8..07f57a5deeb 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/SqlVectorizedExpressionSanityTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/SqlVectorizedExpressionSanityTest.java @@ -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() ); } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java index 4374f1e86e2..84c8a299024 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/SqlTestFramework.java @@ -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( diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index fa3c1edcea2..665d5873b46 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -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()