From d970243c61e28c1fb40e7ef752780760fb46af11 Mon Sep 17 00:00:00 2001 From: Kai Huang Date: Thu, 10 Sep 2026 11:31:49 -0700 Subject: [PATCH] Carry the partial_result override off ThreadContext so it survives the security handoff The per-request partial_result override was stored in Log4j ThreadContext (QueryContext.setPartialResultOverride), set on the transport thread and read on the worker/planning thread by CalciteLogicalIndexScan.tryPartialResultAggregate. The security plugin's transport interceptor does not preserve Log4j ThreadContext across that handoff, so on a security-enabled cluster the override is lost and isPartialResultEnabled silently falls back to the cluster setting. A request that sends partial_result=false to opt out therefore still gets a partial result (with a PARTIAL_RESULT warning) when the cluster setting is on -- the setting's own path is unaffected because it is read from osIndex.getSettings() on the worker thread. This is the same transport->worker handoff bug #5739 fixed for warningsSupported; #5739 moved warningsSupported onto the object graph but left the override on ThreadContext (reasoning the settings fallback made it resilient -- but the fallback masks the drop). Carry the override the same way warningsSupported is carried: PPLQueryRequest.partialResult -> AbstractPlan -> QueryPlan.execute() (worker thread) -> CalcitePlanContext (in its thread-local snapshot, reset per query). tryPartialResultAggregate now reads the override from CalcitePlanContext and falls back to the cluster setting. Removes the dead QueryContext override methods. Reproduced and verified with a new PartialResultSecurityIT case (perRequestPartialResultFalseOverridesClusterSettingUnderSecurity): fails before this change (complete expected, partial returned), passes after. Signed-off-by: Kai Huang --- .../sql/common/utils/QueryContext.java | 32 -------------- .../sql/calcite/CalcitePlanContext.java | 35 +++++++++++++++- .../sql/executor/execution/AbstractPlan.java | 8 ++++ .../sql/executor/execution/QueryPlan.java | 7 ++-- .../sql/security/PartialResultSecurityIT.java | 42 +++++++++++++++++++ .../storage/scan/CalciteLogicalIndexScan.java | 13 +++++- .../transport/TransportPPLQueryAction.java | 12 +++--- .../org/opensearch/sql/ppl/PPLService.java | 1 + 8 files changed, 105 insertions(+), 45 deletions(-) diff --git a/common/src/main/java/org/opensearch/sql/common/utils/QueryContext.java b/common/src/main/java/org/opensearch/sql/common/utils/QueryContext.java index 3a8e06cc9bf..4d4301df473 100644 --- a/common/src/main/java/org/opensearch/sql/common/utils/QueryContext.java +++ b/common/src/main/java/org/opensearch/sql/common/utils/QueryContext.java @@ -8,7 +8,6 @@ import java.util.Map; import java.util.UUID; import org.apache.logging.log4j.ThreadContext; -import org.opensearch.sql.common.setting.Settings; /** * Utility class for recording and accessing context for the query being executed. Implementation @@ -23,8 +22,6 @@ public class QueryContext { private static final String PROFILE_KEY = "profile"; - private static final String PARTIAL_RESULT_OVERRIDE_KEY = "partial_result_override"; - /** * Generates a random UUID and adds to the {@link ThreadContext} as the request id. * @@ -87,33 +84,4 @@ public static void setProfile(boolean profileEnabled) { public static boolean isProfileEnabled() { return Boolean.parseBoolean(ThreadContext.get(PROFILE_KEY)); } - - /** - * Record a per-request override for partial-result mode. When set, it takes precedence over the - * cluster setting: {@code true} forces partial mode on for this request, {@code false} forces it - * off. A {@code null} value (the default) leaves the decision to the cluster setting. - * - * @param override the per-request preference, or null to defer to the cluster setting - */ - public static void setPartialResultOverride(Boolean override) { - if (override == null) { - ThreadContext.remove(PARTIAL_RESULT_OVERRIDE_KEY); - } else { - ThreadContext.put(PARTIAL_RESULT_OVERRIDE_KEY, Boolean.toString(override)); - } - } - - /** - * Whether partial-result mode applies to the current query. The per-request override wins when - * present; otherwise the cluster setting decides. - * - * @param settings the plugin settings to read the cluster default from - */ - public static boolean isPartialResultEnabled(Settings settings) { - String override = ThreadContext.get(PARTIAL_RESULT_OVERRIDE_KEY); - if (override != null) { - return Boolean.parseBoolean(override); - } - return settings.getSettingValue(Settings.Key.PARTIAL_RESULT_ON_MAPPING_CONFLICT); - } } diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java index 11c74edb4dc..7ba2c1df582 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java @@ -82,6 +82,15 @@ public class CalcitePlanContext { private static final ThreadLocal warningsSupported = ThreadLocal.withInitial(() -> false); + /** + * Per-request partial-result override, carried off Log4j {@code ThreadContext} onto the plan (see + * {@code QueryPlan#execute}) for the same reason as {@link #warningsSupported}: the security + * plugin's interceptor drops {@code ThreadContext} on the transport→worker handoff. {@code null} + * defers to the cluster setting; {@code true}/{@code false} force partial mode on/off for this + * query. Cleared per query. + */ + private static final ThreadLocal partialResultOverride = new ThreadLocal<>(); + /** Thread-local switch that tells whether the current query prefers legacy behavior. */ private static final ThreadLocal legacyPreferredFlag = ThreadLocal.withInitial(() -> true); @@ -271,6 +280,7 @@ public static void clearTimewrapSignals() { executionPool.set(null); pendingWarnings.remove(); warningsSupported.set(false); + partialResultOverride.remove(); } /** Records a non-fatal warning to be attached to the response for the current query. */ @@ -291,6 +301,22 @@ public static boolean isWarningsSupported() { return warningsSupported.get(); } + /** + * Records the per-request partial-result override for the current query. {@code null} defers to + * the cluster setting; {@code true}/{@code false} force partial mode on/off. + */ + public static void setPartialResultOverride(Boolean override) { + partialResultOverride.set(override); + } + + /** + * @return the per-request partial-result override, or {@code null} to defer to the cluster + * setting. + */ + public static Boolean getPartialResultOverride() { + return partialResultOverride.get(); + } + /** * Returns and clears the warnings collected for the current query, de-duplicated by value. The * planner may fire a rule that raises a warning more than once for equivalent plan alternatives, @@ -317,6 +343,7 @@ public static class ThreadLocalSnapshot { final String timewrapSeries; final String executionPool; final boolean warningsSupported; + final Boolean partialResultOverride; private ThreadLocalSnapshot( boolean skipEncoding, @@ -324,13 +351,15 @@ private ThreadLocalSnapshot( String timewrapUnitName, String timewrapSeries, String executionPool, - boolean warningsSupported) { + boolean warningsSupported, + Boolean partialResultOverride) { this.skipEncoding = skipEncoding; this.stripNullColumns = stripNullColumns; this.timewrapUnitName = timewrapUnitName; this.timewrapSeries = timewrapSeries; this.executionPool = executionPool; this.warningsSupported = warningsSupported; + this.partialResultOverride = partialResultOverride; } } @@ -342,7 +371,8 @@ public static ThreadLocalSnapshot snapshotThreadLocals() { timewrapUnitName.get(), timewrapSeries.get(), executionPool.get(), - warningsSupported.get()); + warningsSupported.get(), + partialResultOverride.get()); } /** Restore thread-local state from a snapshot. */ @@ -353,6 +383,7 @@ public static void restoreThreadLocals(ThreadLocalSnapshot snapshot) { timewrapSeries.set(snapshot.timewrapSeries); executionPool.set(snapshot.executionPool); warningsSupported.set(snapshot.warningsSupported); + partialResultOverride.set(snapshot.partialResultOverride); } public void pushForeachBindings( diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/AbstractPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/AbstractPlan.java index 2b12c4f9779..1e8835b4815 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/AbstractPlan.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/AbstractPlan.java @@ -32,6 +32,14 @@ public abstract class AbstractPlan { */ @Getter @Setter private boolean warningsSupported = false; + /** + * Per-request partial-result override, carried from the request the same way as {@link + * #warningsSupported}. {@code null} defers to the cluster setting; {@code true}/{@code false} + * force partial mode on/off. Set on the transport thread, applied on the worker (see {@code + * QueryPlan#execute}) so it survives the security transport→worker handoff. + */ + @Getter @Setter private Boolean partialResultOverride = null; + /** Start query execution. */ public abstract void execute(); diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java index 904b35a45be..570bf21f97f 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java @@ -106,10 +106,11 @@ public QueryPlan( @Override public void execute() { - // Runs on the worker thread; carry warnings support from the request off the plan so the - // partial-result gate reads it without depending on Log4j ThreadContext (dropped under - // security). + // Runs on the worker thread; carry warnings support and the per-request partial-result override + // from the request off the plan so the partial-result gate reads them without depending on + // Log4j ThreadContext (dropped under security on the transport→worker handoff). CalcitePlanContext.setWarningsSupported(isWarningsSupported()); + CalcitePlanContext.setPartialResultOverride(getPartialResultOverride()); if (pageSize.isPresent()) { queryService.execute( new Paginate(pageSize.get(), plan), diff --git a/integ-test/src/test/java/org/opensearch/sql/security/PartialResultSecurityIT.java b/integ-test/src/test/java/org/opensearch/sql/security/PartialResultSecurityIT.java index f925b396851..5cd300fe225 100644 --- a/integ-test/src/test/java/org/opensearch/sql/security/PartialResultSecurityIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/security/PartialResultSecurityIT.java @@ -12,11 +12,14 @@ import static org.opensearch.sql.util.TestUtils.performRequest; import java.io.IOException; +import java.util.Locale; import org.json.JSONArray; import org.json.JSONObject; import org.junit.After; import org.junit.Test; import org.opensearch.client.Request; +import org.opensearch.client.RequestOptions; +import org.opensearch.client.Response; import org.opensearch.sql.common.setting.Settings; /** @@ -112,6 +115,45 @@ public void completeResultCarriesNoWarningWithSecurity() throws IOException { assertFalse("a complete result carries no warning", result.has("warnings")); } + @Test + public void perRequestPartialResultFalseOverridesClusterSettingUnderSecurity() + throws IOException { + // Cluster setting ON, but the request explicitly opts OUT via partial_result=false. The + // per-request override must win -> complete result over all indices, no warning. On the buggy + // code the override lives in Log4j ThreadContext (QueryContext.setPartialResultOverride) and is + // dropped by the security transport->worker handoff -- the same drop #5739 fixed for + // warningsSupported but left in place for the override -- so it silently falls back to the ON + // cluster setting and returns a partial result with a warning. + setPartialResult(true); + JSONObject result = + executeQueryAsUserWithPartialResult( + String.format("source=%s | stats count() by env | sort env", PATTERN), USER, false); + // Every index contributes: keyword (prod=2, dev=1) + text (prod=1, qa=1). + verifyDataRows(result, rows(1, "dev"), rows(1, "qa"), rows(3, "prod")); + assertFalse( + "partial_result=false must override the ON cluster setting -> complete result, no warning", + result.has("warnings")); + } + + /** Like {@link #executeQueryAsUser}, but also sends the per-request {@code partial_result}. */ + private JSONObject executeQueryAsUserWithPartialResult( + String query, String username, boolean partialResult) throws IOException { + Request request = new Request("POST", "/_plugins/_ppl"); + request.setJsonEntity( + String.format( + Locale.ROOT, + "{ \"query\": \"%s\", \"partial_result\": %s }", + query, + Boolean.toString(partialResult))); + RequestOptions.Builder options = RequestOptions.DEFAULT.toBuilder(); + options.addHeader("Content-Type", "application/json"); + options.addHeader("Authorization", createBasicAuthHeader(username, STRONG_PASSWORD)); + request.setOptions(options); + Response response = client().performRequest(request); + assertEquals(200, response.getStatusLine().getStatusCode()); + return new JSONObject(org.opensearch.sql.legacy.TestUtils.getResponseBody(response, true)); + } + private void setPartialResult(boolean enabled) throws IOException { updateClusterSettings( new ClusterSetting( diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java index 1a84935a5d1..09cd5278dc5 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java @@ -51,7 +51,6 @@ import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory; import org.opensearch.sql.calcite.utils.PPLHintUtils; import org.opensearch.sql.common.setting.Settings; -import org.opensearch.sql.common.utils.QueryContext; import org.opensearch.sql.data.type.ExprCoreType; import org.opensearch.sql.data.type.ExprType; import org.opensearch.sql.expression.HighlightExpression; @@ -510,7 +509,17 @@ public Void visitInputRef(RexInputRef ref) { */ private AbstractRelNode tryPartialResultAggregate( Aggregate aggregate, @Nullable Project project, List partitionFields) { - if (!QueryContext.isPartialResultEnabled(osIndex.getSettings())) { + // The per-request override wins when present; otherwise the cluster setting decides. Both the + // override and warnings-support are read from CalcitePlanContext (carried onto the plan), not + // Log4j ThreadContext, so they survive the security transport→worker handoff. + Boolean override = CalcitePlanContext.getPartialResultOverride(); + boolean partialResultEnabled = + override != null + ? override + : osIndex + .getSettings() + .getSettingValue(Settings.Key.PARTIAL_RESULT_ON_MAPPING_CONFLICT); + if (!partialResultEnabled) { return null; } // A format with no warnings channel (CSV/RAW/VIZ) must not silently drop indices. diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java index 660ecb23af6..5a16c60b99f 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java @@ -198,8 +198,9 @@ protected void doExecute( // data silently. Carried on the request (not Log4j ThreadContext) so it survives the // transport→worker handoff, which the security plugin's interceptor does not preserve. transformedRequest.warningsSupported(warningsSupported(transformedRequest)); - // Per-request override (e.g. a Dashboards toggle); null defers to the cluster setting. - QueryContext.setPartialResultOverride(transformedRequest.partialResult()); + // The per-request partial-result override (e.g. a Dashboards toggle) rides on the request → + // plan → worker thread (see PPLService/QueryPlan), not Log4j ThreadContext, for the same + // handoff-survival reason as warningsSupported. null defers to the cluster setting. // Start root span with OTel DB semantic convention attributes Span rootSpan = @@ -446,12 +447,11 @@ public void onFailure(Exception e) { /** * Clear the per-request state carried in {@link QueryContext}'s thread-locals. Transport threads - * are pooled, so anything left behind is inherited by the next query to run on this thread -- a - * request that expressed no partial-result preference would otherwise pick up the previous - * request's override. + * are pooled, so anything left behind is inherited by the next query to run on this thread. (The + * partial-result override no longer lives here -- it rides on the plan to the worker thread and + * is reset per query in {@code CalcitePlanContext}.) */ private static void clearRequestScopedState() { QueryProfiling.clear(); - QueryContext.setPartialResultOverride(null); } } diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java index 54d57370168..224cf9c518d 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java @@ -211,6 +211,7 @@ private AbstractPlan plan( AbstractPlan plan = queryExecutionFactory.create(statement, queryListener, explainListener); plan.setWarningsSupported(request.warningsSupported()); + plan.setPartialResultOverride(request.partialResult()); return plan; } }