Skip to content

Commit b13d88b

Browse files
ahkcsahkcs
committed
Refactor context classloader wrapping into CalciteClassLoaderHelper
Extract the CALCITE-3745 context classloader fix into a dedicated CalciteClassLoaderHelper utility class. Remove redundant wrapper from OpenSearchRelRunners.run() (already covered by QueryService callers). Three call sites remain: - QueryService.executeWithCalcite() — covers planning + execution - QueryService.explainWithCalcite() — covers planning + explain - CalciteScriptEngine.compile() — covers pushdown scripts (shard thread) Signed-off-by: Kai Huang <ahkcs@amazon.com>
1 parent f902540 commit b13d88b

4 files changed

Lines changed: 95 additions & 55 deletions

File tree

‎core/src/main/java/org/opensearch/sql/calcite/utils/CalciteClassLoaderHelper.java‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.utils;
7+
8+
import java.util.concurrent.Callable;
9+
10+
/**
11+
* Helper for setting the thread context classloader before Calcite operations. This is needed for
12+
* patched Calcite (CALCITE-3745): when analytics-engine is the parent classloader, Janino uses the
13+
* parent's classloader which can't see SQL plugin classes. The patched Calcite checks {@code
14+
* Thread.currentThread().getContextClassLoader()} first. This helper sets it to the SQL plugin's
15+
* classloader (child) which can see both parent and child classes.
16+
*
17+
* @see <a href="https://issues.apache.org/jira/browse/CALCITE-3745">CALCITE-3745</a>
18+
* @see <a href="https://github.com/opensearch-project/sql/issues/5306">sql#5306</a>
19+
*/
20+
public final class CalciteClassLoaderHelper {
21+
22+
private CalciteClassLoaderHelper() {}
23+
24+
/**
25+
* Run an action with the thread context classloader set to the caller's classloader.
26+
*
27+
* @param action the action to run
28+
* @param callerClass the class whose classloader should be used (pass {@code MyClass.class})
29+
* @param <T> the return type
30+
* @return the result of the action
31+
*/
32+
public static <T> T withCalciteClassLoader(Callable<T> action, Class<?> callerClass) {
33+
Thread currentThread = Thread.currentThread();
34+
ClassLoader originalCl = currentThread.getContextClassLoader();
35+
currentThread.setContextClassLoader(callerClass.getClassLoader());
36+
try {
37+
return action.call();
38+
} catch (RuntimeException e) {
39+
throw e;
40+
} catch (Exception e) {
41+
throw new RuntimeException(e);
42+
} finally {
43+
currentThread.setContextClassLoader(originalCl);
44+
}
45+
}
46+
47+
/**
48+
* Run a void action with the thread context classloader set to the caller's classloader.
49+
*
50+
* @see #withCalciteClassLoader(Callable, Class)
51+
*/
52+
public static void withCalciteClassLoader(Runnable action, Class<?> callerClass) {
53+
withCalciteClassLoader(
54+
() -> {
55+
action.run();
56+
return null;
57+
},
58+
callerClass);
59+
}
60+
}

‎core/src/main/java/org/opensearch/sql/calcite/utils/CalciteToolsHelper.java‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -422,11 +422,8 @@ public RelNode visit(TableScan scan) {
422422
}
423423
};
424424
rel = rel.accept(shuttle);
425-
// Set thread context classloader so patched Calcite (CALCITE-3745) uses the SQL
426-
// plugin's classloader for Janino compilation instead of the parent classloader.
427-
Thread currentThread = Thread.currentThread();
428-
ClassLoader originalCl = currentThread.getContextClassLoader();
429-
currentThread.setContextClassLoader(CalciteToolsHelper.class.getClassLoader());
425+
// Context classloader is set by callers (QueryService.executeWithCalcite/explainWithCalcite)
426+
// to support patched Calcite (CALCITE-3745) Janino compilation.
430427
try (Connection connection = context.connection) {
431428
final RelRunner runner = connection.unwrap(RelRunner.class);
432429
PreparedStatement preparedStatement = runner.prepareStatement(rel);
@@ -445,8 +442,6 @@ public RelNode visit(TableScan scan) {
445442
+ " count() by @timestamp').");
446443
}
447444
throw Util.throwAsRuntime(e);
448-
} finally {
449-
currentThread.setContextClassLoader(originalCl);
450445
}
451446
}
452447
}

‎core/src/main/java/org/opensearch/sql/executor/QueryService.java‎

Lines changed: 28 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import lombok.RequiredArgsConstructor;
1414
import lombok.extern.log4j.Log4j2;
1515
import org.apache.calcite.jdbc.CalciteSchema;
16+
import org.opensearch.sql.calcite.utils.CalciteClassLoaderHelper;
1617
import org.apache.calcite.plan.RelTraitDef;
1718
import org.apache.calcite.rel.RelCollation;
1819
import org.apache.calcite.rel.RelCollations;
@@ -139,23 +140,18 @@ public void executeWithCalcite(
139140
QueryProfiling.activate(QueryContext.isProfileEnabled());
140141
ProfileMetric analyzeMetric = profileContext.getOrCreateMetric(MetricName.ANALYZE);
141142
long analyzeStart = System.nanoTime();
142-
// Set thread context classloader so patched Calcite (CALCITE-3745) uses the
143-
// SQL plugin's classloader for Janino compilation during planning/execution.
144-
Thread currentThread = Thread.currentThread();
145-
ClassLoader originalCl = currentThread.getContextClassLoader();
146-
currentThread.setContextClassLoader(QueryService.class.getClassLoader());
147-
try {
148-
CalcitePlanContext context =
149-
CalcitePlanContext.create(
150-
buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType);
151-
context.setHighlightConfig(highlightConfig);
152-
RelNode relNode = analyze(plan, context);
153-
RelNode calcitePlan = convertToCalcitePlan(relNode, context);
154-
analyzeMetric.set(System.nanoTime() - analyzeStart);
155-
executionEngine.execute(calcitePlan, context, listener);
156-
} finally {
157-
currentThread.setContextClassLoader(originalCl);
158-
}
143+
CalciteClassLoaderHelper.withCalciteClassLoader(
144+
() -> {
145+
CalcitePlanContext context =
146+
CalcitePlanContext.create(
147+
buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType);
148+
context.setHighlightConfig(highlightConfig);
149+
RelNode relNode = analyze(plan, context);
150+
RelNode calcitePlan = convertToCalcitePlan(relNode, context);
151+
analyzeMetric.set(System.nanoTime() - analyzeStart);
152+
executionEngine.execute(calcitePlan, context, listener);
153+
},
154+
QueryService.class);
159155
} catch (Throwable t) {
160156
if (isCalciteFallbackAllowed(t) && !(t instanceof NonFallbackCalciteException)) {
161157
log.warn("Fallback to V2 query engine since got exception", t);
@@ -178,26 +174,21 @@ public void explainWithCalcite(
178174
() -> {
179175
try {
180176
QueryProfiling.noop();
181-
// Set thread context classloader so patched Calcite (CALCITE-3745) uses the
182-
// SQL plugin's classloader for Janino compilation during planning/execution.
183-
Thread currentThread = Thread.currentThread();
184-
ClassLoader originalCl = currentThread.getContextClassLoader();
185-
currentThread.setContextClassLoader(QueryService.class.getClassLoader());
186-
try {
187-
CalcitePlanContext context =
188-
CalcitePlanContext.create(
189-
buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType);
190-
context.setHighlightConfig(highlightConfig);
191-
context.run(
192-
() -> {
193-
RelNode relNode = analyze(plan, context);
194-
RelNode calcitePlan = convertToCalcitePlan(relNode, context);
195-
executionEngine.explain(calcitePlan, mode, context, listener);
196-
},
197-
settings);
198-
} finally {
199-
currentThread.setContextClassLoader(originalCl);
200-
}
177+
CalciteClassLoaderHelper.withCalciteClassLoader(
178+
() -> {
179+
CalcitePlanContext context =
180+
CalcitePlanContext.create(
181+
buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType);
182+
context.setHighlightConfig(highlightConfig);
183+
context.run(
184+
() -> {
185+
RelNode relNode = analyze(plan, context);
186+
RelNode calcitePlan = convertToCalcitePlan(relNode, context);
187+
executionEngine.explain(calcitePlan, mode, context, listener);
188+
},
189+
settings);
190+
},
191+
QueryService.class);
201192
} catch (Throwable t) {
202193
if (isCalciteFallbackAllowed(t)) {
203194
log.warn("Fallback to V2 query engine since got exception", t);

‎opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/CalciteScriptEngine.java‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@
6969
import org.apache.calcite.util.Util;
7070
import org.checkerframework.checker.nullness.qual.Nullable;
7171
import org.opensearch.index.fielddata.ScriptDocValues;
72+
import org.opensearch.sql.calcite.utils.CalciteClassLoaderHelper;
7273
import org.opensearch.script.AggregationScript;
7374
import org.opensearch.script.FieldScript;
7475
import org.opensearch.script.FilterScript;
@@ -137,17 +138,10 @@ public <T> T compile(
137138
getter,
138139
new RelRecordType(List.of()));
139140

140-
// Set thread context classloader so patched Calcite (CALCITE-3745) uses the SQL
141-
// plugin's classloader for Janino compilation instead of the parent classloader.
142-
Thread currentThread = Thread.currentThread();
143-
ClassLoader originalCl = currentThread.getContextClassLoader();
144-
currentThread.setContextClassLoader(CalciteScriptEngine.class.getClassLoader());
145-
Function1<DataContext, Object[]> function;
146-
try {
147-
function = new RexExecutable(code, "generated Rex code").getFunction();
148-
} finally {
149-
currentThread.setContextClassLoader(originalCl);
150-
}
141+
Function1<DataContext, Object[]> function =
142+
CalciteClassLoaderHelper.withCalciteClassLoader(
143+
() -> new RexExecutable(code, "generated Rex code").getFunction(),
144+
CalciteScriptEngine.class);
151145

152146
if (CONTEXTS.containsKey(context)) {
153147
return context.factoryClazz.cast(CONTEXTS.get(context).apply(function, rexNode.getType()));

0 commit comments

Comments
 (0)