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

[ahkcs]

Description

Fixes the with-security-only integ-test failure in #5739 (https://github.com/opensearch-project/sql/issues/5739#issuecomment-5516921306)(`CalcitePartialResultOnMappingConflictIT`, 7 methods). The partial-result feature was added in #5657; this is a regression in that feature's warning channel under the security plugin.

QueryContext.isWarningsSupported() was backed by Log4j ThreadContext and set once on the transport thread (TransportPPLQueryAction). The partial-result gate reads it on the planning thread (CalciteLogicalIndexScan.tryPartialResultAggregate). Under the security plugin, the transport→worker handoff does not preserve Log4j's ThreadContext, so the flag read false at planning time, tryPartialResultAggregate bailed, and the query returned the complete result with no warnings — the exact symptom in the failing tests (extra text-index buckets; missing warnings). It reproduces only in the release distribution's security config; without-security passed, which is why the plugin's own integTestWithSecurity (which only runs org.opensearch.sql.security.*, and the partial-result IT lives elsewhere) went green on the original PR.

Fix

Carry warnings-support off the thread-local and onto the object graph, so no thread handoff can drop it:

PPLQueryRequest.warningsSupported (set from the response format on the transport thread) → AbstractPlan → QueryPlan.execute() (runs on the worker thread) → CalcitePlanContext.setWarningsSupported(...), which the scan already reads. It's also included in CalcitePlanContext's thread-local snapshot (so it survives the planner's own handoff) and reset per query. The now-dead QueryContext warnings-supported methods are removed. isPartialResultEnabled was already resilient (it falls back to the cluster setting), so only warningsSupported needed this.

Testing

• QueryPlanTest: two new cross-thread tests prove the flag rides the plan object and is honored when the configuring thread ≠ the execute() thread (the exact bug class), and that a warnings-unsupported plan resets the flag on a reused pooled worker thread.
• PartialResultAggregatePushdownTest (19) and :opensearch-sql-plugin:compileJava pass.
• PartialResultSecurityIT (new, under org.opensearch.sql.security) exercises the partial-result path with the security plugin installed, so CI's integTestWithSecurity now covers it — closing the gap that let this regression ship (that suite only runs org.opensearch.sql.security.*, and the feature IT lives under calcite/remote). Verified locally against the secure testcluster: it passes on this fix and fails on the pre-fix code (expected:<2> but was:<3> — the complete result leaks the excluded index's rows and drops the warnings), so it is a true differential guard.

Check List

• New functionality includes testing.
• Commits are signed per the DCO.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 9180b99)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ No major issues detected

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 9180b99

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Use remove() for ThreadLocal cleanup The warningsSupported ThreadLocal should be removed using remove() instead of set(false) to prevent memory leaks in thread pools. Setting to false leaves the ThreadLocal entry in the thread's map, while remove() cleans it up completely. core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java [267-274] public static void clearTimewrapSignals() { stripNullColumns.set(false); timewrapUnitName.set(null); timewrapSeries.set(null); executionPool.set(null); pendingWarnings.remove(); - warningsSupported.set(false); + warningsSupported.remove(); } Suggestion importance[1-10]: 9 __ Why: Using remove() instead of set(false) for ThreadLocal cleanup is critical to prevent memory leaks in thread pools. The current implementation leaves the ThreadLocal entry in the thread's map, which can accumulate over time in pooled environments.	High
	Ensure ThreadLocal cleanup after execution The warningsSupported flag is set at the start of execute() but never cleared afterward. If the worker thread is reused for subsequent queries, the flag may persist incorrectly. Consider clearing it in a finally block or ensuring cleanup after query execution completes. core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java [108-117] @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), - getQueryType(), + try { + if (pageSize.isPresent()) { + queryService.execute( + new Paginate(pageSize.get(), plan), + getQueryType(), + } finally { + CalcitePlanContext.setWarningsSupported(false); + } Suggestion importance[1-10]: 4 __ Why: While the concern about ThreadLocal cleanup is valid, the existing test warnings_unsupported_plan_resets_flag_on_reused_worker_thread demonstrates that the flag is properly reset at the start of each execute() call. The suggestion adds defensive cleanup but may not be necessary given the existing reset mechanism.	Low

---

Previous suggestions

Suggestions up to commit ca22496

Category	Suggestion	Impact
Possible issue	Ensure flag cleanup on execution failure The warningsSupported flag is set on the worker thread but never explicitly cleared after query execution completes. This could cause the flag to leak to subsequent queries on the same pooled thread if execution fails before reaching cleanup code. Consider wrapping the execution logic in a try-finally block to ensure the flag is always reset. core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java [108-116] @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), - getQueryType(), + try { + if (pageSize.isPresent()) { + queryService.execute( + new Paginate(pageSize.get(), plan), + getQueryType(), + ... + } + } finally { + CalcitePlanContext.setWarningsSupported(false); + } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies that warningsSupported is set but not explicitly cleared in execute(). However, the PR shows clearTimewrapSignals() at line 273 already resets this flag, and the test at lines 90-109 verifies the reset behavior works. The concern about exception handling is valid but may be addressed elsewhere in the execution flow.	Medium
General	Remove ThreadLocal to prevent memory leaks The warningsSupported ThreadLocal is set to false instead of being removed like pendingWarnings. For consistency and to prevent potential memory leaks in long-running applications with thread pools, consider using remove() instead of set(false) to fully clean up the ThreadLocal entry. core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java [267-274] public static void clearTimewrapSignals() { stripNullColumns.set(false); timewrapUnitName.set(null); timewrapSeries.set(null); executionPool.set(null); pendingWarnings.remove(); - warningsSupported.set(false); + warningsSupported.remove(); } Suggestion importance[1-10]: 6 __ Why: The suggestion raises a valid point about ThreadLocal cleanup consistency. Using remove() instead of set(false) would be more consistent with how pendingWarnings is handled and could prevent potential memory leaks. However, the impact is relatively minor since the ThreadLocal has a primitive wrapper initializer that would be recreated on next access.	Low

Suggestions up to commit 509cc4b

Category	Suggestion	Impact
Possible issue	Use remove() to prevent memory leak The warningsSupported ThreadLocal should be removed using warningsSupported.remove() instead of warningsSupported.set(false) to prevent memory leaks in thread pools. Setting to false leaves the ThreadLocal entry in the thread's map, while remove() cleans it up completely. core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java [273] public static void clearTimewrapSignals() { stripNullColumns.set(false); timewrapUnitName.set(null); timewrapSeries.set(null); executionPool.set(null); pendingWarnings.remove(); - warningsSupported.set(false); + warningsSupported.remove(); } Suggestion importance[1-10]: 8 __ Why: Using warningsSupported.remove() instead of warningsSupported.set(false) is important for preventing memory leaks in thread pools. The remove() method properly cleans up the ThreadLocal entry from the thread's map, while set(false) leaves the entry in place. This is consistent with how pendingWarnings is already being cleaned up on line 272.	Medium
General	Avoid autoboxing in ThreadLocal initialization Consider using ThreadLocal.withInitial(() -> Boolean.FALSE) instead of () -> false to avoid unnecessary autoboxing on every initialization. This improves performance when the ThreadLocal is frequently initialized across pooled threads. core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java [82-83] private static final ThreadLocal<Boolean> warningsSupported = - ThreadLocal.withInitial(() -> false); + ThreadLocal.withInitial(() -> Boolean.FALSE); Suggestion importance[1-10]: 3 __ Why: While using Boolean.FALSE instead of false avoids autoboxing during initialization, the performance impact is minimal since ThreadLocal initialization happens infrequently. This is a minor optimization that improves code consistency but has negligible practical impact.	Low

Suggestions up to commit ca22496

Category	Suggestion	Impact
Possible issue	Ensure flag cleanup on exceptions The warningsSupported flag is set on the worker thread but never explicitly cleared after query execution completes. This could cause the flag to leak to subsequent queries on the same pooled thread if an exception occurs before clearTimewrapSignals() is called. Wrap the execution logic in a try-finally block to ensure the flag is always reset. core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java [108-117] @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), - getQueryType(), + try { + if (pageSize.isPresent()) { + queryService.execute( + new Paginate(pageSize.get(), plan), + getQueryType(), + ... + } + } finally { + CalcitePlanContext.setWarningsSupported(false); + } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential resource leak where warningsSupported could persist on pooled threads if an exception occurs. However, the PR already includes clearTimewrapSignals() at line 273 which resets warningsSupported.set(false), and the test at lines 90-109 validates that the flag is properly reset. The suggestion adds defensive programming but may be redundant given existing cleanup mechanisms.	Medium

[github-actions]

Persistent review updated to latest commit 509cc4b

[github-actions]

Persistent review updated to latest commit ca22496

[RyanL1997]

Since we didnt catch this under the CI with security can we also add these cases under integ-test/src/test/java/org/opensearch/sql/security/?

[github-actions]

Persistent review updated to latest commit 9180b99

[ahkcs]

Since we didnt catch this under the CI with security can we also add these cases under integ-test/src/test/java/org/opensearch/sql/security/?

Added PartialResultSecurityIT under integ-test/src/test/java/org/opensearch/sql/security/, so integTestWithSecurity now exercises the partial-result path with the security plugin. That's the exact gap that hid this — the suite only globs org.opensearch.sql.security.* and the feature IT lives under calcite/remote. Verified locally on the secure testcluster: passes on this fix, fails on the pre-fix code (expected:<2> but was:<3> — excluded index's rows leak in and the warning drops), so it's a real differential guard.