From 9ac26e3a89ee237a38046a7d95949df82b855d4f Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Tue, 30 Jul 2024 12:29:36 +0000 Subject: [PATCH] wire-in hookdispatcher thru connection/etc --- .../druid/quidem/QuidemCaptureResource.java | 6 ++++- .../apache/druid/quidem/QuidemRecorder.java | 9 ++++--- .../calcite/planner/CalciteRulesManager.java | 3 ++- .../sql/calcite/planner/QueryHandler.java | 4 ++-- .../org/apache/druid/sql/hook/DruidHook.java | 12 +++++----- .../druid/sql/hook/DruidHookDispatcher.java | 2 ++ .../druid/quidem/DruidAvaticaTestDriver.java | 5 ++-- .../druid/quidem/DruidConnectionExtras.java | 24 ++++++++++++++++++- .../quidem/DruidQuidemCommandHandler.java | 11 +++++++-- 9 files changed, 58 insertions(+), 18 deletions(-) diff --git a/quidem-it/src/main/java/org/apache/druid/quidem/QuidemCaptureResource.java b/quidem-it/src/main/java/org/apache/druid/quidem/QuidemCaptureResource.java index e1d174ccd7f..557ea5f79b9 100644 --- a/quidem-it/src/main/java/org/apache/druid/quidem/QuidemCaptureResource.java +++ b/quidem-it/src/main/java/org/apache/druid/quidem/QuidemCaptureResource.java @@ -22,6 +22,7 @@ package org.apache.druid.quidem; import com.google.inject.Inject; import org.apache.druid.guice.LazySingleton; import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.sql.hook.DruidHookDispatcher; import javax.inject.Named; import javax.ws.rs.GET; @@ -40,11 +41,13 @@ public class QuidemCaptureResource .getPathFromProjectRoot("quidem-it/src/test/quidem/org.apache.druid.quidem.QTest"); private URI quidemURI; private QuidemRecorder recorder = null; + private DruidHookDispatcher hookDispatcher; @Inject - public QuidemCaptureResource(@Named("quidem") URI quidemURI) + public QuidemCaptureResource(@Named("quidem") URI quidemURI, DruidHookDispatcher hookDispatcher) { this.quidemURI = quidemURI; + this.hookDispatcher = hookDispatcher; } @GET @@ -55,6 +58,7 @@ public class QuidemCaptureResource stopIfRunning(); recorder = new QuidemRecorder( quidemURI, + hookDispatcher, genRecordFilePath() ); return recorder.toString(); diff --git a/quidem-it/src/main/java/org/apache/druid/quidem/QuidemRecorder.java b/quidem-it/src/main/java/org/apache/druid/quidem/QuidemRecorder.java index a05f4fec819..5f2d6a192ed 100644 --- a/quidem-it/src/main/java/org/apache/druid/quidem/QuidemRecorder.java +++ b/quidem-it/src/main/java/org/apache/druid/quidem/QuidemRecorder.java @@ -20,6 +20,7 @@ package org.apache.druid.quidem; import org.apache.druid.sql.hook.DruidHook; +import org.apache.druid.sql.hook.DruidHookDispatcher; import java.io.File; import java.io.FileNotFoundException; @@ -34,9 +35,11 @@ public class QuidemRecorder implements AutoCloseable, DruidHook { private PrintStream printStream; private File file; + private DruidHookDispatcher hookDispatcher; - public QuidemRecorder(URI quidemURI, File file) + public QuidemRecorder(URI quidemURI, DruidHookDispatcher hookDispatcher, File file) { + this.hookDispatcher = hookDispatcher; this.file = file; try { this.printStream = new PrintStream(new FileOutputStream(file), true, StandardCharsets.UTF_8.name()); @@ -47,7 +50,7 @@ public class QuidemRecorder implements AutoCloseable, DruidHook printStream.println("#started " + new Date()); printStream.println("!use " + quidemURI); printStream.println("!set outputformat mysql"); - DruidHook.register(DruidHook.SQL, this); + hookDispatcher.register(DruidHook.SQL, this); } @Override @@ -57,7 +60,7 @@ public class QuidemRecorder implements AutoCloseable, DruidHook printStream.close(); printStream = null; } - DruidHook.unregister(DruidHook.SQL, this); + hookDispatcher.unregister(DruidHook.SQL, this); } @Override diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/CalciteRulesManager.java b/sql/src/main/java/org/apache/druid/sql/calcite/planner/CalciteRulesManager.java index a1855288cd2..c42c237ee36 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/CalciteRulesManager.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/CalciteRulesManager.java @@ -395,7 +395,8 @@ public class CalciteRulesManager public RelNode run(RelOptPlanner planner, RelNode rel, RelTraitSet requiredOutputTraits, List materializations, List lattices) { - DruidHook.dispatch(DruidHook.LOGICAL_PLAN, rel); + PlannerContext pctx = planner.getContext().unwrapOrThrow(PlannerContext.class); + pctx.dispatchHook(DruidHook.LOGICAL_PLAN, rel); return rel; } } 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 b8f25dd5dcf..b543be135cc 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 @@ -155,7 +155,7 @@ public abstract class QueryHandler extends SqlStatementHandler.BaseStatementHand isPrepared = true; SqlNode validatedQueryNode = validatedQueryNode(); rootQueryRel = handlerContext.planner().rel(validatedQueryNode); - DruidHook.dispatch(DruidHook.CONVERTED_PLAN, rootQueryRel.rel); + handlerContext.plannerContext().dispatchHook(DruidHook.CONVERTED_PLAN, rootQueryRel.rel); handlerContext.hook().captureQueryRel(rootQueryRel); final RelDataTypeFactory typeFactory = rootQueryRel.rel.getCluster().getTypeFactory(); final SqlValidator validator = handlerContext.planner().getValidator(); @@ -592,7 +592,7 @@ public abstract class QueryHandler extends SqlStatementHandler.BaseStatementHand handlerContext.hook().captureDruidRel(druidRel); - DruidHook.dispatch(DruidHook.DRUID_PLAN, druidRel); + plannerContext.dispatchHook(DruidHook.DRUID_PLAN, druidRel); if (explain != null) { return planExplanation(possiblyLimitedRoot, druidRel, true); diff --git a/sql/src/main/java/org/apache/druid/sql/hook/DruidHook.java b/sql/src/main/java/org/apache/druid/sql/hook/DruidHook.java index bc9226ed3f4..fd2b32ebb98 100644 --- a/sql/src/main/java/org/apache/druid/sql/hook/DruidHook.java +++ b/sql/src/main/java/org/apache/druid/sql/hook/DruidHook.java @@ -78,31 +78,31 @@ public interface DruidHook @SuppressFBWarnings({"MS_OOI_PKGPROTECT"}) Map, List>> GLOBAL = new HashMap<>(); - static void register(HookKey label, DruidHook hook) + static void register1(HookKey label, DruidHook hook) { GLOBAL.computeIfAbsent(label, k -> new ArrayList<>()).add(hook); } - static void unregister(HookKey key, DruidHook hook) + static void unregister1(HookKey key, DruidHook hook) { GLOBAL.get(key).remove(hook); } - static Closeable withHook(HookKey key, DruidHook hook) + static Closeable withHook1(HookKey key, DruidHook hook) { - register(key, hook); + register1(key, hook); return new Closeable() { @Override public void close() { - unregister(key, hook); + unregister1(key, hook); } }; } @SuppressWarnings({"rawtypes", "unchecked"}) - static void dispatch(HookKey key, T object) + static void dispatch12(HookKey key, T object) { List> hooks = GLOBAL.get(key); if (hooks != null) { 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 688a5e9a3e4..47b80ddca36 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 @@ -20,6 +20,7 @@ package org.apache.druid.sql.hook; import com.google.inject.Inject; +import org.apache.druid.guice.LazySingleton; import org.apache.druid.sql.hook.DruidHook.HookKey; import java.io.Closeable; @@ -28,6 +29,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +@LazySingleton public class DruidHookDispatcher { @Inject diff --git a/sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java b/sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java index be48af21b8f..d2a7b461dfc 100644 --- a/sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java +++ b/sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java @@ -70,6 +70,7 @@ import org.apache.druid.sql.calcite.util.SqlTestFramework.Builder; import org.apache.druid.sql.calcite.util.SqlTestFramework.PlannerComponentSupplier; import org.apache.druid.sql.calcite.util.SqlTestFramework.QueryComponentSupplier; import org.apache.druid.sql.guice.SqlModule; +import org.apache.druid.sql.hook.DruidHookDispatcher; import org.apache.http.client.utils.URIBuilder; import org.eclipse.jetty.server.Server; @@ -138,9 +139,9 @@ public class DruidAvaticaTestDriver implements Driver @Provides @LazySingleton - public DruidConnectionExtras getConnectionExtras(ObjectMapper objectMapper) + public DruidConnectionExtras getConnectionExtras(ObjectMapper objectMapper, DruidHookDispatcher druidHookDispatcher) { - return new DruidConnectionExtras.DruidConnectionExtrasImpl(objectMapper); + return new DruidConnectionExtras.DruidConnectionExtrasImpl(objectMapper, druidHookDispatcher); } @Provides diff --git a/sql/src/test/java/org/apache/druid/quidem/DruidConnectionExtras.java b/sql/src/test/java/org/apache/druid/quidem/DruidConnectionExtras.java index 75bdd4280fa..2ba27bdd7fc 100644 --- a/sql/src/test/java/org/apache/druid/quidem/DruidConnectionExtras.java +++ b/sql/src/test/java/org/apache/druid/quidem/DruidConnectionExtras.java @@ -20,18 +20,25 @@ package org.apache.druid.quidem; import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.druid.sql.hook.DruidHookDispatcher; + +import java.sql.Connection; public interface DruidConnectionExtras { ObjectMapper getObjectMapper(); + DruidHookDispatcher getDruidHookDispatcher(); + class DruidConnectionExtrasImpl implements DruidConnectionExtras { private final ObjectMapper objectMapper; + private final DruidHookDispatcher druidHookDispatcher; - public DruidConnectionExtrasImpl(ObjectMapper objectMapper) + public DruidConnectionExtrasImpl(ObjectMapper objectMapper, DruidHookDispatcher druidHookDispatcher) { this.objectMapper = objectMapper; + this.druidHookDispatcher = druidHookDispatcher; } @Override @@ -39,5 +46,20 @@ public interface DruidConnectionExtras { return objectMapper; } + + @Override + public DruidHookDispatcher getDruidHookDispatcher() + { + return druidHookDispatcher; + } } + + static DruidConnectionExtras unwrapOrThrow(Connection connection) + { + if(connection instanceof DruidConnectionExtras ) { + return (DruidConnectionExtras) connection; + } + throw new UnsupportedOperationException("Expected DruidConnectionExtras to be implemented by connection!"); + } + } 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 58ab62a24f9..141ed0b67e4 100644 --- a/sql/src/test/java/org/apache/druid/quidem/DruidQuidemCommandHandler.java +++ b/sql/src/test/java/org/apache/druid/quidem/DruidQuidemCommandHandler.java @@ -37,10 +37,12 @@ import org.apache.druid.sql.calcite.rel.DruidRel; import org.apache.druid.sql.calcite.util.QueryLogHook; import org.apache.druid.sql.hook.DruidHook; import org.apache.druid.sql.hook.DruidHook.HookKey; +import org.apache.druid.sql.hook.DruidHookDispatcher; import java.io.Closeable; import java.io.IOException; import java.sql.ResultSet; +import java.sql.SQLException; import java.sql.Statement; import java.util.ArrayList; import java.util.List; @@ -167,10 +169,11 @@ public class DruidQuidemCommandHandler implements CommandHandler } @Override - protected final void executeExplain(Context x) throws IOException + protected final void executeExplain(Context x) throws IOException, SQLException { + DruidHookDispatcher dhp = unwrapDruidHookDispatcher(x); List logged = new ArrayList<>(); - try (Closeable unhook = DruidHook.withHook(hook, (key, relNode) -> { + try (Closeable unhook = dhp.withHook(hook, (key, relNode) -> { logged.add(relNode); })) { executeQuery(x); @@ -184,6 +187,10 @@ public class DruidQuidemCommandHandler implements CommandHandler x.echo(ImmutableList.of(str)); } } + + protected final DruidHookDispatcher unwrapDruidHookDispatcher(Context x) { + return DruidConnectionExtras.unwrapOrThrow(x.connection()).getDruidHookDispatcher(); + } } static class LogicalPlanCommand extends AbstractRelPlanCommand