From 78b75d3e8ee776e990094c4e4e5125816bbb6fb9 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Tue, 30 Jul 2024 10:42:41 +0000 Subject: [PATCH] move more to non-static --- .../druid/msq/exec/ResultsContextSerdeTest.java | 4 +++- .../druid/sql/calcite/planner/PlannerContext.java | 6 ++++++ .../druid/sql/calcite/planner/PlannerFactory.java | 7 +++---- .../druid/sql/calcite/planner/PlannerToolbox.java | 11 ++++++++++- .../druid/sql/calcite/planner/QueryHandler.java | 3 ++- .../sql/calcite/expression/ExpressionTestHelper.java | 6 ++++-- .../calcite/external/ExternalTableScanRuleTest.java | 4 +++- .../sql/calcite/planner/CalcitePlannerModuleTest.java | 7 +++++-- .../sql/calcite/planner/DruidRexExecutorTest.java | 4 +++- 9 files changed, 39 insertions(+), 13 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/ResultsContextSerdeTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/ResultsContextSerdeTest.java index b6023392576..72f2c817622 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/ResultsContextSerdeTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/ResultsContextSerdeTest.java @@ -43,6 +43,7 @@ import org.apache.druid.sql.calcite.schema.NamedDruidSchema; import org.apache.druid.sql.calcite.schema.NamedViewSchema; import org.apache.druid.sql.calcite.schema.ViewSchema; import org.apache.druid.sql.calcite.util.CalciteTests; +import org.apache.druid.sql.hook.DruidHookDispatcher; import org.easymock.EasyMock; import org.junit.Assert; import org.junit.Before; @@ -75,7 +76,8 @@ public class ResultsContextSerdeTest "druid", new CalciteRulesManager(ImmutableSet.of()), CalciteTests.TEST_AUTHORIZER_MAPPER, - AuthConfig.newBuilder().build() + AuthConfig.newBuilder().build(), + new DruidHookDispatcher() ); final NativeSqlEngine engine = CalciteTests.createMockSqlEngine( EasyMock.createMock(QuerySegmentWalker.class), diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java index b7e2de3e66b..21ad4c23663 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java @@ -55,6 +55,7 @@ import org.apache.druid.sql.calcite.rule.ReverseLookupRule; import org.apache.druid.sql.calcite.run.EngineFeature; import org.apache.druid.sql.calcite.run.QueryMaker; import org.apache.druid.sql.calcite.run.SqlEngine; +import org.apache.druid.sql.hook.DruidHook.HookKey; import org.joda.time.DateTime; import org.joda.time.DateTimeZone; import org.joda.time.Interval; @@ -668,4 +669,9 @@ public class PlannerContext return lookupCache.getLookup(lookupName); } + + public void dispatchHook(HookKey key, T object) + { + plannerToolbox.getHookDispatcher().dispatch(key, object); + } } 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 43a67bda9fe..514952972d3 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 @@ -66,7 +66,6 @@ public class PlannerFactory extends PlannerToolbox .setConformance(DruidConformance.instance()) .setParserFactory(new DruidSqlParserImplFactory()) // Custom SQL parser factory .build(); - private final DruidHookDispatcher hookDispatcher; @Inject public PlannerFactory( @@ -95,9 +94,9 @@ public class PlannerFactory extends PlannerToolbox druidSchemaName, calciteRuleManager, authorizerMapper, - authConfig + authConfig, + hookDispatcher ); - this.hookDispatcher = hookDispatcher; } /** @@ -117,7 +116,7 @@ public class PlannerFactory extends PlannerToolbox queryContext, hook ); - hookDispatcher.dispatch(DruidHook.SQL, sql); + context.dispatchHook(DruidHook.SQL, sql); return new DruidPlanner(buildFrameworkConfig(context), context, engine, hook); } diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerToolbox.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerToolbox.java index f0009826118..d8e3a25f458 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerToolbox.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerToolbox.java @@ -26,6 +26,7 @@ import org.apache.druid.segment.join.JoinableFactoryWrapper; import org.apache.druid.server.security.AuthConfig; import org.apache.druid.server.security.AuthorizerMapper; import org.apache.druid.sql.calcite.schema.DruidSchemaCatalog; +import org.apache.druid.sql.hook.DruidHookDispatcher; public class PlannerToolbox { @@ -40,6 +41,7 @@ public class PlannerToolbox protected final CalciteRulesManager calciteRuleManager; protected final AuthorizerMapper authorizerMapper; protected final AuthConfig authConfig; + protected final DruidHookDispatcher hookDispatcher; public PlannerToolbox( final DruidOperatorTable operatorTable, @@ -52,7 +54,8 @@ public class PlannerToolbox final String druidSchemaName, final CalciteRulesManager calciteRuleManager, final AuthorizerMapper authorizerMapper, - final AuthConfig authConfig + final AuthConfig authConfig, + final DruidHookDispatcher hookDispatcher ) { this.operatorTable = operatorTable; @@ -66,6 +69,7 @@ public class PlannerToolbox this.calciteRuleManager = calciteRuleManager; this.authorizerMapper = authorizerMapper; this.authConfig = authConfig; + this.hookDispatcher = hookDispatcher; } public DruidOperatorTable operatorTable() @@ -117,4 +121,9 @@ public class PlannerToolbox { return authConfig; } + + public DruidHookDispatcher getHookDispatcher() + { + return hookDispatcher; + } } diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/QueryHandler.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/QueryHandler.java index ed21e929f1f..b8f25dd5dcf 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/QueryHandler.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/QueryHandler.java @@ -563,7 +563,8 @@ public abstract class QueryHandler extends SqlStatementHandler.BaseStatementHand .plus(DruidLogicalConvention.instance()), newRoot ); - DruidHook.dispatch(DruidHook.DRUID_PLAN, newRoot); + + plannerContext.dispatchHook(DruidHook.DRUID_PLAN, newRoot); DruidQueryGenerator generator = new DruidQueryGenerator(plannerContext, (DruidLogicalNode) newRoot, rexBuilder); DruidQuery baseQuery = generator.buildQuery(); diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/expression/ExpressionTestHelper.java b/sql/src/test/java/org/apache/druid/sql/calcite/expression/ExpressionTestHelper.java index 6155678c845..f4ada1f1c17 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/expression/ExpressionTestHelper.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/expression/ExpressionTestHelper.java @@ -62,6 +62,7 @@ import org.apache.druid.sql.calcite.schema.ViewSchema; import org.apache.druid.sql.calcite.table.RowSignatures; 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.easymock.EasyMock; import org.joda.time.DateTime; import org.joda.time.DateTimeZone; @@ -97,7 +98,8 @@ public class ExpressionTestHelper "druid", new CalciteRulesManager(ImmutableSet.of()), CalciteTests.TEST_AUTHORIZER_MAPPER, - AuthConfig.newBuilder().build() + AuthConfig.newBuilder().build(), + new DruidHookDispatcher() ); public static final PlannerContext PLANNER_CONTEXT = PlannerContext.create( PLANNER_TOOLBOX, @@ -336,7 +338,7 @@ public class ExpressionTestHelper } ExprEval result = PLANNER_CONTEXT.parseExpression(expression.getExpression()) - + .eval(expressionBindings); Assert.assertEquals("Result for: " + rexNode, expectedResult, result.value()); diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/external/ExternalTableScanRuleTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/external/ExternalTableScanRuleTest.java index 64a626f682c..fe26ba9db8c 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/external/ExternalTableScanRuleTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/external/ExternalTableScanRuleTest.java @@ -39,6 +39,7 @@ import org.apache.druid.sql.calcite.schema.NamedDruidSchema; import org.apache.druid.sql.calcite.schema.NamedViewSchema; import org.apache.druid.sql.calcite.schema.ViewSchema; import org.apache.druid.sql.calcite.util.CalciteTests; +import org.apache.druid.sql.hook.DruidHookDispatcher; import org.easymock.EasyMock; import org.junit.Assert; import org.junit.Test; @@ -71,7 +72,8 @@ public class ExternalTableScanRuleTest "druid", new CalciteRulesManager(ImmutableSet.of()), CalciteTests.TEST_AUTHORIZER_MAPPER, - AuthConfig.newBuilder().build() + AuthConfig.newBuilder().build(), + new DruidHookDispatcher() ); final PlannerContext plannerContext = PlannerContext.create( toolbox, diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java index 43d47733277..89df405b7f1 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java @@ -50,6 +50,7 @@ import org.apache.druid.sql.calcite.schema.DruidSchemaName; 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.hook.DruidHookDispatcher; import org.easymock.EasyMock; import org.easymock.EasyMockExtension; import org.easymock.Mock; @@ -192,7 +193,8 @@ public class CalcitePlannerModuleTest extends CalciteTestBase "druid", new CalciteRulesManager(ImmutableSet.of()), CalciteTests.TEST_AUTHORIZER_MAPPER, - AuthConfig.newBuilder().build() + AuthConfig.newBuilder().build(), + new DruidHookDispatcher() ); PlannerContext context = PlannerContext.create( @@ -224,7 +226,8 @@ public class CalcitePlannerModuleTest extends CalciteTestBase "druid", new CalciteRulesManager(ImmutableSet.of()), CalciteTests.TEST_AUTHORIZER_MAPPER, - AuthConfig.newBuilder().build() + AuthConfig.newBuilder().build(), + new DruidHookDispatcher() ); PlannerContext contextWithBloat = PlannerContext.create( diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/planner/DruidRexExecutorTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/planner/DruidRexExecutorTest.java index 205f3379d71..072bdd6b6e8 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/planner/DruidRexExecutorTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/planner/DruidRexExecutorTest.java @@ -54,6 +54,7 @@ import org.apache.druid.sql.calcite.schema.ViewSchema; import org.apache.druid.sql.calcite.table.RowSignatures; 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.apache.druid.testing.InitializedNullHandlingTest; import org.easymock.EasyMock; import org.junit.Assert; @@ -100,7 +101,8 @@ public class DruidRexExecutorTest extends InitializedNullHandlingTest "druid", new CalciteRulesManager(ImmutableSet.of()), CalciteTests.TEST_AUTHORIZER_MAPPER, - AuthConfig.newBuilder().build() + AuthConfig.newBuilder().build(), + new DruidHookDispatcher() ); private static final PlannerContext PLANNER_CONTEXT = PlannerContext.create( PLANNER_TOOLBOX,