From ba09e7d1ded45fdf3600831340ad42892f7ceba7 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Tue, 21 May 2024 14:00:02 +0000 Subject: [PATCH] add --- .../apache/druid/quidem/QuidemCapture.java | 24 +++- .../apache/druid/quidem/QuidemRecorder.java | 42 +++++++ .../sql/calcite/planner/PlannerFactory.java | 2 + .../org/apache/druid/quidem/DruidHook.java | 108 ++++++++++++++++++ 4 files changed, 173 insertions(+), 3 deletions(-) create mode 100644 integration-tests/src/main/java/org/apache/druid/quidem/QuidemRecorder.java create mode 100644 sql/src/test/java/org/apache/druid/quidem/DruidHook.java diff --git a/integration-tests/src/main/java/org/apache/druid/quidem/QuidemCapture.java b/integration-tests/src/main/java/org/apache/druid/quidem/QuidemCapture.java index 3b29be13ade..b9719922e04 100644 --- a/integration-tests/src/main/java/org/apache/druid/quidem/QuidemCapture.java +++ b/integration-tests/src/main/java/org/apache/druid/quidem/QuidemCapture.java @@ -28,6 +28,8 @@ import javax.ws.rs.core.MediaType; public class QuidemCapture { + private QuidemRecorder recorder = null; + @GET @Path("/") @Produces(MediaType.TEXT_PLAIN) @@ -37,10 +39,26 @@ public class QuidemCapture } @GET - @Path("/asd") + @Path("/start") @Produces(MediaType.TEXT_PLAIN) - public String getSome1() + public synchronized String getSome1() { - return "Asd"; + stopIfRunning(); + start(); + return null; + } + + private void start() + { + recorder = new QuidemRecorder(); + } + + private void stopIfRunning() + { + if (recorder != null) { + recorder.close(); + recorder = null; + } + } } diff --git a/integration-tests/src/main/java/org/apache/druid/quidem/QuidemRecorder.java b/integration-tests/src/main/java/org/apache/druid/quidem/QuidemRecorder.java new file mode 100644 index 00000000000..45659722cd0 --- /dev/null +++ b/integration-tests/src/main/java/org/apache/druid/quidem/QuidemRecorder.java @@ -0,0 +1,42 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.quidem; + +public class QuidemRecorder implements AutoCloseable, DruidHook +{ + public QuidemRecorder() + { + DruidHook.register(DruidHook.SQL, this); + DruidHook.register(DruidHook.RESULTSET, this); + } + + @Override + public void close() + { + DruidHook.unregister(DruidHook.SQL, this); + DruidHook.unregister(DruidHook.RESULTSET, this); + } + + @Override + public void dispatch1(HookKey key, T object) + { + System.out.println(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 a5cdc028e7f..ba3424cc65b 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 @@ -39,6 +39,7 @@ import org.apache.calcite.tools.FrameworkConfig; import org.apache.calcite.tools.Frameworks; import org.apache.druid.guice.annotations.Json; import org.apache.druid.math.expr.ExprMacroTable; +import org.apache.druid.quidem.DruidHook; import org.apache.druid.segment.join.JoinableFactoryWrapper; import org.apache.druid.server.security.Access; import org.apache.druid.server.security.AuthConfig; @@ -112,6 +113,7 @@ public class PlannerFactory extends PlannerToolbox queryContext, hook ); + DruidHook.dispatch(DruidHook.SQL, sql); return new DruidPlanner(buildFrameworkConfig(context), context, engine, hook); } diff --git a/sql/src/test/java/org/apache/druid/quidem/DruidHook.java b/sql/src/test/java/org/apache/druid/quidem/DruidHook.java new file mode 100644 index 00000000000..d9686cb0581 --- /dev/null +++ b/sql/src/test/java/org/apache/druid/quidem/DruidHook.java @@ -0,0 +1,108 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.quidem; + +import org.apache.calcite.rel.RelNode; +import java.io.Closeable; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +public interface DruidHook +{ + static class HookKey + { + private String label; + private Class type; + + public HookKey(String label, Class type) + { + this.label = label; + this.type = type; + } + + @Override + public int hashCode() + { + return Objects.hash(label, type); + } + + @Override + public boolean equals(Object obj) + { + if (this == obj) { + return true; + } + if (obj == null) { + return false; + } + if (getClass() != obj.getClass()) { + return false; + } + HookKey other = (HookKey) obj; + return Objects.equals(label, other.label) && Objects.equals(type, other.type); + } + + } + + public static final HookKey DRUID_PLAN = new HookKey<>("druidPlan", RelNode.class); + public static final HookKey SQL = new HookKey<>("sql", String.class); + public static final HookKey RESULTSET = new HookKey<>("resultSet", String.class); + + void dispatch1(HookKey key, T object); + + static Map, List> GLOBAL = new HashMap<>(); + + public static void register(HookKey label, DruidHook hook) + { + GLOBAL.computeIfAbsent(label, k -> new ArrayList<>()).add(hook); + } + + public static void unregister(HookKey key, DruidHook hook) + { + GLOBAL.get(key).remove(hook); + } + + public static Closeable withHook(HookKey key, DruidHook hook) + { + register(key, hook); + return new Closeable() + { + @Override + public void close() + { + unregister(key, hook); + } + }; + } + + public static void dispatch(HookKey key, T object) + { + List hooks = GLOBAL.get(key); + if (hooks != null) { + for (DruidHook hook : hooks) { + hook.dispatch1(key, object); + } + } + } + +}