Fix partial-result warnings gate lost across the security thread handoff (#5739)

common/src/main/java/org/opensearch/sql/common/utils/QueryContext.java

@@ -23,8 +23,6 @@ public class QueryContext {
 
   private static final String PROFILE_KEY = "profile";
 
-  private static final String WARNINGS_SUPPORTED_KEY = "warnings_supported";
-
   private static final String PARTIAL_RESULT_OVERRIDE_KEY = "partial_result_override";
 
   /**
@@ -90,26 +88,6 @@ public static boolean isProfileEnabled() {
     return Boolean.parseBoolean(ThreadContext.get(PROFILE_KEY));
   }
 
-  /**
-   * Record whether the requested response format can surface non-fatal warnings. Features that
-   * return a knowingly-partial result gate on this so they never silently drop data into a format
-   * (CSV/RAW) that has no warning channel.
-   *
-   * @param supported whether the response format carries a warnings channel
-   */
-  public static void setWarningsSupported(boolean supported) {
-    ThreadContext.put(WARNINGS_SUPPORTED_KEY, Boolean.toString(supported));
-  }
-
-  /**
-   * @return true if the response format for the current request can surface warnings. Defaults to
-   *     false when unset, so a caller that never declared support cannot get a silent partial
-   *     result.
-   */
-  public static boolean isWarningsSupported() {
-    return Boolean.parseBoolean(ThreadContext.get(WARNINGS_SUPPORTED_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

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

@@ -73,6 +73,15 @@ public class CalcitePlanContext {
   private static final ThreadLocal<List<Warning>> pendingWarnings =
       ThreadLocal.withInitial(ArrayList::new);
 
+  /**
+   * Whether the current query's response format can carry a warnings channel. Set on the worker
+   * thread from the plan (see {@code QueryPlan#execute}) rather than the transport thread, so the
+   * partial-result gate survives the transport→worker handoff — the security plugin's interceptor
+   * drops Log4j {@code ThreadContext}, which is where this used to live. Cleared per query.
+   */
+  private static final ThreadLocal<Boolean> warningsSupported =
+      ThreadLocal.withInitial(() -> false);
+
   /** Thread-local switch that tells whether the current query prefers legacy behavior. */
   private static final ThreadLocal<Boolean> legacyPreferredFlag =
       ThreadLocal.withInitial(() -> true);
@@ -261,13 +270,27 @@ public static void clearTimewrapSignals() {
     timewrapSeries.set(null);
     executionPool.set(null);
     pendingWarnings.remove();
+    warningsSupported.set(false);
   }
 
   /** Records a non-fatal warning to be attached to the response for the current query. */
   public static void addWarning(Warning warning) {
     pendingWarnings.get().add(warning);
   }
 
+  /** Records whether the current query's response format can surface warnings. */
+  public static void setWarningsSupported(boolean supported) {
+    warningsSupported.set(supported);
+  }
+
+  /**
+   * @return whether the current query's response format can surface warnings; false when unset, so
+   *     a caller that never declared support cannot get a silent partial result.
+   */
+  public static boolean isWarningsSupported() {
+    return warningsSupported.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,
@@ -293,18 +316,21 @@ public static class ThreadLocalSnapshot {
     final String timewrapUnitName;
     final String timewrapSeries;
     final String executionPool;
+    final boolean warningsSupported;
 
     private ThreadLocalSnapshot(
         boolean skipEncoding,
         boolean stripNullColumns,
         String timewrapUnitName,
         String timewrapSeries,
-        String executionPool) {
+        String executionPool,
+        boolean warningsSupported) {
       this.skipEncoding = skipEncoding;
       this.stripNullColumns = stripNullColumns;
       this.timewrapUnitName = timewrapUnitName;
       this.timewrapSeries = timewrapSeries;
       this.executionPool = executionPool;
+      this.warningsSupported = warningsSupported;
     }
   }
 
@@ -315,7 +341,8 @@ public static ThreadLocalSnapshot snapshotThreadLocals() {
         stripNullColumns.get(),
         timewrapUnitName.get(),
         timewrapSeries.get(),
-        executionPool.get());
+        executionPool.get(),
+        warningsSupported.get());
   }
 
   /** Restore thread-local state from a snapshot. */
@@ -325,6 +352,7 @@ public static void restoreThreadLocals(ThreadLocalSnapshot snapshot) {
     timewrapUnitName.set(snapshot.timewrapUnitName);
     timewrapSeries.set(snapshot.timewrapSeries);
     executionPool.set(snapshot.executionPool);
+    warningsSupported.set(snapshot.warningsSupported);
   }
 
   public void pushForeachBindings(

core/src/main/java/org/opensearch/sql/executor/execution/AbstractPlan.java

@@ -7,6 +7,7 @@
 
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
+import lombok.Setter;
 import org.opensearch.sql.ast.statement.ExplainMode;
 import org.opensearch.sql.common.response.ResponseListener;
 import org.opensearch.sql.executor.ExecutionEngine;
@@ -23,6 +24,14 @@ public abstract class AbstractPlan {
 
   @Getter protected final QueryType queryType;
 
+  /**
+   * Whether the response format can carry a warnings channel. Set from the request on the transport
+   * thread and read on the worker (see {@code QueryPlan#execute}), so a feature that returns a
+   * partial result never silently drops data into a warnings-incapable format across a thread
+   * handoff. Defaults to false.
+   */
+  @Getter @Setter private boolean warningsSupported = false;
+
   /** Start query execution. */
   public abstract void execute();
 

core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java

@@ -11,6 +11,7 @@
 import org.opensearch.sql.ast.tree.HighlightConfig;
 import org.opensearch.sql.ast.tree.Paginate;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
+import org.opensearch.sql.calcite.CalcitePlanContext;
 import org.opensearch.sql.common.response.ResponseListener;
 import org.opensearch.sql.executor.ExecutionEngine;
 import org.opensearch.sql.executor.QueryId;
@@ -105,6 +106,10 @@ 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).
+    CalcitePlanContext.setWarningsSupported(isWarningsSupported());
     if (pageSize.isPresent()) {
       queryService.execute(
           new Paginate(pageSize.get(), plan),

core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java

@@ -5,6 +5,7 @@
 
 package org.opensearch.sql.executor.execution;
 
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
@@ -16,6 +17,7 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.commons.lang3.NotImplementedException;
 import org.junit.jupiter.api.DisplayNameGeneration;
 import org.junit.jupiter.api.DisplayNameGenerator;
@@ -25,6 +27,7 @@
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.opensearch.sql.ast.statement.ExplainMode;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
+import org.opensearch.sql.calcite.CalcitePlanContext;
 import org.opensearch.sql.common.response.ResponseListener;
 import org.opensearch.sql.executor.DefaultExecutionEngine;
 import org.opensearch.sql.executor.ExecutionEngine;
@@ -58,6 +61,53 @@ public void execute_no_page_size() {
     verify(queryService, times(1)).execute(any(), any(), any(), anyBoolean(), any());
   }
 
+  @Test
+  public void warnings_supported_flag_reaches_execution_thread() throws InterruptedException {
+    QueryPlan query = new QueryPlan(queryId, queryType, plan, queryService, queryListener);
+    query.setWarningsSupported(true);
+
+    // Configure the plan here but run execute() on another thread, mirroring the transport->worker
+    // handoff. The flag must ride the plan object, not a thread-local the handoff can drop.
+    AtomicBoolean defaultedBeforeExecute = new AtomicBoolean(true);
+    AtomicBoolean seenOnWorker = new AtomicBoolean(false);
+    Thread worker =
+        new Thread(
+            () -> {
+              defaultedBeforeExecute.set(CalcitePlanContext.isWarningsSupported());
+              query.execute();
+              seenOnWorker.set(CalcitePlanContext.isWarningsSupported());
+            });
+    worker.start();
+    worker.join();
+
+    assertFalse(
+        defaultedBeforeExecute.get(), "worker thread should default to no warnings support");
+    assertTrue(seenOnWorker.get(), "execute() must carry warningsSupported onto the worker thread");
+    verify(queryService, times(1)).execute(any(), any(), any(), anyBoolean(), any());
+  }
+
+  @Test
+  public void warnings_unsupported_plan_resets_flag_on_reused_worker_thread()
+      throws InterruptedException {
+    // Plan defaults to warningsSupported=false.
+    QueryPlan query = new QueryPlan(queryId, queryType, plan, queryService, queryListener);
+
+    AtomicBoolean seenOnWorker = new AtomicBoolean(true);
+    Thread worker =
+        new Thread(
+            () -> {
+              // Simulate a pooled worker left "supported" by a prior query.
+              CalcitePlanContext.setWarningsSupported(true);
+              query.execute();
+              seenOnWorker.set(CalcitePlanContext.isWarningsSupported());
+            });
+    worker.start();
+    worker.join();
+
+    assertFalse(
+        seenOnWorker.get(), "a warnings-unsupported plan must reset the flag on a reused thread");
+  }
+
   @Test
   public void explain_no_page_size() {
     QueryPlan query = new QueryPlan(queryId, queryType, plan, queryService, queryListener);

integ-test/src/test/java/org/opensearch/sql/security/PartialResultSecurityIT.java

@@ -0,0 +1,122 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.security;
+
+import static org.opensearch.sql.util.MatcherUtils.rows;
+import static org.opensearch.sql.util.MatcherUtils.verifyDataRows;
+import static org.opensearch.sql.util.TestUtils.createIndexByRestClient;
+import static org.opensearch.sql.util.TestUtils.isIndexExist;
+import static org.opensearch.sql.util.TestUtils.performRequest;
+
+import java.io.IOException;
+import org.json.JSONArray;
+import org.json.JSONObject;
+import org.junit.After;
+import org.junit.Test;
+import org.opensearch.client.Request;
+import org.opensearch.sql.common.setting.Settings;
+
+/**
+ * Runs the partial-result-on-mapping-conflict path with the security plugin installed. Regression
+ * guard for #5739: the warnings-supported gate used to live in Log4j {@code ThreadContext}, which
+ * the security plugin's transport interceptor drops on the transport-to-worker handoff, so partial
+ * mode silently bailed and returned a complete result with no warning. The {@code
+ * integTestWithSecurity} suite never exercised this path -- it only runs {@code
+ * org.opensearch.sql.security.*}, and the partial-result IT lives elsewhere -- so only the release
+ * distribution's full-suite-with-security caught it. Placing this test in the security package
+ * closes that gap.
+ */
+public class PartialResultSecurityIT extends SecurityTestBase {
+
+  private static final String KEYWORD_INDEX = "partial_sec_keyword";
+  private static final String TEXT_INDEX = "partial_sec_text";
+  private static final String PATTERN = "partial_sec_*";
+
+  private static final String USER = "partial_sec_user";
+  private static final String ROLE = "partial_sec_role";
+
+  private boolean initialized = false;
+
+  @Override
+  protected void init() throws Exception {
+    super.init();
+    enableCalcite();
+    if (!initialized) {
+      createRoleWithIndexAccess(ROLE, PATTERN);
+      createUser(USER, ROLE);
+      createConflictIndices();
+      initialized = true;
+    }
+  }
+
+  @After
+  public void resetPartialResult() throws IOException {
+    setPartialResult(false);
+  }
+
+  private void createConflictIndices() throws IOException {
+    // env is an aggregatable keyword here...
+    if (!isIndexExist(client(), KEYWORD_INDEX)) {
+      String mapping = "{\"mappings\":{\"properties\":{\"env\":{\"type\":\"keyword\"}}}}";
+      createIndexByRestClient(client(), KEYWORD_INDEX, mapping);
+      Request bulk = new Request("POST", "/" + KEYWORD_INDEX + "/_bulk?refresh=true");
+      bulk.setJsonEntity(
+          "{\"index\":{}}\n{\"env\":\"prod\"}\n"
+              + "{\"index\":{}}\n{\"env\":\"prod\"}\n"
+              + "{\"index\":{}}\n{\"env\":\"dev\"}\n");
+      performRequest(client(), bulk);
+    }
+    // ...and bare text (no .keyword sub-field) here, so the field collapses to non-aggregatable.
+    if (!isIndexExist(client(), TEXT_INDEX)) {
+      String mapping = "{\"mappings\":{\"properties\":{\"env\":{\"type\":\"text\"}}}}";
+      createIndexByRestClient(client(), TEXT_INDEX, mapping);
+      Request bulk = new Request("POST", "/" + TEXT_INDEX + "/_bulk?refresh=true");
+      bulk.setJsonEntity(
+          "{\"index\":{}}\n{\"env\":\"prod\"}\n" + "{\"index\":{}}\n{\"env\":\"qa\"}\n");
+      performRequest(client(), bulk);
+    }
+  }
+
+  @Test
+  public void partialResultWarningSurvivesSecurityHandoff() throws IOException {
+    setPartialResult(true);
+    JSONObject result =
+        executeQueryAsUser(
+            String.format("source=%s | stats count() by env | sort env", PATTERN), USER);
+
+    // Only the aggregatable keyword index contributes (prod=2, dev=1); the text index is excluded.
+    verifyDataRows(result, rows(1, "dev"), rows(2, "prod"));
+
+    assertTrue(
+        "partial result must carry a warnings channel under security", result.has("warnings"));
+    JSONArray warnings = result.getJSONArray("warnings");
+    assertEquals(1, warnings.length());
+    JSONObject warning = warnings.getJSONObject(0);
+    assertEquals("PARTIAL_RESULT", warning.getString("type"));
+    assertTrue(
+        "warning should name the excluded text index",
+        warning.getString("detail").contains(TEXT_INDEX));
+  }
+
+  @Test
+  public void completeResultCarriesNoWarningWithSecurity() throws IOException {
+    setPartialResult(false);
+    JSONObject result =
+        executeQueryAsUser(
+            String.format("source=%s | stats count() by env | sort env", PATTERN), USER);
+    // 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("a complete result carries no warning", result.has("warnings"));
+  }
+
+  private void setPartialResult(boolean enabled) throws IOException {
+    updateClusterSettings(
+        new ClusterSetting(
+            "persistent",
+            Settings.Key.PARTIAL_RESULT_ON_MAPPING_CONFLICT.getKeyValue(),
+            Boolean.toString(enabled)));
+  }
+}

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java

@@ -504,8 +504,8 @@ public Void visitInputRef(RexInputRef ref) {
    * On a text/keyword mapping conflict, narrow the scan to the index subset where the group field
    * is aggregatable, push the aggregation over just that subset, and record a warning naming the
    * excluded indices. Only runs behind the opt-in setting and only when the response format can
-   * carry the warning ({@link QueryContext#isWarningsSupported}); returns {@code null} otherwise.
-   * {@code partitionFields} are the scan fields the group keys resolve to (see {@link
+   * carry the warning ({@link CalcitePlanContext#isWarningsSupported}); returns {@code null}
+   * otherwise. {@code partitionFields} are the scan fields the group keys resolve to (see {@link
    * #resolvePartitionFields}). Partitioning lives in {@link PartialResultAggregatePushdown}.
    */
   private AbstractRelNode tryPartialResultAggregate(
@@ -514,7 +514,7 @@ private AbstractRelNode tryPartialResultAggregate(
       return null;
     }
     // A format with no warnings channel (CSV/RAW/VIZ) must not silently drop indices.
-    if (!QueryContext.isWarningsSupported()) {
+    if (!CalcitePlanContext.isWarningsSupported()) {
       return null;
     }
     try {