Carry the partial_result override off ThreadContext so it survives the security handoff

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);
-  }
 }

core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java

@@ -82,6 +82,15 @@ public class CalcitePlanContext {
   private static final ThreadLocal<Boolean> 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<Boolean> partialResultOverride = new ThreadLocal<>();
+
   /** Thread-local switch that tells whether the current query prefers legacy behavior. */
   private static final ThreadLocal<Boolean> 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,20 +343,23 @@ public static class ThreadLocalSnapshot {
     final String timewrapSeries;
     final String executionPool;
     final boolean warningsSupported;
+    final Boolean partialResultOverride;
 
     private ThreadLocalSnapshot(
         boolean skipEncoding,
         boolean stripNullColumns,
         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(

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();
 

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),

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(

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<String> 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.

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);
   }
 }

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;
   }
 }