Apache
/
druid
mirror of https://github.com/apache/druid.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

use updated hook class

This commit is contained in:
Zoltan Haindrich 2024-07-30 16:11:57 +00:00
parent de207c7295
commit 5f6290eb54
2 changed files with 17 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/MSQDruidMeta.java
View File

@ -36,8 +36,9 @@ import org.apache.druid.segment.column.RowSignature;
import org.apache.druid.server.security.AuthenticatorMapper; import org.apache.druid.server.security.AuthenticatorMapper;
import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.SqlStatementFactory;
import org.apache.druid.sql.calcite.planner.DruidTypeSystem; 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.calcite.table.RowSignatures;
import org.apache.druid.sql.hook.DruidHook;
import org.apache.druid.sql.hook.DruidHookDispatcher;
import java.util.Collections; import java.util.Collections;
import java.util.List; import java.util.List;
@ -45,7 +46,8 @@ import java.util.List;
public class MSQDruidMeta extends DruidMeta public class MSQDruidMeta extends DruidMeta
{ {
protected final MSQTestOverlordServiceClient overlordClient; protected final MSQTestOverlordServiceClient overlordClient;
private final ObjectMapper objectMapper; protected final ObjectMapper objectMapper;
protected final DruidHookDispatcher hookDispatcher;
@Inject @Inject
public MSQDruidMeta( public MSQDruidMeta(
@ -54,11 +56,14 @@ public class MSQDruidMeta extends DruidMeta
final ErrorHandler errorHandler, final ErrorHandler errorHandler,
final AuthenticatorMapper authMapper, final AuthenticatorMapper authMapper,
final MSQTestOverlordServiceClient overlordClient, final MSQTestOverlordServiceClient overlordClient,
final ObjectMapper objectMapper) final ObjectMapper objectMapper,
final DruidHookDispatcher hookDispatcher
)
{ {
super(sqlStatementFactory, config, errorHandler, authMapper); super(sqlStatementFactory, config, errorHandler, authMapper);
this.overlordClient = overlordClient; this.overlordClient = overlordClient;
this.objectMapper = objectMapper; this.objectMapper = objectMapper;
this.hookDispatcher = hookDispatcher;
} }
@Override @Override
@ -89,10 +94,10 @@ public class MSQDruidMeta extends DruidMeta
.writerWithDefaultPrettyPrinter() .writerWithDefaultPrettyPrinter()
.writeValueAsString(payload.getStages()); .writeValueAsString(payload.getStages());
str = str.replaceAll(taskId, "<taskId>"); str = str.replaceAll(taskId, "<taskId>");
DruidHook.dispatch(DruidHook.MSQ_PLAN, str); hookDispatcher.dispatch(DruidHook.MSQ_PLAN, str);
} }
catch (JsonProcessingException e) { 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()); Signature signature = makeSignature(druidStatement, payload.getResults().getSignature());

13
sql/src/test/java/org/apache/druid/quidem/DruidQuidemCommandHandler.java
View File

@ -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; protected abstract void executeExplain(Context x) throws Exception;
} }
@ -190,11 +195,6 @@ public class DruidQuidemCommandHandler implements CommandHandler
x.echo(ImmutableList.of(str)); 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 @Override
protected final void executeExplain(Context x) throws IOException protected final void executeExplain(Context x) throws IOException
{ {
DruidHookDispatcher dhp = unwrapDruidHookDispatcher(x);
List<String> logged = new ArrayList<>(); List<String> logged = new ArrayList<>();
try (Closeable unhook = DruidHook.withHook(hook, (key, relNode) -> { try (Closeable unhook = dhp.withHook(hook, (key, relNode) -> {
logged.add(relNode); logged.add(relNode);
})) { })) {
executeQuery(x); executeQuery(x);