move more to non-static

kgyrtkirk · Jul 30, 2024 · 78b75d3 · 78b75d3
1 parent f6cc540
commit 78b75d3
Show file tree

Hide file tree

Showing 9 changed files with 39 additions and 13 deletions.
diff --git a/...re/multi-stage-query/src/test/java/org/apache/druid/msq/exec/ResultsContextSerdeTest.java b/...re/multi-stage-query/src/test/java/org/apache/druid/msq/exec/ResultsContextSerdeTest.java
@@ -43,6 +43,7 @@
 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 @@ NamedViewSchema.NAME, new NamedViewSchema(EasyMock.createMock(ViewSchema.class))
         "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
@@ -55,6 +55,7 @@
 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 LookupExtractor getLookup(final String lookupName)
 
     return lookupCache.getLookup(lookupName);
   }
+
+  public <T> void dispatchHook(HookKey<T> 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
@@ -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 PlannerFactory(
         druidSchemaName,
         calciteRuleManager,
         authorizerMapper,
-        authConfig
+        authConfig,
+        hookDispatcher
     );
-    this.hookDispatcher = hookDispatcher;
   }
 
   /**
@@ -117,7 +116,7 @@ public DruidPlanner createPlanner(
         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
@@ -26,6 +26,7 @@
 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 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 PlannerToolbox(
     this.calciteRuleManager = calciteRuleManager;
     this.authorizerMapper = authorizerMapper;
     this.authConfig = authConfig;
+    this.hookDispatcher = hookDispatcher;
   }
 
   public DruidOperatorTable operatorTable()
@@ -117,4 +121,9 @@ public AuthConfig getAuthConfig()
   {
     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
@@ -563,7 +563,8 @@ protected PlannerResult planWithDruidConvention() throws ValidationException
                  .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
@@ -62,6 +62,7 @@
 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 @@ NamedViewSchema.NAME, new NamedViewSchema(EasyMock.createMock(ViewSchema.class))
       "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 @@ void testExpression(
     }
 
     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
@@ -39,6 +39,7 @@
 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 @@ NamedViewSchema.NAME, new NamedViewSchema(EasyMock.createMock(ViewSchema.class))
         "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
@@ -50,6 +50,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.hook.DruidHookDispatcher;
 import org.easymock.EasyMock;
 import org.easymock.EasyMockExtension;
 import org.easymock.Mock;
@@ -192,7 +193,8 @@ public void testExtensionCalciteRule()
         "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 void testConfigurableBloat()
             "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
@@ -54,6 +54,7 @@
 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 @@ NamedViewSchema.NAME, new NamedViewSchema(EasyMock.createMock(ViewSchema.class))
       "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,