From 5f6290eb54e6ee6faf6193600c4459b9ac364d4f Mon Sep 17 00:00:00 2001
From: Zoltan Haindrich <kirk@rxd.hu>
Date: Tue, 30 Jul 2024 16:11:57 +0000
Subject: [PATCH] use updated hook class

---
 .../apache/druid/sql/avatica/MSQDruidMeta.java    | 15 ++++++++++-----
 .../druid/quidem/DruidQuidemCommandHandler.java   | 13 +++++++------
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/MSQDruidMeta.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/MSQDruidMeta.java
index 92b0ab4b915..bae85aefdef 100644
--- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/MSQDruidMeta.java
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/MSQDruidMeta.java
@@ -36,8 +36,9 @@ import org.apache.druid.segment.column.RowSignature;
 import org.apache.druid.server.security.AuthenticatorMapper;
 import org.apache.druid.sql.SqlStatementFactory;
 import org.apache.druid.sql.calcite.planner.DruidTypeSystem;
-import org.apache.druid.sql.calcite.run.DruidHook;
 import org.apache.druid.sql.calcite.table.RowSignatures;
+import org.apache.druid.sql.hook.DruidHook;
+import org.apache.druid.sql.hook.DruidHookDispatcher;
 
 import java.util.Collections;
 import java.util.List;
@@ -45,7 +46,8 @@ import java.util.List;
 public class MSQDruidMeta extends DruidMeta
 {
   protected final MSQTestOverlordServiceClient overlordClient;
-  private final ObjectMapper objectMapper;
+  protected final ObjectMapper objectMapper;
+  protected final DruidHookDispatcher hookDispatcher;
 
   @Inject
   public MSQDruidMeta(
@@ -54,11 +56,14 @@ public class MSQDruidMeta extends DruidMeta
       final ErrorHandler errorHandler,
       final AuthenticatorMapper authMapper,
       final MSQTestOverlordServiceClient overlordClient,
-      final ObjectMapper objectMapper)
+      final ObjectMapper objectMapper,
+      final DruidHookDispatcher hookDispatcher
+      )
   {
     super(sqlStatementFactory, config, errorHandler, authMapper);
     this.overlordClient = overlordClient;
     this.objectMapper = objectMapper;
+    this.hookDispatcher = hookDispatcher;
   }
 
   @Override
@@ -89,10 +94,10 @@ public class MSQDruidMeta extends DruidMeta
           .writerWithDefaultPrettyPrinter()
           .writeValueAsString(payload.getStages());
       str = str.replaceAll(taskId, "<taskId>");
-      DruidHook.dispatch(DruidHook.MSQ_PLAN, str);
+      hookDispatcher.dispatch(DruidHook.MSQ_PLAN, str);
     }
     catch (JsonProcessingException e) {
-      DruidHook.dispatch(DruidHook.MSQ_PLAN, "error happened during json serialization");
+      hookDispatcher.dispatch(DruidHook.MSQ_PLAN, "error happened during json serialization");
     }
 
     Signature signature = makeSignature(druidStatement, payload.getResults().getSignature());
diff --git a/sql/src/test/java/org/apache/druid/quidem/DruidQuidemCommandHandler.java b/sql/src/test/java/org/apache/druid/quidem/DruidQuidemCommandHandler.java
index 146958b6c58..699f8a3619d 100644
--- a/sql/src/test/java/org/apache/druid/quidem/DruidQuidemCommandHandler.java
+++ b/sql/src/test/java/org/apache/druid/quidem/DruidQuidemCommandHandler.java
@@ -121,6 +121,11 @@ public class DruidQuidemCommandHandler implements CommandHandler
       }
     }
 
+    protected final DruidHookDispatcher unwrapDruidHookDispatcher(Context x)
+    {
+      return DruidConnectionExtras.unwrapOrThrow(x.connection()).getDruidHookDispatcher();
+    }
+
     protected abstract void executeExplain(Context x) throws Exception;
   }
 
@@ -190,11 +195,6 @@ public class DruidQuidemCommandHandler implements CommandHandler
         x.echo(ImmutableList.of(str));
       }
     }
-
-    protected final DruidHookDispatcher unwrapDruidHookDispatcher(Context x)
-    {
-      return DruidConnectionExtras.unwrapOrThrow(x.connection()).getDruidHookDispatcher();
-    }
   }
 
   /**
@@ -213,8 +213,9 @@ public class DruidQuidemCommandHandler implements CommandHandler
     @Override
     protected final void executeExplain(Context x) throws IOException
     {
+      DruidHookDispatcher dhp = unwrapDruidHookDispatcher(x);
       List<String> logged = new ArrayList<>();
-      try (Closeable unhook = DruidHook.withHook(hook, (key, relNode) -> {
+      try (Closeable unhook = dhp.withHook(hook, (key, relNode) -> {
         logged.add(relNode);
       })) {
         executeQuery(x);