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 a4e262692cc..3a8e06cc9bf 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 @@ -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 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 60c0da9d452..11c74edb4dc 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java @@ -73,6 +73,15 @@ public class CalcitePlanContext { private static final ThreadLocal> 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 warningsSupported = + ThreadLocal.withInitial(() -> false); + /** Thread-local switch that tells whether the current query prefers legacy behavior. */ private static final ThreadLocal legacyPreferredFlag = ThreadLocal.withInitial(() -> true); @@ -261,6 +270,7 @@ 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. */ @@ -268,6 +278,19 @@ 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( 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 fbdabe2fa44..2b12c4f9779 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 @@ -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(); 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 d78036f0faa..904b35a45be 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 @@ -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), diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java index 3220c5d28f7..0080f17c177 100644 --- a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java +++ b/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); 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 new file mode 100644 index 00000000000..f925b396851 --- /dev/null +++ b/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))); + } +} 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 873792a1787..1a84935a5d1 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 @@ -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 { 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 e8147f33ad6..660ecb23af6 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 @@ -195,8 +195,9 @@ protected void doExecute( PPLQueryRequest transformedRequest = transportRequest.toPPLQueryRequest(); QueryContext.setProfile(transformedRequest.profile()); // Only the JSON shape carries warnings; gate partial results on it so CSV/RAW/VIZ never drop - // data silently. - QueryContext.setWarningsSupported(warningsSupported(transformedRequest)); + // 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()); @@ -452,6 +453,5 @@ public void onFailure(Exception e) { private static void clearRequestScopedState() { QueryProfiling.clear(); QueryContext.setPartialResultOverride(null); - QueryContext.setWarningsSupported(false); } } 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 382888274ef..54d57370168 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java @@ -209,6 +209,8 @@ private AbstractPlan plan( log.info("[{}] Incoming request {}", QueryContext.getRequestId(), anonymized); anonymizedQuerySink.accept(anonymized); - return queryExecutionFactory.create(statement, queryListener, explainListener); + AbstractPlan plan = queryExecutionFactory.create(statement, queryListener, explainListener); + plan.setWarningsSupported(request.warningsSupported()); + return plan; } } diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java b/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java index b412c7b828b..d8fa695de3c 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/domain/PPLQueryRequest.java @@ -72,6 +72,16 @@ public class PPLQueryRequest { @Accessors(fluent = true) private Boolean partialResult = null; + /** + * Whether the requested response format can carry a warnings channel. Derived from the format on + * the transport thread and threaded to the worker via the plan, so the partial-result gate does + * not rely on Log4j ThreadContext (dropped across the security plugin's thread handoff). + */ + @Setter + @Getter + @Accessors(fluent = true) + private boolean warningsSupported = false; + public PPLQueryRequest(String pplQuery, JSONObject jsonContent, String path) { this(pplQuery, jsonContent, path, ""); }