From 64bcd083a6917fa060dacd6711ee792da492ca31 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Thu, 10 Sep 2026 15:11:55 -0700 Subject: [PATCH 01/26] feat(ppl): Report completed PPL queries to Query Insights Top N Emit a Query Insights record for every completed PPL query so it appears in the Top N overview alongside DSL queries. Each record carries the query source (PPL), the anonymized query text, the resolved indices (extracted from the parsed AST), the requesting user on a security-enabled cluster, and query-total latency, CPU, and memory from the task resource-tracking framework. Child DSL searches spawned by a PPL query (e.g. a join's per-table scans) are tagged with a parent marker (X-Query-Insights-Parent) so Query Insights can associate them with their parent as sub-queries. The marker is propagated on the OpenSearch ThreadContext header across the engine's worker and complex-worker thread hops and re-applied as a task header on the background scan pool. The report is sent to the local Query Insights node as an internal, system-initiated request; the send stashes the thread context so it is not authorized against the end user's roles on an FGAC domain. This is reporting only: no phase breakdown and no profiling changes. The profiling subsystem is left as-is for a future change. Signed-off-by: Kishore Kumaar Natarajan --- docs/user/ppl/admin/settings.md | 32 +++- .../sql/ppl/QueryInsightsReportingIT.java | 63 +++++++ .../QueryInsightsReportingSecurityIT.java | 81 +++++++++ .../executor/OpenSearchQueryManager.java | 70 ++++++++ .../ThreadPoolExecutionDispatcher.java | 14 ++ .../executor/OpenSearchQueryManagerTest.java | 159 ++++++++++++++++++ .../ThreadPoolExecutionDispatcherTest.java | 150 ++++++++++++++++- .../org/opensearch/sql/plugin/SQLPlugin.java | 11 ++ .../sql/plugin/transport/PPLQueryTask.java | 58 +++++++ .../plugin/transport/QueryInsightsMarker.java | 24 +++ .../transport/QueryInsightsReporter.java | 89 ++++++++++ .../transport/TransportPPLQueryAction.java | 159 +++++++++++++++++- .../plugin/transport/PPLQueryTaskTest.java | 36 ++++ .../org/opensearch/sql/ppl/PPLService.java | 60 ++++--- .../sql/ppl/QueryInsightsMetadata.java | 17 ++ .../sql/ppl/utils/PPLQueryIndexExtractor.java | 158 +++++++++++++++++ .../opensearch/sql/ppl/PPLServiceTest.java | 41 +++++ .../ppl/utils/PPLQueryIndexExtractorTest.java | 87 ++++++++++ 18 files changed, 1276 insertions(+), 33 deletions(-) create mode 100644 integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsReportingIT.java create mode 100644 integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsReportingSecurityIT.java create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsMarker.java create mode 100644 plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java create mode 100644 ppl/src/main/java/org/opensearch/sql/ppl/QueryInsightsMetadata.java create mode 100644 ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java create mode 100644 ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java diff --git a/docs/user/ppl/admin/settings.md b/docs/user/ppl/admin/settings.md index b1d5eef17fa..28986cd732d 100644 --- a/docs/user/ppl/admin/settings.md +++ b/docs/user/ppl/admin/settings.md @@ -472,4 +472,34 @@ Expected output: "transient": {} } ``` - \ No newline at end of file + + +## search.insights.top_queries.ppl.enabled + +### Version + +3.9 + +### Description + +This setting controls whether completed PPL queries are recorded into Query Insights Top N. When enabled, the PPL plugin reports each finished PPL query (with its resolved indices, latency, CPU, memory, and — on a security-enabled cluster — the requesting user) to the Query Insights plugin, and tags the child DSL searches a PPL query spawns so they appear as sub-queries. + +1. The default value is false (opt-in). +2. This setting is node scope, matching the rest of the `search.insights.top_queries.*` family. +3. It is dynamic: update it through the cluster settings API (as in the example below), and the value applies per node. + +Note on user identity: when a security plugin is installed, the recorded record includes the requesting user (username, roles, backend roles). Who can see that identity in Top N results is governed by the Query Insights plugin's own RBAC filter, `search.insights.top_queries.filter_by_mode` — default `none` (no per-user filtering), `username` (a user sees only their own queries; admins see all), or `backend_roles`. Set it to `username` or `backend_roles` if user identities should not be visible to every Top N reader. + +Note: This setting is **registered by the Query Insights plugin** (it belongs to the `search.insights.top_queries.*` family), not by the PPL plugin. The PPL plugin reads its value to decide whether to report. It therefore only takes effect when the Query Insights plugin is installed; when Query Insights is absent the setting is unregistered and PPL reporting is a no-op. + +### Example + +With the Query Insights plugin installed, you can enable PPL recording like this: + +```shell +curl -sS -H 'Content-Type: application/json' \ +-X PUT localhost:9200/_cluster/settings \ +-d '{"persistent" : {"search.insights.top_queries.ppl.enabled" : "true"}}' +``` + +This example is not executed as part of the documentation tests because the setting is registered by the Query Insights plugin, which is not installed in the PPL documentation test cluster. diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsReportingIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsReportingIT.java new file mode 100644 index 00000000000..1e8fe2e9038 --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsReportingIT.java @@ -0,0 +1,63 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ppl; + +import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK; +import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_DOG; + +import java.io.IOException; +import org.json.JSONObject; +import org.junit.Test; +import org.opensearch.client.ResponseException; +import org.opensearch.sql.legacy.TestUtils; + +/** + * Verifies that reporting completed PPL queries to Query Insights does not interfere with query + * execution. + * + *

Query Insights is a separate plugin and is not installed in the SQL integ-test cluster. The + * reporting gate ({@code isQueryInsightsRecordingEnabled}) therefore returns false and the whole + * reporting path is skipped, so these tests assert the important contract: with Query Insights + * absent, plain and multi-scan (join) PPL queries still run and return correct results. The + * end-to-end assertion that records and sub-queries land in Query Insights lives with the Query + * Insights plugin, which can co-install both plugins. + */ +public class QueryInsightsReportingIT extends PPLIntegTestCase { + + @Override + public void init() throws Exception { + super.init(); + enableCalcite(); // join runs on the Calcite path + loadIndex(Index.BANK); + loadIndex(Index.DOG); + } + + @Test + public void singleSourceQueryReturnsResultsWithReportingPathActive() throws IOException { + JSONObject result = + executeQuery( + String.format("source=%s | where age > 30 | fields firstname", TEST_INDEX_BANK)); + assertTrue(result.getJSONArray("datarows").length() > 0); + } + + @Test + public void joinQuerySpawningChildScansReturnsResults() throws IOException { + // A join fans out into per-side child scans on the background pool — the path that stamps the + // parent marker. It must still return results when Query Insights is not installed. + JSONObject result; + try { + result = + executeQuery( + String.format( + "source=%s | join on firstname=holdersName %s", TEST_INDEX_BANK, TEST_INDEX_DOG)); + } catch (ResponseException e) { + // Surface the server error body to make failures debuggable rather than opaque. + throw new AssertionError( + "join query failed: " + TestUtils.getResponseBody(e.getResponse()), e); + } + assertTrue(result.has("datarows")); + } +} diff --git a/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsReportingSecurityIT.java b/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsReportingSecurityIT.java new file mode 100644 index 00000000000..0ec7903d63c --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsReportingSecurityIT.java @@ -0,0 +1,81 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.security; + +import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK; +import static org.opensearch.sql.util.MatcherUtils.columnName; +import static org.opensearch.sql.util.MatcherUtils.verifyColumn; + +import java.io.IOException; +import org.json.JSONObject; +import org.junit.Test; + +/** + * Security tests for the PPL → Query Insights reporting path on an FGAC (security-enabled) cluster. + * + *

Reporting a completed PPL query is an internal {@code cluster:admin/...} action the end user + * never invokes; the send is wrapped in {@code stashContext()} so it runs user-less rather than + * being authorized against the querying user (see {@code QueryInsightsReporter}). These tests pin + * that contract: a low-privilege user who holds only PPL + index-read permissions — and explicitly + * NOT the report action's {@code cluster:admin} permission — can still run PPL queries, including + * the multi-scan (join) path that stamps the parent marker for child sub-query tagging. + * + *

Read-side visibility of the recorded user identity is governed by the Query Insights plugin's + * own RBAC filter ({@code search.insights.top_queries.filter_by_mode}); that plugin is not + * installed in the SQL integ-test cluster, so it is covered by Query Insights' own tests rather + * than here. + */ +public class QueryInsightsReportingSecurityIT extends SecurityTestBase { + + private static final String PPL_USER = "qi_ppl_user"; + private static final String PPL_ROLE = "qi_ppl_role"; + + private boolean initialized = false; + + @Override + protected void init() throws Exception { + super.init(); + if (!initialized) { + // PPL + index read/mapping/settings, but deliberately NO report action permission + // (cluster:admin/opensearch/query_insights/report_query_bytes) and no all_access. + createRoleWithPermissions( + PPL_ROLE, + TEST_INDEX_BANK, + new String[] {"cluster:admin/opensearch/ppl"}, + new String[] { + "indices:data/read/search*", + "indices:admin/mappings/get", + "indices:monitor/settings/get", + "indices:data/read/point_in_time/create", + "indices:data/read/point_in_time/delete" + }); + createUser(PPL_USER, PPL_ROLE); + loadIndex(Index.BANK); + enableCalcite(); + allowCalciteFallback(); + initialized = true; + } + } + + @Test + public void lowPrivilegeUserCanRunPplQueryWithoutReportPermission() throws IOException { + // The report send is stashed to Origin.LOCAL, so lacking the report cluster:admin permission + // must not cause the user's own query to fail. + JSONObject result = + executeQueryAsUser( + String.format("source=%s | where age > 30 | fields firstname", TEST_INDEX_BANK), + PPL_USER); + verifyColumn(result, columnName("firstname")); + } + + @Test + public void lowPrivilegeUserCanRunStatsQueryWithoutReportPermission() throws IOException { + JSONObject result = + executeQueryAsUser( + String.format("source=%s | stats count() by gender", TEST_INDEX_BANK), PPL_USER); + verifyColumn(result, columnName("gender"), columnName("count()")); + } +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java index c391153fca6..def9d195d5b 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java @@ -5,6 +5,8 @@ package org.opensearch.sql.opensearch.executor; +import com.sun.management.ThreadMXBean; +import java.lang.management.ManagementFactory; import java.util.Map; import lombok.RequiredArgsConstructor; import org.apache.logging.log4j.LogManager; @@ -12,6 +14,9 @@ import org.apache.logging.log4j.ThreadContext; import org.opensearch.OpenSearchTimeoutException; import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.tasks.resourcetracker.ResourceStats; +import org.opensearch.core.tasks.resourcetracker.ResourceStatsType; +import org.opensearch.core.tasks.resourcetracker.ResourceUsageMetric; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.executor.QueryId; import org.opensearch.sql.executor.QueryManager; @@ -27,6 +32,20 @@ public class OpenSearchQueryManager implements QueryManager { private static final Logger LOG = LogManager.getLogger(OpenSearchQueryManager.class); + /** Samples per-thread CPU/memory for resource tracking; null when the JVM bean is unavailable. */ + private static final ThreadMXBean THREAD_MX_BEAN = resolveThreadMXBean(); + + private static ThreadMXBean resolveThreadMXBean() { + try { + if (ManagementFactory.getThreadMXBean() instanceof ThreadMXBean bean) { + return bean; + } + } catch (Exception e) { + LOG.warn("Per-thread resource metrics unavailable; PPL task resource tracking disabled", e); + } + return null; + } + private final NodeClient nodeClient; private final Settings settings; @@ -81,6 +100,14 @@ private void schedule( setCancellableTask(cancelTask); + // Bracket resource tracking for the inline path (runs on this thread). Script plans + // hand off to the complex-worker pool, which brackets itself. + final boolean trackResources = + cancelTask != null && cancelTask.supportsResourceTracking(); + final long trackedThreadId = Thread.currentThread().getId(); + final boolean trackingStarted = + trackResources && startThreadResourceTracking(cancelTask, trackedThreadId); + try { task.run(); timeoutTask.cancel(); @@ -98,6 +125,9 @@ private void schedule( throw e; } finally { + if (trackingStarted) { + stopThreadResourceTracking(cancelTask, trackedThreadId); + } clearCancellableTask(); } }); @@ -112,4 +142,44 @@ private Runnable withCurrentContext(final Runnable task) { task.run(); }; } + + /** + * Records the starting CPU/memory snapshot for {@code threadId}. + * + * @return true if tracking started; only then should {@link #stopThreadResourceTracking} be + * called. + */ + static boolean startThreadResourceTracking(CancellableTask task, long threadId) { + try { + task.startThreadResourceTracking( + threadId, ResourceStatsType.WORKER_STATS, currentThreadResourceMetrics(threadId)); + return true; + } catch (Exception e) { + LOG.warn("Failed to start resource tracking for task [{}]", task.getId(), e); + return false; + } + } + + /** Records the final CPU/memory snapshot for {@code threadId}. */ + static void stopThreadResourceTracking(CancellableTask task, long threadId) { + try { + task.stopThreadResourceTracking( + threadId, ResourceStatsType.WORKER_STATS, currentThreadResourceMetrics(threadId)); + } catch (Exception e) { + LOG.warn("Failed to stop resource tracking for task [{}]", task.getId(), e); + } + } + + /** Per-thread memory and CPU usage; empty array when the JVM bean is unavailable. */ + private static ResourceUsageMetric[] currentThreadResourceMetrics(long threadId) { + if (THREAD_MX_BEAN == null) { + return new ResourceUsageMetric[0]; + } + ResourceUsageMetric memory = + new ResourceUsageMetric( + ResourceStats.MEMORY, THREAD_MX_BEAN.getThreadAllocatedBytes(threadId)); + ResourceUsageMetric cpu = + new ResourceUsageMetric(ResourceStats.CPU, THREAD_MX_BEAN.getThreadCpuTime(threadId)); + return new ResourceUsageMetric[] {memory, cpu}; + } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java index 3de8bb85c2b..e8f86ae6f36 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java @@ -91,6 +91,9 @@ private void dispatchInternal( ThreadPool.Names.GENERIC); Cancellable cancelPoller = scheduleCancellationPoller(cancellableTask, executionThread); Hook.Closeable hookHandle = null; + boolean trackResources = false; + long trackedThreadId = -1L; + boolean trackingStarted = false; try { // Restore state from caller thread ThreadContext.putAll(ctx); @@ -106,6 +109,14 @@ private void dispatchInternal( hookHandle = Hook.CURRENT_TIME.addThread((Consumer>) h -> h.set(currentTime)); } + // Script plans do their real work on this thread, so bracket tracking here. + trackResources = + cancellableTask != null && cancellableTask.supportsResourceTracking(); + trackedThreadId = Thread.currentThread().getId(); + trackingStarted = + trackResources + && OpenSearchQueryManager.startThreadResourceTracking( + cancellableTask, trackedThreadId); task.run(); } catch (Exception e) { LOG.error("Exception during task execution on complex pool", e); @@ -119,6 +130,9 @@ private void dispatchInternal( if (hookHandle != null) { hookHandle.close(); } + if (trackingStarted) { + OpenSearchQueryManager.stopThreadResourceTracking(cancellableTask, trackedThreadId); + } OpenSearchQueryManager.clearCancellableTask(); RelMetadataQueryBase.THREAD_PROVIDERS.remove(); CalcitePlanContext.clearTimewrapSignals(); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java index 1463cf48fff..237ac21ad7f 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java @@ -5,18 +5,30 @@ package org.opensearch.sql.opensearch.executor; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import java.util.Collections; +import java.util.List; +import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.core.tasks.resourcetracker.ResourceStatsType; +import org.opensearch.core.tasks.resourcetracker.ResourceUsageMetric; +import org.opensearch.core.tasks.resourcetracker.ThreadResourceInfo; import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; @@ -26,6 +38,7 @@ import org.opensearch.sql.executor.QueryType; import org.opensearch.sql.executor.execution.AbstractPlan; import org.opensearch.sql.executor.execution.QueryPlan; +import org.opensearch.tasks.CancellableTask; import org.opensearch.threadpool.Scheduler; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.client.node.NodeClient; @@ -51,6 +64,7 @@ public void submitQuery() { Scheduler.ScheduledCancellable mockScheduledTask = mock(Scheduler.ScheduledCancellable.class); when(nodeClient.threadPool()).thenReturn(threadPool); + when(settings.getSettingValue(Settings.Key.PPL_QUERY_TIMEOUT)) .thenReturn(TimeValue.timeValueSeconds(60)); @@ -76,4 +90,149 @@ public void execute() { assertTrue(isRun.get()); } + + @AfterEach + public void clearTask() { + OpenSearchQueryManager.clearCancellableTask(); + } + + @Test + public void tracksResourceUsageWhenTaskSupportsIt() { + TrackingTask task = new TrackingTask(true); + OpenSearchQueryManager.setCancellableTask(task); + + runSubmit(); + + Map> stats = task.getResourceStats(); + assertEquals(1, stats.size()); + List infos = stats.values().iterator().next(); + assertEquals(1, infos.size()); + assertFalse("resource entry should be closed after execution", infos.get(0).isActive()); + } + + @Test + public void skipsResourceTrackingWhenTaskDoesNotSupportIt() { + TrackingTask task = new TrackingTask(false); + OpenSearchQueryManager.setCancellableTask(task); + + runSubmit(); + + assertTrue("no resource entries expected", task.getResourceStats().isEmpty()); + } + + @Test + public void propagatesExecutionException() { + RuntimeException boom = new RuntimeException("execute failed"); + RuntimeException thrown = + org.junit.jupiter.api.Assertions.assertThrows( + RuntimeException.class, () -> runSubmitThatThrows(boom)); + assertEquals(boom, thrown); + } + + /** Runs a trivial query through the manager with the schedule hook executing tasks inline. */ + private void runSubmit() { + NodeClient nodeClient = mock(NodeClient.class); + ThreadPool threadPool = mock(ThreadPool.class); + Settings settings = mock(Settings.class); + Scheduler.ScheduledCancellable mockScheduledTask = mock(Scheduler.ScheduledCancellable.class); + + when(nodeClient.threadPool()).thenReturn(threadPool); + + when(settings.getSettingValue(Settings.Key.PPL_QUERY_TIMEOUT)) + .thenReturn(TimeValue.timeValueSeconds(60)); + + AbstractPlan queryPlan = + new QueryPlan(queryId, queryType, plan, queryService, listener) { + @Override + public void execute() {} + }; + + doAnswer( + invocation -> { + Runnable task = invocation.getArgument(0); + task.run(); + return mockScheduledTask; + }) + .when(threadPool) + .schedule(any(), any(), any()); + + new OpenSearchQueryManager(nodeClient, settings).submit(queryPlan); + } + + @Test + public void startThreadResourceTrackingReturnsFalseOnFailure() { + CancellableTask task = mock(CancellableTask.class); + doThrow(new RuntimeException("bean failure")) + .when(task) + .startThreadResourceTracking( + anyLong(), any(ResourceStatsType.class), any(ResourceUsageMetric[].class)); + assertFalse(OpenSearchQueryManager.startThreadResourceTracking(task, 1L)); + } + + @Test + public void stopThreadResourceTrackingSwallowsFailure() { + CancellableTask task = mock(CancellableTask.class); + doThrow(new RuntimeException("bean failure")) + .when(task) + .stopThreadResourceTracking( + anyLong(), any(ResourceStatsType.class), any(ResourceUsageMetric[].class)); + OpenSearchQueryManager.stopThreadResourceTracking(task, 1L); + } + + /** Runs a query whose execute() throws, so the wrapped task's catch/finally path is exercised. */ + private void runSubmitThatThrows(RuntimeException toThrow) { + NodeClient nodeClient = mock(NodeClient.class); + ThreadPool threadPool = mock(ThreadPool.class); + Settings settings = mock(Settings.class); + Scheduler.ScheduledCancellable mockScheduledTask = mock(Scheduler.ScheduledCancellable.class); + + when(nodeClient.threadPool()).thenReturn(threadPool); + + when(settings.getSettingValue(Settings.Key.PPL_QUERY_TIMEOUT)) + .thenReturn(TimeValue.timeValueSeconds(60)); + + AbstractPlan queryPlan = + new QueryPlan(queryId, queryType, plan, queryService, listener) { + @Override + public void execute() { + throw toThrow; + } + }; + + // Run only the worker task inline; do NOT run the timeout task (running it would interrupt the + // thread and turn the thrown exception into an OpenSearchTimeoutException). + doAnswer( + invocation -> { + Object executor = invocation.getArgument(2); + if ("sql-worker".equals(executor)) { + Runnable task = invocation.getArgument(0); + task.run(); + } + return mockScheduledTask; + }) + .when(threadPool) + .schedule(any(), any(), any()); + + new OpenSearchQueryManager(nodeClient, settings).submit(queryPlan); + } + + /** Minimal CancellableTask whose resource-tracking support is configurable. */ + private static class TrackingTask extends CancellableTask { + private final boolean supportsTracking; + + TrackingTask(boolean supportsTracking) { + super(1L, "ppl", "action", "desc", TaskId.EMPTY_TASK_ID, Collections.emptyMap()); + this.supportsTracking = supportsTracking; + } + + @Override + public boolean supportsResourceTracking() { + return supportsTracking; + } + + @Override + public boolean shouldCancelChildrenOnCancellation() { + return true; + } + } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java index bf6adbf985e..510214cda4a 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java @@ -5,10 +5,15 @@ package org.opensearch.sql.opensearch.executor; +import static org.apache.logging.log4j.ThreadContext.clearAll; +import static org.apache.logging.log4j.ThreadContext.get; +import static org.apache.logging.log4j.ThreadContext.put; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doThrow; @@ -16,6 +21,7 @@ import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import static org.opensearch.common.settings.Settings.EMPTY; import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_COMPLEX_WORKER_THREAD_POOL_NAME; import java.util.List; @@ -24,7 +30,6 @@ import org.apache.calcite.rel.RelVisitor; import org.apache.calcite.rel.metadata.JaninoRelMetadataProvider; import org.apache.calcite.rel.metadata.RelMetadataQueryBase; -import org.apache.logging.log4j.ThreadContext; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -34,6 +39,9 @@ import org.mockito.junit.jupiter.MockitoSettings; import org.mockito.quality.Strictness; import org.opensearch.common.unit.TimeValue; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.tasks.resourcetracker.ResourceStatsType; +import org.opensearch.core.tasks.resourcetracker.ResourceUsageMetric; import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; @@ -60,6 +68,7 @@ class ThreadPoolExecutionDispatcherTest { @BeforeEach void setUp() { dispatcher = new ThreadPoolExecutionDispatcher(threadPool, settings); + when(threadPool.getThreadContext()).thenReturn(new ThreadContext(EMPTY)); // Mock schedule calls to return non-null cancellables (for both outer dispatch and inner // timeout) when(threadPool.schedule(any(Runnable.class), any(TimeValue.class), any())) @@ -70,7 +79,7 @@ void setUp() { @AfterEach void tearDown() { - ThreadContext.clearAll(); + clearAll(); OpenSearchQueryManager.clearCancellableTask(); RelMetadataQueryBase.THREAD_PROVIDERS.remove(); CalcitePlanContext.clearTimewrapSignals(); @@ -175,8 +184,8 @@ void propagatesLog4jThreadContextToSlowPool() { .thenReturn(true); when(settings.getSettingValue(Settings.Key.PPL_QUERY_TIMEOUT)) .thenReturn(new TimeValue(60000)); - ThreadContext.put("request.id", "test-123"); - ThreadContext.put("user", "admin"); + put("request.id", "test-123"); + put("user", "admin"); AtomicReference requestIdOnSlowPool = new AtomicReference<>(); AtomicReference userOnSlowPool = new AtomicReference<>(); @@ -184,10 +193,10 @@ void propagatesLog4jThreadContextToSlowPool() { invocation -> { Runnable task = invocation.getArgument(0); // Simulate a different thread — clear MDC - ThreadContext.clearAll(); + clearAll(); task.run(); - requestIdOnSlowPool.set(ThreadContext.get("request.id")); - userOnSlowPool.set(ThreadContext.get("user")); + requestIdOnSlowPool.set(get("request.id")); + userOnSlowPool.set(get("user")); return mock(ScheduledCancellable.class); }) .when(threadPool) @@ -368,6 +377,133 @@ void cancellableTaskAvailableDuringExecution() { assertEquals(mockTask, taskDuringExecution.get()); } + @Test + void bracketsResourceTrackingOnComplexPool() { + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) + .thenReturn(true); + when(settings.getSettingValue(Settings.Key.PPL_QUERY_TIMEOUT)) + .thenReturn(new TimeValue(60000)); + CancellableTask trackingTask = mock(CancellableTask.class); + when(trackingTask.supportsResourceTracking()).thenReturn(true); + OpenSearchQueryManager.setCancellableTask(trackingTask); + + doAnswer( + invocation -> { + Object executor = invocation.getArgument(2); + if (SQL_COMPLEX_WORKER_THREAD_POOL_NAME.equals(executor)) { + OpenSearchQueryManager.clearCancellableTask(); + invocation.getArgument(0).run(); + } + return mock(ScheduledCancellable.class); + }) + .when(threadPool) + .schedule(any(Runnable.class), any(TimeValue.class), any()); + + AbstractCalciteIndexScan scan = createMockScanWithScripts(); + dispatcher.dispatch(scan, context, listener, engine); + + // start/stopThreadResourceTracking take (threadId, statsType, ResourceUsageMetric...) varargs. + verify(trackingTask).supportsResourceTracking(); + verify(trackingTask) + .startThreadResourceTracking( + anyLong(), any(ResourceStatsType.class), any(ResourceUsageMetric[].class)); + verify(trackingTask) + .stopThreadResourceTracking( + anyLong(), any(ResourceStatsType.class), any(ResourceUsageMetric[].class)); + } + + @Test + void cancellationPollerInterruptsWhenTaskCancelled() { + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) + .thenReturn(true); + when(settings.getSettingValue(Settings.Key.PPL_QUERY_TIMEOUT)) + .thenReturn(new TimeValue(60000)); + CancellableTask cancelledTask = mock(CancellableTask.class); + when(cancelledTask.isCancelled()).thenReturn(true); + OpenSearchQueryManager.setCancellableTask(cancelledTask); + + // Capture the poller runnable so we can fire it directly and assert it interrupts the thread. + AtomicReference poller = new AtomicReference<>(); + when(threadPool.scheduleWithFixedDelay(any(Runnable.class), any(TimeValue.class), any())) + .thenAnswer( + invocation -> { + poller.set(invocation.getArgument(0)); + return mock(Cancellable.class); + }); + + AtomicReference interruptedInExecute = new AtomicReference<>(false); + doAnswer( + invocation -> { + Object executor = invocation.getArgument(2); + if (SQL_COMPLEX_WORKER_THREAD_POOL_NAME.equals(executor)) { + invocation.getArgument(0).run(); + } + return mock(ScheduledCancellable.class); + }) + .when(threadPool) + .schedule(any(Runnable.class), any(TimeValue.class), any()); + doAnswer( + invocation -> { + // Fire the poller on this (execution) thread; it should interrupt us. + poller.get().run(); + interruptedInExecute.set(Thread.currentThread().isInterrupted()); + Thread.interrupted(); // clear so we don't leak the interrupt into the pool thread + return null; + }) + .when(engine) + .execute(any(RelNode.class), any(), any(ResponseListener.class)); + + AbstractCalciteIndexScan scan = createMockScanWithScripts(); + dispatcher.dispatch(scan, context, listener, engine); + + assertTrue( + interruptedInExecute.get(), + "cancellation poller should interrupt the execution thread when the task is cancelled"); + } + + @Test + void dispatchTaskRunsInlineWhenNoScripts() { + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) + .thenReturn(true); + RelNode plan = createMockNode(); + AtomicReference ran = new AtomicReference<>(false); + + dispatcher.dispatchTask(plan, context, () -> ran.set(true)); + + assertTrue(ran.get()); + verify(threadPool, never()).schedule(any(), any(TimeValue.class), any()); + } + + @Test + void dispatchTaskSwallowsFailureWhenNoListener() { + when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) + .thenReturn(true); + when(settings.getSettingValue(Settings.Key.PPL_QUERY_TIMEOUT)) + .thenReturn(new TimeValue(60000)); + doAnswer( + invocation -> { + Object executor = invocation.getArgument(2); + if (SQL_COMPLEX_WORKER_THREAD_POOL_NAME.equals(executor)) { + invocation.getArgument(0).run(); + } + return mock(ScheduledCancellable.class); + }) + .when(threadPool) + .schedule(any(Runnable.class), any(TimeValue.class), any()); + + AbstractCalciteIndexScan scan = createMockScanWithScripts(); + + dispatcher.dispatchTask( + scan, + context, + () -> { + throw new RuntimeException("boom"); + }); + + // No exception escaped and the task was cleared: the null-listener catch + finally ran. + assertNull(OpenSearchQueryManager.getCancellableTask()); + } + private static RelNode createMockNode(RelNode... children) { RelNode node = mock(RelNode.class); List childList = List.of(children); diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java index 93ea787c617..67684208a45 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java @@ -120,6 +120,7 @@ import org.opensearch.sql.plugin.rest.RestQuerySettingsAction; import org.opensearch.sql.plugin.rest.RestUnifiedQueryAction; import org.opensearch.sql.plugin.transport.PPLQueryAction; +import org.opensearch.sql.plugin.transport.QueryInsightsMarker; import org.opensearch.sql.plugin.transport.TransportPPLQueryAction; import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse; import org.opensearch.sql.prometheus.storage.PrometheusStorageFactory; @@ -357,6 +358,16 @@ public void onFailure(Exception e) { }; } + /** + * Register the Query Insights parent-marker header so OpenSearch copies it from a SQL/PPL query's + * thread context into the child DSL search tasks it spawns (including on remote data nodes). This + * lets Query Insights classify each child's source and associate it with the originating query. + */ + @Override + public Collection getTaskHeaders() { + return List.of(QueryInsightsMarker.PARENT_HEADER); + } + /** Register action and handler so that transportClient can find proxy for action. */ @Override public List> getActions() { diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java index 2df96bdbd12..fa4385cd114 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java @@ -5,12 +5,32 @@ package org.opensearch.sql.plugin.transport; +import java.util.List; import java.util.Map; import org.opensearch.core.tasks.TaskId; import org.opensearch.tasks.CancellableTask; public class PPLQueryTask extends CancellableTask { + // The fields below are captured on the request thread and read later by the report listener + // (which runs on a different thread), for the Query Insights record. + + /** Security {@code _opendistro_security_user_info} string; null when unsecured. */ + private volatile String queryInsightsUserInfo; + + /** Source index name(s) resolved from the AST; empty when unresolved. */ + private volatile List queryInsightsIndices = List.of(); + + /** Anonymized query text ({@code PPLQueryDataAnonymizer}); null when anonymization didn't run. */ + private volatile String queryInsightsAnonymizedQuery; + + /** + * Whether per-thread CPU/memory accounting is enabled for this task. Off by default so a normal + * PPL query behaves like any {@link CancellableTask}; turned on only when Query Insights + * recording and core's task resource tracking are both enabled. + */ + private volatile boolean resourceTrackingEnabled = false; + public PPLQueryTask( long id, String type, @@ -21,8 +41,46 @@ public PPLQueryTask( super(id, type, action, description, parentTaskId, headers); } + public void setQueryInsightsUserInfo(String userInfo) { + this.queryInsightsUserInfo = userInfo; + } + + public String getQueryInsightsUserInfo() { + return queryInsightsUserInfo; + } + + public void setQueryInsightsIndices(List indices) { + this.queryInsightsIndices = indices == null ? List.of() : indices; + } + + public List getQueryInsightsIndices() { + return queryInsightsIndices; + } + + public void setQueryInsightsAnonymizedQuery(String anonymizedQuery) { + this.queryInsightsAnonymizedQuery = anonymizedQuery; + } + + public String getQueryInsightsAnonymizedQuery() { + return queryInsightsAnonymizedQuery; + } + + public void setResourceTrackingEnabled(boolean enabled) { + this.resourceTrackingEnabled = enabled; + } + @Override public boolean shouldCancelChildrenOnCancellation() { return true; } + + /** + * Per-thread CPU/memory accounting for the coordinator task, enabled only when Query Insights + * recording and core's {@code task_resource_tracking.enabled} are both on (see {@code + * TransportPPLQueryAction}). {@link CancellableTask} defaults to {@code false}. + */ + @Override + public boolean supportsResourceTracking() { + return resourceTrackingEnabled; + } } diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsMarker.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsMarker.java new file mode 100644 index 00000000000..2725873692a --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsMarker.java @@ -0,0 +1,24 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +/** + * Task header linking a child DSL search back to the SQL/PPL query that spawned it, for Query + * Insights. Value is {@code ::}; registered via {@code + * SQLPlugin.getTaskHeaders()} so OpenSearch copies it onto child search tasks. + */ +public final class QueryInsightsMarker { + + /** Task header name carrying {@code ::} of the originating query. */ + public static final String PARENT_HEADER = "X-Query-Insights-Parent"; + + private QueryInsightsMarker() {} + + /** Build the {@code ::} header value for a coordinator query. */ + public static String value(String source, String nodeId, long taskId) { + return source + ":" + nodeId + ":" + taskId; + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java new file mode 100644 index 00000000000..51dd4bfe3bc --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java @@ -0,0 +1,89 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import java.util.List; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.Version; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.transport.BytesTransportRequest; +import org.opensearch.transport.EmptyTransportResponseHandler; +import org.opensearch.transport.TransportService; + +/** + * Serializes a completed PPL/SQL query record and sends it to Query Insights over the transport + * layer. Uses the core {@link BytesTransportRequest} (loaded by both plugin classloaders, unlike a + * plugin-defined type) sent via {@link TransportService#sendRequest}. + */ +public final class QueryInsightsReporter { + + private static final Logger LOG = LogManager.getLogger(QueryInsightsReporter.class); + + public static final String ACTION_NAME = + "cluster:admin/opensearch/query_insights/report_query_bytes"; + + public static final int FORMAT_VERSION = 2; + + private QueryInsightsReporter() {} + + /** + * Serialize the record and send it to the local node's Query Insights handler; errors are ignored + * (e.g. Query Insights not installed). + * + * @param parentMarker {@code ::}; child DSL records reference it via + * {@code DERIVED_FROM} so their cpu/memory roll up into this record + * @param userInfo security {@code _opendistro_security_user_info} string; empty when unsecured + */ + public static void report( + TransportService transportService, + DiscoveryNode localNode, + String querySource, + String parentMarker, + String nodeId, + String queryText, + long timestampMillis, + long latencyMillis, + long cpuNanos, + long memoryBytes, + List indices, + String userInfo) { + try { + final BytesStreamOutput out = new BytesStreamOutput(); + out.writeVInt(FORMAT_VERSION); + out.writeString(querySource == null ? "" : querySource); + out.writeString(parentMarker == null ? "" : parentMarker); + out.writeString(nodeId == null ? "" : nodeId); + out.writeString(queryText == null ? "" : queryText); + out.writeVLong(timestampMillis); + out.writeVLong(Math.max(0L, latencyMillis)); + out.writeVLong(Math.max(0L, cpuNanos)); + out.writeVLong(Math.max(0L, memoryBytes)); + + final List safeIndices = indices == null ? List.of() : indices; + out.writeVInt(safeIndices.size()); + for (final String index : safeIndices) { + out.writeString(index == null ? "" : index); + } + + out.writeString(userInfo == null ? "" : userInfo); // v2 + + final BytesTransportRequest request = new BytesTransportRequest(out.bytes(), Version.CURRENT); + // Stash the context so this internal cluster:admin send runs user-less (security's + // Origin.LOCAL path) and isn't authorized against the end user. The handler reads the user + // off the wire, so it needs no user context. + final ThreadContext threadContext = transportService.getThreadPool().getThreadContext(); + try (ThreadContext.StoredContext ignored = threadContext.stashContext()) { + transportService.sendRequest( + localNode, ACTION_NAME, request, EmptyTransportResponseHandler.INSTANCE_SAME); + } + } catch (Exception e) { + LOG.debug("Failed to report PPL query to Query Insights", e); + } + } +} 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 6754995ea06..0db16ef9507 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 @@ -11,6 +11,7 @@ import static org.opensearch.sql.protocol.response.format.JsonResponseFormatter.Style.PRETTY; import java.util.LinkedHashMap; +import java.util.List; import java.util.Locale; import java.util.Map; import java.util.Optional; @@ -28,7 +29,11 @@ import org.opensearch.common.inject.Inject; import org.opensearch.common.inject.Injector; import org.opensearch.common.inject.ModulesBuilder; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; +import org.opensearch.core.action.NotifyOnceListener; +import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.common.utils.QueryContext; @@ -50,6 +55,7 @@ import org.opensearch.sql.plugin.rest.AnalyticsExecutorHolder; import org.opensearch.sql.plugin.rest.RestUnifiedQueryAction; import org.opensearch.sql.ppl.PPLService; +import org.opensearch.sql.ppl.QueryInsightsMetadata; import org.opensearch.sql.ppl.domain.PPLQueryRequest; import org.opensearch.sql.protocol.response.QueryResult; import org.opensearch.sql.protocol.response.format.CsvResponseFormatter; @@ -61,6 +67,7 @@ import org.opensearch.sql.protocol.response.format.VisualizationResponseFormatter; import org.opensearch.sql.protocol.response.format.YamlResponseFormatter; import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskResourceTrackingService; import org.opensearch.telemetry.tracing.Span; import org.opensearch.telemetry.tracing.SpanCreationContext; import org.opensearch.telemetry.tracing.SpanScope; @@ -88,6 +95,7 @@ public class TransportPPLQueryAction private final NodeClient clientRef; private final ClusterService clusterServiceRef; private final org.opensearch.sql.common.setting.Settings pluginSettingsRef; + private final TransportService transportServiceRef; @Inject public TransportPPLQueryAction( @@ -102,6 +110,7 @@ public TransportPPLQueryAction( super(PPLQueryAction.NAME, transportService, actionFilters, TransportPPLQueryRequest::new); this.clientRef = client; this.clusterServiceRef = clusterService; + this.transportServiceRef = transportService; ModulesBuilder modules = new ModulesBuilder(); modules.add(new OpenSearchPluginModule(extensionsHolder.engines(), tracer)); @@ -160,6 +169,126 @@ private void buildUnifiedQueryHandlerIfReady() { } } + /** Stamp the parent marker so child DSL searches carry it back to this PPL query. */ + private void stampQueryInsightsParentHeader(PPLQueryTask pplQueryTask) { + try { + ThreadContext threadContext = clientRef.threadPool().getThreadContext(); + String value = + QueryInsightsMarker.value( + "PPL", clusterServiceRef.localNode().getId(), pplQueryTask.getId()); + if (threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER) == null) { + threadContext.putHeader(QueryInsightsMarker.PARENT_HEADER, value); + } + // Capture the user here; the report listener runs on a thread without the transient. + pplQueryTask.setQueryInsightsUserInfo( + threadContext.getTransient(SECURITY_USER_INFO_TRANSIENT)); + } catch (Exception e) { + LOG.warn("Failed to stamp Query Insights parent header for query association", e); + } + } + + /** Security plugin's authenticated-user transient ({@code name|backendroles|roles|...}). */ + private static final String SECURITY_USER_INFO_TRANSIENT = "_opendistro_security_user_info"; + + /** + * Toggle for recording PPL queries into Query Insights. Registered by the Query Insights plugin + * ({@code QueryInsightsSettings.TOP_N_PPL_QUERIES_ENABLED}, default false). + */ + static final String QUERY_INSIGHTS_PPL_ENABLED_KEY = "search.insights.top_queries.ppl.enabled"; + + /** True when Query Insights is installed and the toggle is on; disabled on any failure. */ + private boolean isQueryInsightsRecordingEnabled() { + try { + Setting setting = + clusterServiceRef.getClusterSettings().get(QUERY_INSIGHTS_PPL_ENABLED_KEY); + if (setting == null) { + return false; + } + Object value = clusterServiceRef.getClusterSettings().get(setting); + return Boolean.TRUE.equals(value); + } catch (Exception e) { + LOG.debug("Failed to evaluate Query Insights PPL recording gate; defaulting to disabled", e); + return false; + } + } + + /** + * Mirrors core's {@code task_resource_tracking.enabled} gate so PPL resource tracking is not + * enabled when an operator has disabled task resource tracking cluster-wide. Defaults to the core + * setting's own default ({@code true}) if the value can't be read. + */ + private boolean isTaskResourceTrackingEnabled() { + try { + return clusterServiceRef + .getClusterSettings() + .get(TaskResourceTrackingService.TASK_RESOURCE_TRACKING_ENABLED); + } catch (Exception e) { + LOG.debug("Failed to read task_resource_tracking.enabled; defaulting to enabled", e); + return true; + } + } + + /** Report the finished PPL query once its resource tracking completes (stats are final). */ + private void registerQueryInsightsReport(PPLQueryTask pplQueryTask) { + try { + boolean registered = + pplQueryTask.addResourceTrackingCompletionListener( + new NotifyOnceListener<>() { + @Override + protected void innerOnResponse(Task task) { + writeQueryInsightsRecord((PPLQueryTask) task); + } + + @Override + protected void innerOnFailure(Exception e) {} + }); + if (registered == false) { + LOG.debug( + "Query Insights report listener not registered; resource tracking already complete"); + } + } catch (Exception e) { + LOG.warn("Failed to register Query Insights report listener", e); + } + } + + /** Serialize a completed PPL query and send it to Query Insights; errors are ignored. */ + private void writeQueryInsightsRecord(PPLQueryTask pplQueryTask) { + try { + String nodeId = clusterServiceRef.localNode().getId(); + String queryText = pplQueryTask.getQueryInsightsAnonymizedQuery(); + + TaskResourceUsage usage = pplQueryTask.getTotalResourceStats(); + long cpuNanos = usage == null ? 0L : usage.getCpuTimeInNanos(); + long memoryBytes = usage == null ? 0L : usage.getMemoryInBytes(); + + // Wall-clock latency (monotonic); the profile total collapses to ~0 across the finish thread. + long latencyMillis = + Math.max(0L, (System.nanoTime() - pplQueryTask.getStartTimeNanos()) / 1_000_000L); + + // Same marker stamped on child DSL tasks, so their cpu/memory rolls up into this parent. + String parentMarker = QueryInsightsMarker.value("PPL", nodeId, pplQueryTask.getId()); + + List indices = pplQueryTask.getQueryInsightsIndices(); + String userInfo = pplQueryTask.getQueryInsightsUserInfo(); + + QueryInsightsReporter.report( + transportServiceRef, + clusterServiceRef.localNode(), + "PPL", + parentMarker, + nodeId, + queryText, + System.currentTimeMillis(), + latencyMillis, + cpuNanos, + memoryBytes, + indices, + userInfo); + } catch (Exception e) { + LOG.debug("Failed to write PPL query to Query Insights", e); + } + } + /** * {@inheritDoc} Transform the request and call super.doExecute() to support call from other * plugins. @@ -183,8 +312,17 @@ protected void doExecute( return; } - if (task instanceof PPLQueryTask pplQueryTask) { + final PPLQueryTask pplQueryTask = task instanceof PPLQueryTask ? (PPLQueryTask) task : null; + if (pplQueryTask != null) { OpenSearchQueryManager.setCancellableTask(pplQueryTask); + if (isQueryInsightsRecordingEnabled()) { + // Only track resource usage when Query Insights recording and core's task resource tracking + // are both on, so a disabled feature adds no ThreadMXBean overhead and we don't populate + // _tasks resource_stats behind the operator's back. + pplQueryTask.setResourceTrackingEnabled(isTaskResourceTrackingEnabled()); + stampQueryInsightsParentHeader(pplQueryTask); + registerQueryInsightsReport(pplQueryTask); + } } Metrics.getInstance().getNumericalMetric(MetricName.PPL_REQ_TOTAL).increment(); Metrics.getInstance().getNumericalMetric(MetricName.PPL_REQ_COUNT_TOTAL).increment(); @@ -250,30 +388,39 @@ protected void doExecute( return; } - Consumer anonymizedQuerySink = - anonymized -> rootSpan.addAttribute("db.query.text", anonymized); + final PPLQueryTask metadataTarget = pplQueryTask; + Consumer queryInsightsSink = + metadata -> { + rootSpan.addAttribute("db.query.text", metadata.anonymizedQuery()); + if (metadataTarget != null) { + metadataTarget.setQueryInsightsAnonymizedQuery(metadata.anonymizedQuery()); + metadataTarget.setQueryInsightsIndices(metadata.indices()); + } + }; PPLService pplService = injector.getInstance(PPLService.class); if (transformedRequest.isExplainRequest()) { pplService.explain( transformedRequest, createExplainResponseListener(transformedRequest, clearingListener), - anonymizedQuerySink); + queryInsightsSink); } else if (transformedRequest.analyze()) { pplService.analyze( transformedRequest, createAnalyzeResponseListener(transformedRequest, clearingListener), - anonymizedQuerySink); + queryInsightsSink); } else { pplService.execute( transformedRequest, createListener(transformedRequest, clearingListener), createExplainResponseListener(transformedRequest, clearingListener), - anonymizedQuerySink); + queryInsightsSink); } } catch (Exception e) { clearingListener.onFailure(e); } finally { spanScope.close(); + // Clear the task in case an early return/throw skipped submit() and left it on this thread. + OpenSearchQueryManager.clearCancellableTask(); } } diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java index c9502ac3bbf..b5bf47a0d5a 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java @@ -62,4 +62,40 @@ public void testCooperativeModel() { task.cancel("Test"); assertTrue(task.isCancelled()); } + + private PPLQueryTask newTask() { + return new PPLQueryTask( + 1, + "transport", + "cluster:admin/opensearch/ppl", + "test query", + TaskId.EMPTY_TASK_ID, + Map.of()); + } + + @Test + public void testSupportsResourceTrackingOffByDefault() { + // Off by default so a PPL query adds no resource-tracking overhead unless Query Insights + // recording (and core's task resource tracking) are enabled. + assertFalse(newTask().supportsResourceTracking()); + } + + @Test + public void testSupportsResourceTrackingWhenEnabled() { + PPLQueryTask task = newTask(); + task.setResourceTrackingEnabled(true); + assertTrue(task.supportsResourceTracking()); + } + + @Test + public void testQueryInsightsParentHeaderName() { + assertEquals("X-Query-Insights-Parent", QueryInsightsMarker.PARENT_HEADER); + } + + @Test + public void testQueryInsightsParentHeaderValueIsSourcePrefixed() { + // Value format is :: so QI reads both source and parent id from one + // header. SQL will reuse the same helper with source "SQL". + assertEquals("PPL:node-1:42", QueryInsightsMarker.value("PPL", "node-1", 42L)); + } } 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 8d82b550d44..eee3bb8efdf 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java @@ -8,9 +8,11 @@ import static org.opensearch.sql.executor.ExecutionEngine.QueryResponse; import static org.opensearch.sql.executor.execution.QueryPlanFactory.NO_CONSUMER_RESPONSE_LISTENER; +import java.util.List; import java.util.function.Consumer; import lombok.extern.log4j.Log4j2; import org.antlr.v4.runtime.tree.ParseTree; +import org.opensearch.sql.ast.statement.Explain; import org.opensearch.sql.ast.statement.Query; import org.opensearch.sql.ast.statement.Statement; import org.opensearch.sql.ast.tree.UnresolvedPlan; @@ -29,13 +31,14 @@ import org.opensearch.sql.ppl.parser.AstBuilder; import org.opensearch.sql.ppl.parser.AstStatementBuilder; import org.opensearch.sql.ppl.utils.PPLQueryDataAnonymizer; +import org.opensearch.sql.ppl.utils.PPLQueryIndexExtractor; /** PPLService. */ @Log4j2 public class PPLService { - /** Callers that don't care about the anonymized query pass this. */ - public static final Consumer NO_ANONYMIZED_QUERY_SINK = s -> {}; + /** Callers that don't need the query-insights metadata pass this. */ + public static final Consumer NO_QUERY_INSIGHTS_SINK = m -> {}; private final PPLSyntaxParser parser; @@ -72,17 +75,17 @@ public void execute( PPLQueryRequest request, ResponseListener queryListener, ResponseListener explainListener) { - execute(request, queryListener, explainListener, NO_ANONYMIZED_QUERY_SINK); + execute(request, queryListener, explainListener, NO_QUERY_INSIGHTS_SINK); } - /** Variant that hands the anonymized query text to {@code anonymizedQuerySink}. */ + /** Variant that hands query-insights metadata (anonymized query + indices) to {@code sink}. */ public void execute( PPLQueryRequest request, ResponseListener queryListener, ResponseListener explainListener, - Consumer anonymizedQuerySink) { + Consumer sink) { try { - queryManager.submit(plan(request, queryListener, explainListener, anonymizedQuerySink)); + queryManager.submit(plan(request, queryListener, explainListener, sink)); } catch (Exception e) { queryListener.onFailure(e); } @@ -96,17 +99,16 @@ public void execute( * @param listener {@link ResponseListener} for explain response */ public void explain(PPLQueryRequest request, ResponseListener listener) { - explain(request, listener, NO_ANONYMIZED_QUERY_SINK); + explain(request, listener, NO_QUERY_INSIGHTS_SINK); } - /** Variant that hands the anonymized query text to {@code anonymizedQuerySink}. */ + /** Variant that hands query-insights metadata (anonymized query + indices) to {@code sink}. */ public void explain( PPLQueryRequest request, ResponseListener listener, - Consumer anonymizedQuerySink) { + Consumer sink) { try { - queryManager.submit( - plan(request, NO_CONSUMER_RESPONSE_LISTENER, listener, anonymizedQuerySink)); + queryManager.submit(plan(request, NO_CONSUMER_RESPONSE_LISTENER, listener, sink)); } catch (Exception e) { listener.onFailure(e); } @@ -119,24 +121,23 @@ public void explain( * @param listener {@link ResponseListener} for analyze response */ public void analyze(PPLQueryRequest request, ResponseListener listener) { - analyze(request, listener, NO_ANONYMIZED_QUERY_SINK); + analyze(request, listener, NO_QUERY_INSIGHTS_SINK); } - /** Variant that hands the anonymized query text to {@code anonymizedQuerySink}. */ + /** Variant that hands query-insights metadata (anonymized query + indices) to {@code sink}. */ public void analyze( PPLQueryRequest request, ResponseListener listener, - Consumer anonymizedQuerySink) { + Consumer sink) { try { String queryText = request.getRequest(); - ParseTree cst; Statement statement; String anonymized; // Transport-thread work — parse, AST build, anonymize. Trace-only: QueryProfiling isn't // active yet on this thread. Cold-start ANTLR grammar init dominates this region. try (ProfileScope preparePhase = ProfileScope.openTraceOnly("prepare")) { try { - cst = parser.parse(queryText); + ParseTree cst = parser.parse(queryText); statement = cst.accept( new AstStatementBuilder( @@ -159,9 +160,9 @@ public void analyze( } } log.info("[{}] Incoming request {}", QueryContext.getRequestId(), anonymized); - anonymizedQuerySink.accept(anonymized); UnresolvedPlan unresolvedPlan = ((Query) statement).getPlan(); + sink.accept(new QueryInsightsMetadata(anonymized, extractIndexNames(statement))); AbstractPlan analyzePlan = queryExecutionFactory.createAnalyzePlan(unresolvedPlan, PPL_QUERY, listener); queryManager.submit(analyzePlan); @@ -174,7 +175,7 @@ private AbstractPlan plan( PPLQueryRequest request, ResponseListener queryListener, ResponseListener explainListener, - Consumer anonymizedQuerySink) { + Consumer sink) { Statement statement; String anonymized; // Transport-thread work — parse, AST build, anonymize. Trace-only: QueryProfiling isn't @@ -208,9 +209,30 @@ private AbstractPlan plan( } } log.info("[{}] Incoming request {}", QueryContext.getRequestId(), anonymized); - anonymizedQuerySink.accept(anonymized); + + // Extract indices from the same AST that will execute, so recorded metadata matches the + // executed query and the query is parsed only once. + sink.accept(new QueryInsightsMetadata(anonymized, extractIndexNames(statement))); AbstractPlan plan = queryExecutionFactory.create(statement, queryListener, explainListener); return plan; } + + /** + * Best-effort source index names from an already-parsed statement (unwrapping an {@code explain} + * to its inner query); empty on any failure. + */ + private static List extractIndexNames(Statement statement) { + try { + Statement inner = + statement instanceof Explain ? ((Explain) statement).getStatement() : statement; + if (inner instanceof Query query) { + return PPLQueryIndexExtractor.extractIndexNames(query.getPlan()); + } + return List.of(); + } catch (Exception e) { + log.debug("[{}] Failed to extract PPL index names", QueryContext.getRequestId(), e); + return List.of(); + } + } } diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/QueryInsightsMetadata.java b/ppl/src/main/java/org/opensearch/sql/ppl/QueryInsightsMetadata.java new file mode 100644 index 00000000000..4325c6de06b --- /dev/null +++ b/ppl/src/main/java/org/opensearch/sql/ppl/QueryInsightsMetadata.java @@ -0,0 +1,17 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ppl; + +import java.util.List; + +/** + * Best-effort metadata about a PPL query, derived once from the parsed AST that is executed, for + * consumers such as Query Insights. + * + * @param anonymizedQuery the query with literals masked by {@code PPLQueryDataAnonymizer} + * @param indices the source index name(s) the query reads from; empty when none are resolvable + */ +public record QueryInsightsMetadata(String anonymizedQuery, List indices) {} diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java b/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java new file mode 100644 index 00000000000..839abf4b48e --- /dev/null +++ b/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java @@ -0,0 +1,158 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ppl.utils; + +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; +import org.opensearch.sql.ast.AbstractNodeVisitor; +import org.opensearch.sql.ast.expression.Let; +import org.opensearch.sql.ast.expression.QualifiedName; +import org.opensearch.sql.ast.expression.subquery.ExistsSubquery; +import org.opensearch.sql.ast.expression.subquery.InSubquery; +import org.opensearch.sql.ast.expression.subquery.ScalarSubquery; +import org.opensearch.sql.ast.tree.Append; +import org.opensearch.sql.ast.tree.AppendCol; +import org.opensearch.sql.ast.tree.AppendPipe; +import org.opensearch.sql.ast.tree.Eval; +import org.opensearch.sql.ast.tree.Filter; +import org.opensearch.sql.ast.tree.GraphLookup; +import org.opensearch.sql.ast.tree.Join; +import org.opensearch.sql.ast.tree.Lookup; +import org.opensearch.sql.ast.tree.Relation; +import org.opensearch.sql.ast.tree.UnresolvedPlan; + +/** + * Collects the source index name(s) a PPL query reads from by walking the parsed AST (via {@link + * Relation#getQualifiedNames()}) instead of regex-matching text, which mishandles multi-index and + * datasource-qualified sources. + * + *

The default visitor only follows {@code getChild()}, so a source held in any other field is + * reached by an explicit override: a {@link Join}'s left/right branches, a {@link Lookup} table, + * the sub-search of {@link Append}/{@link AppendCol}/{@link AppendPipe}, a {@link GraphLookup}'s + * {@code from} table, and subquery relations embedded in {@link Filter}/{@link Eval} condition + * expressions. Returns distinct names in encounter order. This is best-effort metadata; a source it + * cannot resolve is simply omitted. + */ +public final class PPLQueryIndexExtractor { + + private PPLQueryIndexExtractor() {} + + /** Return the distinct source index name(s) referenced by {@code plan}, in encounter order. */ + public static List extractIndexNames(UnresolvedPlan plan) { + Set names = new LinkedHashSet<>(); + if (plan != null) { + plan.accept(new Collector(), names); + } + return new ArrayList<>(names); + } + + private static final class Collector extends AbstractNodeVisitor> { + @Override + public Void visitRelation(Relation node, Set names) { + for (QualifiedName qualifiedName : node.getQualifiedNames()) { + names.add(qualifiedName.toString()); + } + return null; + } + + @Override + public Void visitJoin(Join node, Set names) { + // getChild() exposes only the left branch (and only once attached, so it may be null); visit + // both sides explicitly, guarding the nullable left. + if (node.getLeft() != null) { + node.getLeft().accept(this, names); + } + node.getRight().accept(this, names); + return null; + } + + @Override + public Void visitLookup(Lookup node, Set names) { + // getChild() returns only the piped input; the lookup table is a separate relation. + super.visitChildren(node, names); + node.getLookupRelation().accept(this, names); + return null; + } + + @Override + public Void visitAppend(Append node, Set names) { + super.visitChildren(node, names); + acceptIfPresent(node.getSubSearch(), names); + return null; + } + + @Override + public Void visitAppendCol(AppendCol node, Set names) { + super.visitChildren(node, names); + acceptIfPresent(node.getSubSearch(), names); + return null; + } + + @Override + public Void visitAppendPipe(AppendPipe node, Set names) { + super.visitChildren(node, names); + acceptIfPresent(node.getSubQuery(), names); + return null; + } + + @Override + public Void visitGraphLookup(GraphLookup node, Set names) { + super.visitChildren(node, names); + acceptIfPresent(node.getFromTable(), names); + return null; + } + + @Override + public Void visitFilter(Filter node, Set names) { + // The condition can hold a subquery (where id in [ source=b ... ]); the plan walk doesn't + // descend into it, so bridge into the expression here. + super.visitChildren(node, names); + if (node.getCondition() != null) { + node.getCondition().accept(this, names); + } + return null; + } + + @Override + public Void visitEval(Eval node, Set names) { + super.visitChildren(node, names); + if (node.getExpressionList() != null) { + for (Let let : node.getExpressionList()) { + if (let != null && let.getExpression() != null) { + let.getExpression().accept(this, names); + } + } + } + return null; + } + + @Override + public Void visitScalarSubquery(ScalarSubquery node, Set names) { + node.getQuery().accept(this, names); + return null; + } + + @Override + public Void visitInSubquery(InSubquery node, Set names) { + node.getQuery().accept(this, names); + return null; + } + + @Override + public Void visitExistsSubquery(ExistsSubquery node, Set names) { + node.getQuery().accept(this, names); + return null; + } + + private void acceptIfPresent(UnresolvedPlan plan, Set names) { + if (plan != null) { + plan.accept(this, names); + } + } + } +} diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java index fea5e3d4030..68f88c6acd4 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java @@ -11,6 +11,7 @@ import java.util.Collections; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import org.junit.After; import org.junit.Assert; import org.junit.Before; @@ -115,6 +116,46 @@ public void testExecuteShouldPass() { getExplainListener(false)); } + @Test + public void testExecutePassesAnonymizedQueryToSink() { + doAnswer( + invocation -> { + ResponseListener listener = invocation.getArgument(4); + listener.onResponse(new QueryResponse(schema, Collections.emptyList(), Cursor.None)); + return null; + }) + .when(queryService) + .execute(any(), any(), any(), anyBoolean(), any()); + + AtomicReference metadata = new AtomicReference<>(); + pplService.execute( + new PPLQueryRequest("search source=t a=42", null, QUERY), + getQueryListener(false), + getExplainListener(false), + metadata::set); + + // The sink receives the anonymized query (literal masked, never the raw value) and the source + // index resolved from the same AST that executes. + Assert.assertNotNull(metadata.get()); + Assert.assertTrue(metadata.get().anonymizedQuery().contains("***")); + Assert.assertFalse(metadata.get().anonymizedQuery().contains("42")); + Assert.assertEquals(Collections.singletonList("t"), metadata.get().indices()); + } + + @Test + public void testExplainPassesAnonymizedQueryToSink() { + AtomicReference metadata = new AtomicReference<>(); + pplService.explain( + new PPLQueryRequest("search source=t a=42", null, EXPLAIN), + getExplainListener(false), + metadata::set); + + Assert.assertNotNull(metadata.get()); + Assert.assertTrue(metadata.get().anonymizedQuery().contains("***")); + Assert.assertFalse(metadata.get().anonymizedQuery().contains("42")); + Assert.assertEquals(Collections.singletonList("t"), metadata.get().indices()); + } + @Test public void testExecuteCsvFormatShouldPass() { doAnswer( diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java new file mode 100644 index 00000000000..b65f3f6b781 --- /dev/null +++ b/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java @@ -0,0 +1,87 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ppl.utils; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import java.util.List; +import org.junit.Test; +import org.opensearch.sql.ast.tree.UnresolvedPlan; +import org.opensearch.sql.ppl.AstPlanningTestBase; + +/** Unit tests for {@link PPLQueryIndexExtractor}. */ +public class PPLQueryIndexExtractorTest extends AstPlanningTestBase { + + private List indices(String query) { + return PPLQueryIndexExtractor.extractIndexNames((UnresolvedPlan) plan(query)); + } + + @Test + public void singleSource() { + assertEquals(List.of("accounts"), indices("source=accounts | fields firstname")); + } + + @Test + public void indexKeyword() { + assertEquals(List.of("logs"), indices("search index=logs")); + } + + @Test + public void multipleCommaSeparatedSourcesWithSpace() { + // The reviewer's case: "source=accounts, account2" must keep BOTH indices — the old regex + // stopped at the whitespace after the comma and dropped account2. + assertEquals(List.of("accounts", "account2"), indices("source=accounts, account2")); + } + + @Test + public void multipleCommaSeparatedSourcesNoSpace() { + assertEquals(List.of("a", "b", "c"), indices("source=a,b,c | stats count()")); + } + + @Test + public void joinCapturesBothSides() { + List result = indices("source=t1 | join on t1.id = t2.id t2"); + assertTrue(result.contains("t1")); + assertTrue(result.contains("t2")); + } + + @Test + public void lookupCapturesLookupTable() { + List result = indices("source=t1 | lookup t2 id"); + assertTrue(result.contains("t1")); + assertTrue(result.contains("t2")); + } + + @Test + public void datasourceQualifiedName() { + assertEquals(List.of("myds.myindex"), indices("source=myds.myindex | head 1")); + } + + @Test + public void filterAndFieldsDoNotAddIndices() { + assertEquals(List.of("accounts"), indices("source=accounts | where age > 30 | fields name")); + } + + @Test + public void appendCapturesSubSearchSource() { + List result = indices("source=accounts | append [ source=account2 | stats count() ]"); + assertTrue(result.contains("accounts")); + assertTrue(result.contains("account2")); + } + + @Test + public void subqueryInWhereCapturesInnerSource() { + List result = indices("source=outer | where a in [ source=inner | fields b ]"); + assertTrue(result.contains("outer")); + assertTrue(result.contains("inner")); + } + + @Test + public void nullPlanYieldsEmpty() { + assertTrue(PPLQueryIndexExtractor.extractIndexNames(null).isEmpty()); + } +} From 3bf5d5aca2f77cd534506341244d47339284e3af Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Thu, 1 Oct 2026 00:52:43 -0700 Subject: [PATCH 02/26] feat(ppl): Report failed PPL queries to Query Insights Carry a failed flag from the PPL coordinator into the Query Insights record so a failed PPL/SQL query is classified as failed in Top N, matching how native DSL searches are recorded. TransportPPLQueryAction wraps its listener so any onFailure marks the task failed; the flag is read when the Query Insights record is emitted and serialized as the trailing wire field by QueryInsightsReporter. FORMAT_VERSION is 1 (the format is unreleased and built as a unit). Also add PPLQueryIndexExtractor coverage for the appendcol, appendpipe, graphlookup, and eval-subquery source bridges, and document that a subquery in a join ON condition is intentionally not walked. Signed-off-by: Kishore Kumaar Natarajan --- .../sql/plugin/transport/PPLQueryTask.java | 11 ++++ .../transport/QueryInsightsReporter.java | 9 ++-- .../transport/TransportPPLQueryAction.java | 26 ++++++++- .../plugin/transport/PPLQueryTaskTest.java | 12 +++++ .../sql/ppl/utils/PPLQueryIndexExtractor.java | 3 ++ .../ppl/utils/PPLQueryIndexExtractorTest.java | 53 +++++++++++++++++++ 6 files changed, 110 insertions(+), 4 deletions(-) diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java index fa4385cd114..b049b1613b8 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java @@ -24,6 +24,9 @@ public class PPLQueryTask extends CancellableTask { /** Anonymized query text ({@code PPLQueryDataAnonymizer}); null when anonymization didn't run. */ private volatile String queryInsightsAnonymizedQuery; + /** Whether the query failed; recorded so a failed PPL query is flagged in Top N. */ + private volatile boolean queryInsightsFailed = false; + /** * Whether per-thread CPU/memory accounting is enabled for this task. Off by default so a normal * PPL query behaves like any {@link CancellableTask}; turned on only when Query Insights @@ -65,6 +68,14 @@ public String getQueryInsightsAnonymizedQuery() { return queryInsightsAnonymizedQuery; } + public void setQueryInsightsFailed(boolean failed) { + this.queryInsightsFailed = failed; + } + + public boolean isQueryInsightsFailed() { + return queryInsightsFailed; + } + public void setResourceTrackingEnabled(boolean enabled) { this.resourceTrackingEnabled = enabled; } diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java index 51dd4bfe3bc..7adaf64109f 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java @@ -28,7 +28,7 @@ public final class QueryInsightsReporter { public static final String ACTION_NAME = "cluster:admin/opensearch/query_insights/report_query_bytes"; - public static final int FORMAT_VERSION = 2; + public static final int FORMAT_VERSION = 1; private QueryInsightsReporter() {} @@ -39,6 +39,7 @@ private QueryInsightsReporter() {} * @param parentMarker {@code ::}; child DSL records reference it via * {@code DERIVED_FROM} so their cpu/memory roll up into this record * @param userInfo security {@code _opendistro_security_user_info} string; empty when unsecured + * @param failed whether the query failed, so Query Insights can flag it in Top N */ public static void report( TransportService transportService, @@ -52,7 +53,8 @@ public static void report( long cpuNanos, long memoryBytes, List indices, - String userInfo) { + String userInfo, + boolean failed) { try { final BytesStreamOutput out = new BytesStreamOutput(); out.writeVInt(FORMAT_VERSION); @@ -71,7 +73,8 @@ public static void report( out.writeString(index == null ? "" : index); } - out.writeString(userInfo == null ? "" : userInfo); // v2 + out.writeString(userInfo == null ? "" : userInfo); + out.writeBoolean(failed); final BytesTransportRequest request = new BytesTransportRequest(out.bytes(), Version.CURRENT); // Stash the context so this internal cluster:admin send runs user-less (security's 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 0db16ef9507..41df0786054 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 @@ -270,6 +270,7 @@ private void writeQueryInsightsRecord(PPLQueryTask pplQueryTask) { List indices = pplQueryTask.getQueryInsightsIndices(); String userInfo = pplQueryTask.getQueryInsightsUserInfo(); + boolean failed = pplQueryTask.isQueryInsightsFailed(); QueryInsightsReporter.report( transportServiceRef, @@ -283,7 +284,8 @@ private void writeQueryInsightsRecord(PPLQueryTask pplQueryTask) { cpuNanos, memoryBytes, indices, - userInfo); + userInfo, + failed); } catch (Exception e) { LOG.debug("Failed to write PPL query to Query Insights", e); } @@ -355,6 +357,8 @@ protected void doExecute( TraceableActionListener.create(listener, rootSpan, tracer); ActionListener clearingListener = wrapWithProfilingClear(tracedListener); + // Read back later by writeQueryInsightsRecord, which runs after this listener completes. + clearingListener = markFailedOnFailure(clearingListener, pplQueryTask); try { // Route to analytics engine for non-Lucene (e.g., Parquet-backed) indices. @@ -546,6 +550,26 @@ private Format format(PPLQueryRequest pplRequest) { } } + /** Returns {@code delegate} unchanged when there is no PPL task to mark. */ + private ActionListener markFailedOnFailure( + ActionListener delegate, PPLQueryTask pplQueryTask) { + if (pplQueryTask == null) { + return delegate; + } + return new ActionListener<>() { + @Override + public void onResponse(TransportPPLQueryResponse response) { + delegate.onResponse(response); + } + + @Override + public void onFailure(Exception e) { + pplQueryTask.setQueryInsightsFailed(true); + delegate.onFailure(e); + } + }; + } + private ActionListener wrapWithProfilingClear( ActionListener delegate) { return new ActionListener<>() { diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java index b5bf47a0d5a..676c304d258 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java @@ -87,6 +87,18 @@ public void testSupportsResourceTrackingWhenEnabled() { assertTrue(task.supportsResourceTracking()); } + @Test + public void testQueryInsightsFailedOffByDefault() { + assertFalse(newTask().isQueryInsightsFailed()); + } + + @Test + public void testQueryInsightsFailedFlagIsSettable() { + PPLQueryTask task = newTask(); + task.setQueryInsightsFailed(true); + assertTrue(task.isQueryInsightsFailed()); + } + @Test public void testQueryInsightsParentHeaderName() { assertEquals("X-Query-Insights-Parent", QueryInsightsMarker.PARENT_HEADER); diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java b/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java index 839abf4b48e..a876e094988 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java @@ -37,6 +37,9 @@ * {@code from} table, and subquery relations embedded in {@link Filter}/{@link Eval} condition * expressions. Returns distinct names in encounter order. This is best-effort metadata; a source it * cannot resolve is simply omitted. + * + *

A subquery in a {@link Join}'s ON condition is not walked (only left/right branches are), + * unlike {@link Filter}/{@link Eval}; such a source is dropped from this best-effort list. */ public final class PPLQueryIndexExtractor { diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java index b65f3f6b781..d36faee2ce5 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java @@ -10,6 +10,13 @@ import java.util.List; import org.junit.Test; +import org.opensearch.sql.ast.expression.DataType; +import org.opensearch.sql.ast.expression.Literal; +import org.opensearch.sql.ast.expression.QualifiedName; +import org.opensearch.sql.ast.tree.AppendCol; +import org.opensearch.sql.ast.tree.AppendPipe; +import org.opensearch.sql.ast.tree.Filter; +import org.opensearch.sql.ast.tree.Relation; import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.ppl.AstPlanningTestBase; @@ -80,6 +87,52 @@ public void subqueryInWhereCapturesInnerSource() { assertTrue(result.contains("inner")); } + @Test + public void evalScalarSubqueryCapturesInnerSource() { + List result = indices("source=outer | eval m = [ source=inner | stats max(b) ]"); + assertTrue(result.contains("outer")); + assertTrue(result.contains("inner")); + } + + @Test + public void graphLookupCapturesFromTable() { + List result = + indices( + "source=t | graphLookup employees start=reportsTo edge=manager-->name" + + " as reportingHierarchy"); + assertTrue(result.contains("t")); + assertTrue(result.contains("employees")); + } + + @Test + public void appendColCapturesSubSearchSource() { + // Built directly, not parsed: the grammar runs appendcol's sub-search over the piped input, so + // a nested source= can't be expressed in text. + UnresolvedPlan plan = new AppendCol(false, subPipelineOver("appended")).attach(baseRelation()); + List result = PPLQueryIndexExtractor.extractIndexNames(plan); + assertTrue(result.contains("base")); + assertTrue(result.contains("appended")); + } + + @Test + public void appendPipeCapturesSubQuerySource() { + // Built directly for the same reason as appendcol above. + UnresolvedPlan plan = new AppendPipe(subPipelineOver("appended")).attach(baseRelation()); + List result = PPLQueryIndexExtractor.extractIndexNames(plan); + assertTrue(result.contains("base")); + assertTrue(result.contains("appended")); + } + + private UnresolvedPlan baseRelation() { + return new Relation(new QualifiedName("base")); + } + + /** A {@link Relation} nested under a piped command, like real parser output (not a bare root). */ + private UnresolvedPlan subPipelineOver(String index) { + return new Filter(new Literal(true, DataType.BOOLEAN)) + .attach(new Relation(new QualifiedName(index))); + } + @Test public void nullPlanYieldsEmpty() { assertTrue(PPLQueryIndexExtractor.extractIndexNames(null).isEmpty()); From cb33a90707d3bff9c76501a1096175687bef679f Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 6 Oct 2026 02:37:02 -0700 Subject: [PATCH 03/26] fix(ppl): Scope Query Insights setup to the execute route Only an executed PPL query produces a Top N record, so move the parent header stamp and the report listener into the execute branch. Explain, analyze, and the analytics engine run no search: they were emitting records with an empty query string, no indices and zero cpu/memory, and stamping their child searches orphaned those children under a parent that never reports. Stamp the header inside a stored context that is restored when doExecute returns, and hand the caller a context-preserving listener. An in-process caller that issues several PPL queries (or a plain DSL search) through NodeClient no longer inherits the first query's marker, either on the calling thread or in the response callback. Child searches still see the header because the executor captures the thread context when the work is scheduled, inside the stored-context block. Signed-off-by: Kishore Kumaar Natarajan --- .../transport/TransportPPLQueryAction.java | 43 +++++++++++++++---- 1 file changed, 35 insertions(+), 8 deletions(-) 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 41df0786054..a42be7508b8 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 @@ -22,6 +22,7 @@ import org.apache.logging.log4j.Logger; import org.opensearch.action.ActionRequest; import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.ContextPreservingActionListener; import org.opensearch.action.support.HandledTransportAction; import org.opensearch.analytics.exec.QueryPlanExecutor; import org.opensearch.cluster.service.ClusterService; @@ -320,10 +321,9 @@ protected void doExecute( if (isQueryInsightsRecordingEnabled()) { // Only track resource usage when Query Insights recording and core's task resource tracking // are both on, so a disabled feature adds no ThreadMXBean overhead and we don't populate - // _tasks resource_stats behind the operator's back. + // _tasks resource_stats behind the operator's back. The header stamp and report listener + // are set up only on the execute route below. pplQueryTask.setResourceTrackingEnabled(isTaskResourceTrackingEnabled()); - stampQueryInsightsParentHeader(pplQueryTask); - registerQueryInsightsReport(pplQueryTask); } } Metrics.getInstance().getNumericalMetric(MetricName.PPL_REQ_TOTAL).increment(); @@ -412,12 +412,27 @@ protected void doExecute( transformedRequest, createAnalyzeResponseListener(transformedRequest, clearingListener), queryInsightsSink); + } else if (pplQueryTask != null && isQueryInsightsRecordingEnabled()) { + // Query Insights is set up only on this execute route: explain, analyze, and the analytics + // engine run no search, so stamping their child searches would orphan them under a parent + // that never reports, and reporting the query itself would put an empty record in Top N. + final ThreadContext threadContext = clientRef.threadPool().getThreadContext(); + // Capture the caller's context before stamping, so the callback runs without our header. + // The worker thread that completes the query carries it, and an in-process caller invoked + // via NodeClient would otherwise inherit it in onResponse. + ActionListener callerListener = + ContextPreservingActionListener.wrapPreservingContext(clearingListener, threadContext); + // Stamp inside a stored context so the header is dropped from this thread when doExecute + // returns, keeping a caller that issues several PPL queries from tagging later ones (or a + // plain DSL search) with this query's marker. Child searches still see it: the executor + // captures the thread context when the work is scheduled, which happens inside this block. + try (ThreadContext.StoredContext ignored = threadContext.newStoredContext(true)) { + stampQueryInsightsParentHeader(pplQueryTask); + registerQueryInsightsReport(pplQueryTask); + executePplQuery(pplService, transformedRequest, callerListener, queryInsightsSink); + } } else { - pplService.execute( - transformedRequest, - createListener(transformedRequest, clearingListener), - createExplainResponseListener(transformedRequest, clearingListener), - queryInsightsSink); + executePplQuery(pplService, transformedRequest, clearingListener, queryInsightsSink); } } catch (Exception e) { clearingListener.onFailure(e); @@ -550,6 +565,18 @@ private Format format(PPLQueryRequest pplRequest) { } } + private void executePplQuery( + PPLService pplService, + PPLQueryRequest request, + ActionListener listener, + Consumer queryInsightsSink) { + pplService.execute( + request, + createListener(request, listener), + createExplainResponseListener(request, listener), + queryInsightsSink); + } + /** Returns {@code delegate} unchanged when there is no PPL task to mark. */ private ActionListener markFailedOnFailure( ActionListener delegate, PPLQueryTask pplQueryTask) { From b72cdb4385a811e7847227fe951143a4000938c1 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 6 Oct 2026 02:43:02 -0700 Subject: [PATCH 04/26] test(ppl): Pin the parent-header thread-context contract Assert the three properties the scoped stamp depends on: work scheduled during execution still sees the parent header, the header does not outlive the call on the calling thread, and the caller's callback does not inherit it. Also cover a caller that already carries a marker. Signed-off-by: Kishore Kumaar Natarajan --- .../QueryInsightsParentHeaderScopeTest.java | 88 +++++++++++++++++++ 1 file changed, 88 insertions(+) create mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsParentHeaderScopeTest.java diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsParentHeaderScopeTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsParentHeaderScopeTest.java new file mode 100644 index 00000000000..6d514c3d2ca --- /dev/null +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsParentHeaderScopeTest.java @@ -0,0 +1,88 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; + +import java.util.concurrent.atomic.AtomicReference; +import org.junit.Test; +import org.opensearch.action.support.ContextPreservingActionListener; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.action.ActionListener; + +/** + * Pins the thread-context contract {@code TransportPPLQueryAction} relies on when it stamps the + * Query Insights parent header: the header must reach the searches scheduled during execution, but + * must not outlive the call on the calling thread or reappear in the caller's callback. Without + * that, a caller issuing several PPL queries through {@code NodeClient} would tag later queries (or + * a plain DSL search) with the first query's marker. + */ +public class QueryInsightsParentHeaderScopeTest { + + private static final String MARKER = "PPL:node-1:42"; + + @Test + public void headerReachesScheduledWorkButDoesNotLeak() { + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + AtomicReference seenByScheduledWork = new AtomicReference<>("unset"); + AtomicReference seenByCallback = new AtomicReference<>("unset"); + + ActionListener caller = + ActionListener.wrap( + response -> + seenByCallback.set(threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER)), + e -> {}); + // Captured before stamping, mirroring doExecute. + ActionListener callerListener = + ContextPreservingActionListener.wrapPreservingContext(caller, threadContext); + + Runnable scheduledWork; + try (ThreadContext.StoredContext ignored = threadContext.newStoredContext(true)) { + threadContext.putHeader(QueryInsightsMarker.PARENT_HEADER, MARKER); + // Executors capture the thread context when work is scheduled; this is what carries the + // header to the sql-worker and background-scan threads. + scheduledWork = + threadContext.preserveContext( + () -> + seenByScheduledWork.set( + threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER))); + } + + assertNull( + "header must not outlive the call on the calling thread", + threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER)); + + scheduledWork.run(); + assertEquals( + "searches scheduled during execution must still see the header", + MARKER, + seenByScheduledWork.get()); + + callerListener.onResponse("ok"); + assertNull("the caller's callback must not inherit the header", seenByCallback.get()); + } + + @Test + public void existingCallerHeaderIsPreservedAndRestored() { + ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + threadContext.putHeader(QueryInsightsMarker.PARENT_HEADER, "caller-supplied"); + + try (ThreadContext.StoredContext ignored = threadContext.newStoredContext(true)) { + // doExecute only stamps when absent, so a caller that already carries a marker keeps it. + if (threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER) == null) { + threadContext.putHeader(QueryInsightsMarker.PARENT_HEADER, MARKER); + } + assertEquals("caller-supplied", threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER)); + } + + assertEquals( + "the caller's own header must survive the scoped stamp", + "caller-supplied", + threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER)); + } +} From 0e9b98905ffbe4e82fa247325f1cdd13d60ec2ec Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 6 Oct 2026 03:23:48 -0700 Subject: [PATCH 05/26] fix(ppl): Skip Top N records for explains and unparsed queries Decide whether to report from what the metadata sink captured instead of from the route, because the route cannot see either case. An inline explain ("explain source=t | ...") reaches the execute route because isExplainRequest() only inspects the path, so it was reported as an executed query. QueryInsightsMetadata now carries whether the parsed statement is an Explain, and the task records it. A query that fails to parse throws before the sink runs, so the report fired with a null query text and no indices, putting a blank entry in Top N for every malformed request. Reporting now requires the sink to have captured metadata. A query that plans and then fails at runtime still has its metadata, so it keeps being reported as failed. Signed-off-by: Kishore Kumaar Natarajan --- .../sql/plugin/transport/PPLQueryTask.java | 11 ++++++++ .../transport/TransportPPLQueryAction.java | 15 +++++++++++ .../plugin/transport/PPLQueryTaskTest.java | 25 +++++++++++++++++++ .../org/opensearch/sql/ppl/PPLService.java | 8 ++++-- .../sql/ppl/QueryInsightsMetadata.java | 5 +++- 5 files changed, 61 insertions(+), 3 deletions(-) diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java index b049b1613b8..a463376d59e 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java @@ -27,6 +27,9 @@ public class PPLQueryTask extends CancellableTask { /** Whether the query failed; recorded so a failed PPL query is flagged in Top N. */ private volatile boolean queryInsightsFailed = false; + /** Whether the statement is an {@code explain}, which plans but runs no search. */ + private volatile boolean queryInsightsExplain = false; + /** * Whether per-thread CPU/memory accounting is enabled for this task. Off by default so a normal * PPL query behaves like any {@link CancellableTask}; turned on only when Query Insights @@ -68,6 +71,14 @@ public String getQueryInsightsAnonymizedQuery() { return queryInsightsAnonymizedQuery; } + public void setQueryInsightsExplain(boolean explain) { + this.queryInsightsExplain = explain; + } + + public boolean isQueryInsightsExplain() { + return queryInsightsExplain; + } + public void setQueryInsightsFailed(boolean failed) { this.queryInsightsFailed = failed; } 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 a42be7508b8..e2ea33bc6fb 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 @@ -253,7 +253,21 @@ protected void innerOnFailure(Exception e) {} } /** Serialize a completed PPL query and send it to Query Insights; errors are ignored. */ + /** + * Whether a finished task should produce a Top N record. Decided from what the metadata sink + * captured rather than from the route, because the route cannot see either case: {@code explain} + * is valid inside the query text (not just on {@code /_explain}), and a query that fails to parse + * never reaches the sink at all, so it has no query text or indices to report. + */ + static boolean shouldReportToQueryInsights(PPLQueryTask pplQueryTask) { + return pplQueryTask.getQueryInsightsAnonymizedQuery() != null + && !pplQueryTask.isQueryInsightsExplain(); + } + private void writeQueryInsightsRecord(PPLQueryTask pplQueryTask) { + if (!shouldReportToQueryInsights(pplQueryTask)) { + return; + } try { String nodeId = clusterServiceRef.localNode().getId(); String queryText = pplQueryTask.getQueryInsightsAnonymizedQuery(); @@ -399,6 +413,7 @@ protected void doExecute( if (metadataTarget != null) { metadataTarget.setQueryInsightsAnonymizedQuery(metadata.anonymizedQuery()); metadataTarget.setQueryInsightsIndices(metadata.indices()); + metadataTarget.setQueryInsightsExplain(metadata.explain()); } }; PPLService pplService = injector.getInstance(PPLService.class); diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java index 676c304d258..33178445deb 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java @@ -99,6 +99,31 @@ public void testQueryInsightsFailedFlagIsSettable() { assertTrue(task.isQueryInsightsFailed()); } + @Test + public void testInlineExplainIsNotReported() { + // "explain source=t | ..." reaches the execute route because isExplainRequest() only looks at + // the path, so the sink's flag is what keeps it out of Top N. + PPLQueryTask task = newTask(); + task.setQueryInsightsAnonymizedQuery("explain source=table | fields identifier"); + task.setQueryInsightsExplain(true); + assertFalse(TransportPPLQueryAction.shouldReportToQueryInsights(task)); + } + + @Test + public void testSyntaxErrorIsNotReported() { + // A parse failure throws before the sink runs, leaving no query text to report. + assertFalse(TransportPPLQueryAction.shouldReportToQueryInsights(newTask())); + } + + @Test + public void testRuntimeFailureAfterPlanningIsReported() { + // Planning succeeded, so the metadata is present; a later execution failure must still report. + PPLQueryTask task = newTask(); + task.setQueryInsightsAnonymizedQuery("source=table | where identifier > ***"); + task.setQueryInsightsFailed(true); + assertTrue(TransportPPLQueryAction.shouldReportToQueryInsights(task)); + } + @Test public void testQueryInsightsParentHeaderName() { assertEquals("X-Query-Insights-Parent", QueryInsightsMarker.PARENT_HEADER); 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 eee3bb8efdf..f8a2483a94f 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java @@ -162,7 +162,9 @@ public void analyze( log.info("[{}] Incoming request {}", QueryContext.getRequestId(), anonymized); UnresolvedPlan unresolvedPlan = ((Query) statement).getPlan(); - sink.accept(new QueryInsightsMetadata(anonymized, extractIndexNames(statement))); + sink.accept( + new QueryInsightsMetadata( + anonymized, extractIndexNames(statement), statement instanceof Explain)); AbstractPlan analyzePlan = queryExecutionFactory.createAnalyzePlan(unresolvedPlan, PPL_QUERY, listener); queryManager.submit(analyzePlan); @@ -212,7 +214,9 @@ private AbstractPlan plan( // Extract indices from the same AST that will execute, so recorded metadata matches the // executed query and the query is parsed only once. - sink.accept(new QueryInsightsMetadata(anonymized, extractIndexNames(statement))); + sink.accept( + new QueryInsightsMetadata( + anonymized, extractIndexNames(statement), statement instanceof Explain)); AbstractPlan plan = queryExecutionFactory.create(statement, queryListener, explainListener); return plan; diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/QueryInsightsMetadata.java b/ppl/src/main/java/org/opensearch/sql/ppl/QueryInsightsMetadata.java index 4325c6de06b..c638990bf2c 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/QueryInsightsMetadata.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/QueryInsightsMetadata.java @@ -13,5 +13,8 @@ * * @param anonymizedQuery the query with literals masked by {@code PPLQueryDataAnonymizer} * @param indices the source index name(s) the query reads from; empty when none are resolvable + * @param explain whether the statement is an {@code explain}, which plans but runs no search. The + * endpoint alone cannot tell: {@code explain} is also valid inside the query text. */ -public record QueryInsightsMetadata(String anonymizedQuery, List indices) {} +public record QueryInsightsMetadata( + String anonymizedQuery, List indices, boolean explain) {} From 2a552ab7b3da720e5962d1fed8ee881fce88bea6 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 6 Oct 2026 05:00:00 -0700 Subject: [PATCH 06/26] refactor(ast): Let plan nodes declare their source plans Add UnresolvedPlan.getSources() so a node declares any plan it holds in a field rather than in getChild(), and override it on the six commands that do: Append, AppendCol, AppendPipe, GraphLookup, Lookup and Join. PPLQueryIndexExtractor now walks getChild() plus getSources() instead of carrying a visitor override per command, so a new command only has to declare its sources on the node to be picked up. This removes the class of silent miss where a source held in a field was never reached. Subqueries carried in expressions are not reachable through a plan-level contract, so the Filter and Eval condition bridges remain. Signed-off-by: Kishore Kumaar Natarajan --- .../org/opensearch/sql/ast/tree/Append.java | 6 ++ .../opensearch/sql/ast/tree/AppendCol.java | 6 ++ .../opensearch/sql/ast/tree/AppendPipe.java | 6 ++ .../opensearch/sql/ast/tree/GraphLookup.java | 6 ++ .../org/opensearch/sql/ast/tree/Join.java | 9 ++ .../org/opensearch/sql/ast/tree/Lookup.java | 6 ++ .../sql/ast/tree/UnresolvedPlan.java | 17 ++++ .../sql/ppl/utils/PPLQueryIndexExtractor.java | 91 ++++++------------- 8 files changed, 83 insertions(+), 64 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Append.java b/core/src/main/java/org/opensearch/sql/ast/tree/Append.java index e5d4f54228b..edb8b907455 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/Append.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Append.java @@ -33,6 +33,12 @@ public UnresolvedPlan attach(UnresolvedPlan child) { return this; } + /** {@inheritDoc} Adds the appended sub-search, which is not a child. */ + @Override + public List getSources() { + return this.subSearch == null ? ImmutableList.of() : ImmutableList.of(this.subSearch); + } + @Override public List getChild() { return this.child == null ? ImmutableList.of() : ImmutableList.of(this.child); diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/AppendCol.java b/core/src/main/java/org/opensearch/sql/ast/tree/AppendCol.java index 39bfe255749..ef011ad2627 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/AppendCol.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/AppendCol.java @@ -37,6 +37,12 @@ public UnresolvedPlan attach(UnresolvedPlan child) { return this; } + /** {@inheritDoc} Adds the appended sub-search, which is not a child. */ + @Override + public List getSources() { + return this.subSearch == null ? ImmutableList.of() : ImmutableList.of(this.subSearch); + } + @Override public List getChild() { return this.child == null ? ImmutableList.of() : ImmutableList.of(this.child); diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/AppendPipe.java b/core/src/main/java/org/opensearch/sql/ast/tree/AppendPipe.java index 0ea1cb9b453..6ad7fb4eefd 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/AppendPipe.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/AppendPipe.java @@ -33,6 +33,12 @@ public AppendPipe attach(UnresolvedPlan child) { return this; } + /** {@inheritDoc} Adds the appended sub-pipeline, which is not a child. */ + @Override + public List getSources() { + return this.subQuery == null ? ImmutableList.of() : ImmutableList.of(this.subQuery); + } + @Override public List getChild() { return this.child == null ? ImmutableList.of() : ImmutableList.of(this.child); diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/GraphLookup.java b/core/src/main/java/org/opensearch/sql/ast/tree/GraphLookup.java index f1457a8a540..d38ec09fc98 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/GraphLookup.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/GraphLookup.java @@ -95,6 +95,12 @@ public UnresolvedPlan attach(UnresolvedPlan child) { return this; } + /** {@inheritDoc} Adds the graph-lookup source table, which is not a child. */ + @Override + public List getSources() { + return this.fromTable == null ? ImmutableList.of() : ImmutableList.of(this.fromTable); + } + @Override public List getChild() { return child == null ? ImmutableList.of() : ImmutableList.of(child); diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Join.java b/core/src/main/java/org/opensearch/sql/ast/tree/Join.java index a0f58d35124..3c2e9497521 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/Join.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Join.java @@ -73,6 +73,15 @@ public UnresolvedPlan attach(UnresolvedPlan child) { return this; } + /** + * {@inheritDoc} Adds the right branch ({@code getChild()} exposes only the left), which is not a + * child. + */ + @Override + public List getSources() { + return this.right == null ? ImmutableList.of() : ImmutableList.of(this.right); + } + @Override public List getChild() { return this.left == null ? ImmutableList.of() : ImmutableList.of(this.left); diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Lookup.java b/core/src/main/java/org/opensearch/sql/ast/tree/Lookup.java index 5b0c5d585d8..dc2144232dd 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/Lookup.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Lookup.java @@ -45,6 +45,12 @@ public UnresolvedPlan attach(UnresolvedPlan child) { return this; } + /** {@inheritDoc} Adds the lookup table, which is not a child. */ + @Override + public List getSources() { + return this.lookupRelation == null ? ImmutableList.of() : ImmutableList.of(this.lookupRelation); + } + @Override public List getChild() { return this.child == null ? ImmutableList.of() : ImmutableList.of(this.child); diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/UnresolvedPlan.java b/core/src/main/java/org/opensearch/sql/ast/tree/UnresolvedPlan.java index 3074303b0d7..d4f616ff3ed 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/UnresolvedPlan.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/UnresolvedPlan.java @@ -5,6 +5,8 @@ package org.opensearch.sql.ast.tree; +import com.google.common.collect.ImmutableList; +import java.util.List; import lombok.EqualsAndHashCode; import lombok.ToString; import org.opensearch.sql.ast.AbstractNodeVisitor; @@ -19,5 +21,20 @@ public T accept(AbstractNodeVisitor nodeVisitor, C context) { return nodeVisitor.visitChildren(this, context); } + /** + * Plans this node reads from that {@link #getChild()} does not expose, such as a sub-search or a + * lookup table held in a separate field. + * + *

Override this when adding a command that holds a plan outside its child list. Consumers that + * need every source a query reads from (for example index-name extraction for Query Insights) + * walk {@code getChild()} plus this list, so a plan left out of both is silently invisible to + * them. The returned plans must not duplicate {@code getChild()}. + * + * @return the additional source plans; empty by default + */ + public List getSources() { + return ImmutableList.of(); + } + public abstract UnresolvedPlan attach(UnresolvedPlan child); } diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java b/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java index a876e094988..44c1af0fd09 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java @@ -10,19 +10,15 @@ import java.util.List; import java.util.Set; import org.opensearch.sql.ast.AbstractNodeVisitor; +import org.opensearch.sql.ast.Node; import org.opensearch.sql.ast.expression.Let; import org.opensearch.sql.ast.expression.QualifiedName; import org.opensearch.sql.ast.expression.subquery.ExistsSubquery; import org.opensearch.sql.ast.expression.subquery.InSubquery; import org.opensearch.sql.ast.expression.subquery.ScalarSubquery; -import org.opensearch.sql.ast.tree.Append; -import org.opensearch.sql.ast.tree.AppendCol; -import org.opensearch.sql.ast.tree.AppendPipe; import org.opensearch.sql.ast.tree.Eval; import org.opensearch.sql.ast.tree.Filter; -import org.opensearch.sql.ast.tree.GraphLookup; import org.opensearch.sql.ast.tree.Join; -import org.opensearch.sql.ast.tree.Lookup; import org.opensearch.sql.ast.tree.Relation; import org.opensearch.sql.ast.tree.UnresolvedPlan; @@ -31,15 +27,17 @@ * Relation#getQualifiedNames()}) instead of regex-matching text, which mishandles multi-index and * datasource-qualified sources. * - *

The default visitor only follows {@code getChild()}, so a source held in any other field is - * reached by an explicit override: a {@link Join}'s left/right branches, a {@link Lookup} table, - * the sub-search of {@link Append}/{@link AppendCol}/{@link AppendPipe}, a {@link GraphLookup}'s - * {@code from} table, and subquery relations embedded in {@link Filter}/{@link Eval} condition - * expressions. Returns distinct names in encounter order. This is best-effort metadata; a source it - * cannot resolve is simply omitted. + *

Plan sources are found generically: the walk follows {@code getChild()} plus {@link + * UnresolvedPlan#getSources()}, which each node uses to declare a plan it holds in a field (a + * sub-search, a lookup table, a join's right branch). A new command therefore only has to declare + * its sources on the node itself to be covered here. * - *

A subquery in a {@link Join}'s ON condition is not walked (only left/right branches are), - * unlike {@link Filter}/{@link Eval}; such a source is dropped from this best-effort list. + *

Subqueries carried in expressions are not reachable that way, so {@link Filter} and + * {@link Eval} conditions are bridged explicitly. A subquery in a {@link Join}'s ON condition is + * still not walked and is dropped. + * + *

Returns distinct names in encounter order. This is best-effort metadata; a source it cannot + * resolve is simply omitted and query execution is unaffected. */ public final class PPLQueryIndexExtractor { @@ -63,58 +61,29 @@ public Void visitRelation(Relation node, Set names) { return null; } + /** + * Walks {@code getChild()} as usual, then the node's declared {@link + * UnresolvedPlan#getSources()}. Commands whose source lives in a field are covered without this + * collector knowing each one. + */ @Override - public Void visitJoin(Join node, Set names) { - // getChild() exposes only the left branch (and only once attached, so it may be null); visit - // both sides explicitly, guarding the nullable left. - if (node.getLeft() != null) { - node.getLeft().accept(this, names); - } - node.getRight().accept(this, names); - return null; - } - - @Override - public Void visitLookup(Lookup node, Set names) { - // getChild() returns only the piped input; the lookup table is a separate relation. - super.visitChildren(node, names); - node.getLookupRelation().accept(this, names); - return null; - } - - @Override - public Void visitAppend(Append node, Set names) { - super.visitChildren(node, names); - acceptIfPresent(node.getSubSearch(), names); - return null; - } - - @Override - public Void visitAppendCol(AppendCol node, Set names) { - super.visitChildren(node, names); - acceptIfPresent(node.getSubSearch(), names); - return null; - } - - @Override - public Void visitAppendPipe(AppendPipe node, Set names) { - super.visitChildren(node, names); - acceptIfPresent(node.getSubQuery(), names); - return null; - } - - @Override - public Void visitGraphLookup(GraphLookup node, Set names) { + public Void visitChildren(Node node, Set names) { super.visitChildren(node, names); - acceptIfPresent(node.getFromTable(), names); + if (node instanceof UnresolvedPlan plan) { + for (UnresolvedPlan source : plan.getSources()) { + if (source != null) { + source.accept(this, names); + } + } + } return null; } @Override public Void visitFilter(Filter node, Set names) { // The condition can hold a subquery (where id in [ source=b ... ]); the plan walk doesn't - // descend into it, so bridge into the expression here. - super.visitChildren(node, names); + // descend into expressions, so bridge into it here. + visitChildren(node, names); if (node.getCondition() != null) { node.getCondition().accept(this, names); } @@ -123,7 +92,7 @@ public Void visitFilter(Filter node, Set names) { @Override public Void visitEval(Eval node, Set names) { - super.visitChildren(node, names); + visitChildren(node, names); if (node.getExpressionList() != null) { for (Let let : node.getExpressionList()) { if (let != null && let.getExpression() != null) { @@ -151,11 +120,5 @@ public Void visitExistsSubquery(ExistsSubquery node, Set names) { node.getQuery().accept(this, names); return null; } - - private void acceptIfPresent(UnresolvedPlan plan, Set names) { - if (plan != null) { - plan.accept(this, names); - } - } } } From ffcda088b47162a9d03a12eade59ac71f74485c2 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 6 Oct 2026 05:01:58 -0700 Subject: [PATCH 07/26] test(ast): Pin the getSources() contract Cover the parts the PPL extractor tests exercise only indirectly: the empty default, each declaring node returning its field, that sources do not repeat getChild(), and that an absent sub-search yields an empty list instead of an NPE. Signed-off-by: Kishore Kumaar Natarajan --- .../ast/tree/UnresolvedPlanSourcesTest.java | 68 +++++++++++++++++++ 1 file changed, 68 insertions(+) create mode 100644 core/src/test/java/org/opensearch/sql/ast/tree/UnresolvedPlanSourcesTest.java diff --git a/core/src/test/java/org/opensearch/sql/ast/tree/UnresolvedPlanSourcesTest.java b/core/src/test/java/org/opensearch/sql/ast/tree/UnresolvedPlanSourcesTest.java new file mode 100644 index 00000000000..fddb6cd298e --- /dev/null +++ b/core/src/test/java/org/opensearch/sql/ast/tree/UnresolvedPlanSourcesTest.java @@ -0,0 +1,68 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ast.tree; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.List; +import org.junit.jupiter.api.Test; +import org.opensearch.sql.ast.expression.QualifiedName; + +/** + * Pins the {@link UnresolvedPlan#getSources()} contract: a node declares the plans it reads from + * that {@code getChild()} does not expose, so consumers can find every source generically. A node + * that drops its declaration here silently hides a source from them. + */ +class UnresolvedPlanSourcesTest { + + private static UnresolvedPlan relation(String name) { + return new Relation(new QualifiedName(name)); + } + + @Test + void defaultsToEmpty() { + // A node whose sources are all reachable via getChild() declares none. + assertTrue(relation("t").getSources().isEmpty()); + } + + @Test + void appendDeclaresSubSearch() { + UnresolvedPlan sub = relation("appended"); + assertEquals(List.of(sub), new Append(sub).attach(relation("base")).getSources()); + } + + @Test + void appendColDeclaresSubSearch() { + UnresolvedPlan sub = relation("appended"); + assertEquals(List.of(sub), new AppendCol(false, sub).attach(relation("base")).getSources()); + } + + @Test + void appendPipeDeclaresSubQuery() { + UnresolvedPlan sub = relation("appended"); + assertEquals(List.of(sub), new AppendPipe(sub).attach(relation("base")).getSources()); + } + + @Test + void sourcesAreDisjointFromChild() { + // The contract requires getSources() not to repeat getChild(), or consumers walking both would + // visit the piped input twice. + UnresolvedPlan child = relation("base"); + UnresolvedPlan sub = relation("appended"); + Append append = (Append) new Append(sub).attach(child); + assertEquals(List.of(child), append.getChild()); + assertEquals(List.of(sub), append.getSources()); + } + + @Test + void nullSourceYieldsEmptyRatherThanNpe() { + // Nodes can be inspected before attach()/with an absent sub-search; the walk must not break. + assertTrue(new AppendPipe(null).getSources().isEmpty()); + assertTrue(new Append(null).getSources().isEmpty()); + assertTrue(new AppendCol(false, null).getSources().isEmpty()); + } +} From 59f2ebb6655b218eb085814522decdc8bfa127ef Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 6 Oct 2026 05:06:25 -0700 Subject: [PATCH 08/26] fix(ppl): Enable resource tracking only where a record consumes it Move setResourceTrackingEnabled into the execute branch, next to the header stamp and the report listener. It was still being set before routing, so explain, analyze and the analytics engine paid the ThreadMXBean sampling and per-thread bookkeeping without ever producing a record to use it. Signed-off-by: Kishore Kumaar Natarajan --- .../plugin/transport/TransportPPLQueryAction.java | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) 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 e2ea33bc6fb..7fd468559c8 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 @@ -332,13 +332,6 @@ protected void doExecute( final PPLQueryTask pplQueryTask = task instanceof PPLQueryTask ? (PPLQueryTask) task : null; if (pplQueryTask != null) { OpenSearchQueryManager.setCancellableTask(pplQueryTask); - if (isQueryInsightsRecordingEnabled()) { - // Only track resource usage when Query Insights recording and core's task resource tracking - // are both on, so a disabled feature adds no ThreadMXBean overhead and we don't populate - // _tasks resource_stats behind the operator's back. The header stamp and report listener - // are set up only on the execute route below. - pplQueryTask.setResourceTrackingEnabled(isTaskResourceTrackingEnabled()); - } } Metrics.getInstance().getNumericalMetric(MetricName.PPL_REQ_TOTAL).increment(); Metrics.getInstance().getNumericalMetric(MetricName.PPL_REQ_COUNT_TOTAL).increment(); @@ -442,6 +435,12 @@ protected void doExecute( // plain DSL search) with this query's marker. Child searches still see it: the executor // captures the thread context when the work is scheduled, which happens inside this block. try (ThreadContext.StoredContext ignored = threadContext.newStoredContext(true)) { + // Enable per-thread accounting only here, where a record will actually consume it, and + // only when core's task_resource_tracking is on so we never populate _tasks + // resource_stats behind the operator's back. The engine's brackets gate on + // supportsResourceTracking(), which stays false for every other route and when the + // feature is off, so they cost nothing then. + pplQueryTask.setResourceTrackingEnabled(isTaskResourceTrackingEnabled()); stampQueryInsightsParentHeader(pplQueryTask); registerQueryInsightsReport(pplQueryTask); executePplQuery(pplService, transformedRequest, callerListener, queryInsightsSink); From a8755caa97e438f9101d9ae0e2feba3657a17f3a Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 6 Oct 2026 05:13:26 -0700 Subject: [PATCH 09/26] fix(ppl): Mirror core's full resource-tracking gate Core refuses to track unless the task supports it, the cluster setting task_resource_tracking.enabled is on, AND the JVM reports per-thread allocation (TaskResourceTrackingService.isTaskResourceTrackingSupported). We honoured the first two but only null-checked the ThreadMXBean, so on a JVM where allocation reporting is unsupported or disabled a PPL task would have recorded stats that core itself would never record. Gate startThreadResourceTracking on the same support check. Signed-off-by: Kishore Kumaar Natarajan --- .../executor/OpenSearchQueryManager.java | 15 +++++++++++++++ .../executor/OpenSearchQueryManagerTest.java | 15 +++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java index def9d195d5b..0b7a3ca1cd4 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java @@ -149,7 +149,22 @@ private Runnable withCurrentContext(final Runnable task) { * @return true if tracking started; only then should {@link #stopThreadResourceTracking} be * called. */ + /** + * Whether the JVM can report per-thread allocation, mirroring core's {@code + * TaskResourceTrackingService.isTaskResourceTrackingSupported()}. Core refuses to track when this + * is false, so we must too: otherwise a PPL task would record stats on a JVM where core + * deliberately records none, and {@code getThreadAllocatedBytes} is not meaningful there. + */ + static boolean isThreadResourceTrackingSupported() { + return THREAD_MX_BEAN != null + && THREAD_MX_BEAN.isThreadAllocatedMemorySupported() + && THREAD_MX_BEAN.isThreadAllocatedMemoryEnabled(); + } + static boolean startThreadResourceTracking(CancellableTask task, long threadId) { + if (!isThreadResourceTrackingSupported()) { + return false; + } try { task.startThreadResourceTracking( threadId, ResourceStatsType.WORKER_STATS, currentThreadResourceMetrics(threadId)); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java index 237ac21ad7f..2d09bbf84d8 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java @@ -159,6 +159,21 @@ public void execute() {} new OpenSearchQueryManager(nodeClient, settings).submit(queryPlan); } + @Test + public void trackingIsSkippedWhenJvmDoesNotSupportIt() { + // Core refuses to track unless per-thread allocation is supported AND enabled + // (TaskResourceTrackingService.isTaskResourceTrackingSupported). We mirror that, so a PPL task + // never records stats on a JVM where core would record none. + TrackingTask task = new TrackingTask(true); + boolean started = OpenSearchQueryManager.startThreadResourceTracking(task, 1L); + if (OpenSearchQueryManager.isThreadResourceTrackingSupported()) { + assertTrue("supported JVM should track", started); + } else { + assertFalse("unsupported JVM must not track", started); + assertTrue("no entries recorded when unsupported", task.getResourceStats().isEmpty()); + } + } + @Test public void startThreadResourceTrackingReturnsFalseOnFailure() { CancellableTask task = mock(CancellableTask.class); From 872ba54756bbcd367edc30e292fbd8367c5fa84d Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 6 Oct 2026 12:42:10 +0000 Subject: [PATCH 10/26] refactor(ppl): Account PPL thread usage privately instead of via task resource stats Signed-off-by: Kishore Kumaar Natarajan --- .../executor/OpenSearchQueryManager.java | 89 +---------- .../ThreadPoolExecutionDispatcher.java | 18 +-- .../executor/ThreadResourceAccounting.java | 39 +++++ .../executor/OpenSearchQueryManagerTest.java | 139 +++++------------- .../ThreadPoolExecutionDispatcherTest.java | 24 ++- .../sql/plugin/transport/PPLQueryTask.java | 97 +++++++++--- .../transport/TransportPPLQueryAction.java | 32 ++-- .../plugin/transport/PPLQueryTaskTest.java | 69 ++++++++- 8 files changed, 249 insertions(+), 258 deletions(-) create mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadResourceAccounting.java diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java index 0b7a3ca1cd4..6734b68c5e0 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManager.java @@ -5,8 +5,6 @@ package org.opensearch.sql.opensearch.executor; -import com.sun.management.ThreadMXBean; -import java.lang.management.ManagementFactory; import java.util.Map; import lombok.RequiredArgsConstructor; import org.apache.logging.log4j.LogManager; @@ -14,9 +12,6 @@ import org.apache.logging.log4j.ThreadContext; import org.opensearch.OpenSearchTimeoutException; import org.opensearch.common.unit.TimeValue; -import org.opensearch.core.tasks.resourcetracker.ResourceStats; -import org.opensearch.core.tasks.resourcetracker.ResourceStatsType; -import org.opensearch.core.tasks.resourcetracker.ResourceUsageMetric; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.executor.QueryId; import org.opensearch.sql.executor.QueryManager; @@ -32,20 +27,6 @@ public class OpenSearchQueryManager implements QueryManager { private static final Logger LOG = LogManager.getLogger(OpenSearchQueryManager.class); - /** Samples per-thread CPU/memory for resource tracking; null when the JVM bean is unavailable. */ - private static final ThreadMXBean THREAD_MX_BEAN = resolveThreadMXBean(); - - private static ThreadMXBean resolveThreadMXBean() { - try { - if (ManagementFactory.getThreadMXBean() instanceof ThreadMXBean bean) { - return bean; - } - } catch (Exception e) { - LOG.warn("Per-thread resource metrics unavailable; PPL task resource tracking disabled", e); - } - return null; - } - private final NodeClient nodeClient; private final Settings settings; @@ -100,13 +81,10 @@ private void schedule( setCancellableTask(cancelTask); - // Bracket resource tracking for the inline path (runs on this thread). Script plans - // hand off to the complex-worker pool, which brackets itself. - final boolean trackResources = - cancelTask != null && cancelTask.supportsResourceTracking(); - final long trackedThreadId = Thread.currentThread().getId(); - final boolean trackingStarted = - trackResources && startThreadResourceTracking(cancelTask, trackedThreadId); + // Account this thread's work to the query. Script plans hand off to the + // complex-worker pool, which opens its own scope. + final ThreadResourceAccounting.Scope accounting = + ThreadResourceAccounting.enter(cancelTask); try { task.run(); @@ -125,9 +103,7 @@ private void schedule( throw e; } finally { - if (trackingStarted) { - stopThreadResourceTracking(cancelTask, trackedThreadId); - } + accounting.close(); clearCancellableTask(); } }); @@ -142,59 +118,4 @@ private Runnable withCurrentContext(final Runnable task) { task.run(); }; } - - /** - * Records the starting CPU/memory snapshot for {@code threadId}. - * - * @return true if tracking started; only then should {@link #stopThreadResourceTracking} be - * called. - */ - /** - * Whether the JVM can report per-thread allocation, mirroring core's {@code - * TaskResourceTrackingService.isTaskResourceTrackingSupported()}. Core refuses to track when this - * is false, so we must too: otherwise a PPL task would record stats on a JVM where core - * deliberately records none, and {@code getThreadAllocatedBytes} is not meaningful there. - */ - static boolean isThreadResourceTrackingSupported() { - return THREAD_MX_BEAN != null - && THREAD_MX_BEAN.isThreadAllocatedMemorySupported() - && THREAD_MX_BEAN.isThreadAllocatedMemoryEnabled(); - } - - static boolean startThreadResourceTracking(CancellableTask task, long threadId) { - if (!isThreadResourceTrackingSupported()) { - return false; - } - try { - task.startThreadResourceTracking( - threadId, ResourceStatsType.WORKER_STATS, currentThreadResourceMetrics(threadId)); - return true; - } catch (Exception e) { - LOG.warn("Failed to start resource tracking for task [{}]", task.getId(), e); - return false; - } - } - - /** Records the final CPU/memory snapshot for {@code threadId}. */ - static void stopThreadResourceTracking(CancellableTask task, long threadId) { - try { - task.stopThreadResourceTracking( - threadId, ResourceStatsType.WORKER_STATS, currentThreadResourceMetrics(threadId)); - } catch (Exception e) { - LOG.warn("Failed to stop resource tracking for task [{}]", task.getId(), e); - } - } - - /** Per-thread memory and CPU usage; empty array when the JVM bean is unavailable. */ - private static ResourceUsageMetric[] currentThreadResourceMetrics(long threadId) { - if (THREAD_MX_BEAN == null) { - return new ResourceUsageMetric[0]; - } - ResourceUsageMetric memory = - new ResourceUsageMetric( - ResourceStats.MEMORY, THREAD_MX_BEAN.getThreadAllocatedBytes(threadId)); - ResourceUsageMetric cpu = - new ResourceUsageMetric(ResourceStats.CPU, THREAD_MX_BEAN.getThreadCpuTime(threadId)); - return new ResourceUsageMetric[] {memory, cpu}; - } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java index e8f86ae6f36..e210368304a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcher.java @@ -91,9 +91,7 @@ private void dispatchInternal( ThreadPool.Names.GENERIC); Cancellable cancelPoller = scheduleCancellationPoller(cancellableTask, executionThread); Hook.Closeable hookHandle = null; - boolean trackResources = false; - long trackedThreadId = -1L; - boolean trackingStarted = false; + ThreadResourceAccounting.Scope accounting = ThreadResourceAccounting.Scope.NOOP; try { // Restore state from caller thread ThreadContext.putAll(ctx); @@ -109,14 +107,8 @@ private void dispatchInternal( hookHandle = Hook.CURRENT_TIME.addThread((Consumer>) h -> h.set(currentTime)); } - // Script plans do their real work on this thread, so bracket tracking here. - trackResources = - cancellableTask != null && cancellableTask.supportsResourceTracking(); - trackedThreadId = Thread.currentThread().getId(); - trackingStarted = - trackResources - && OpenSearchQueryManager.startThreadResourceTracking( - cancellableTask, trackedThreadId); + // Script plans do their real work on this thread, so account it to the query here. + accounting = ThreadResourceAccounting.enter(cancellableTask); task.run(); } catch (Exception e) { LOG.error("Exception during task execution on complex pool", e); @@ -130,9 +122,7 @@ private void dispatchInternal( if (hookHandle != null) { hookHandle.close(); } - if (trackingStarted) { - OpenSearchQueryManager.stopThreadResourceTracking(cancellableTask, trackedThreadId); - } + accounting.close(); OpenSearchQueryManager.clearCancellableTask(); RelMetadataQueryBase.THREAD_PROVIDERS.remove(); CalcitePlanContext.clearTimewrapSignals(); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadResourceAccounting.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadResourceAccounting.java new file mode 100644 index 00000000000..a7ebb88f03d --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadResourceAccounting.java @@ -0,0 +1,39 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.executor; + +/** + * Per-query accounting of the CPU and memory a query uses on the SQL plugin's own threads. + * + *

Core's task resource tracking only brackets threads on its task-aware pools, which a plugin + * pool cannot opt into, so the engine brackets its own threads instead. It does so through this + * interface rather than {@code Task.startThreadResourceTracking}, so nothing is written into the + * task's {@code resource_stats}: whether and how to measure is the task's decision, not the + * engine's. + */ +public interface ThreadResourceAccounting { + + /** + * Starts measuring the current thread for this query. The returned scope must be closed on the + * same thread when the thread's work for the query ends. + */ + Scope enterThread(); + + /** Opens a scope for {@code task}, or a no-op scope if the task does no accounting. */ + static Scope enter(Object task) { + return task instanceof ThreadResourceAccounting accounting + ? accounting.enterThread() + : Scope.NOOP; + } + + /** A measured stretch of one thread's work; closing it records the usage. */ + interface Scope extends AutoCloseable { + Scope NOOP = () -> {}; + + @Override + void close(); + } +} diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java index 2d09bbf84d8..396bf29f239 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java @@ -6,19 +6,16 @@ package org.opensearch.sql.opensearch.executor; import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.doAnswer; -import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import java.util.Collections; -import java.util.List; -import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -26,9 +23,6 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.common.unit.TimeValue; import org.opensearch.core.tasks.TaskId; -import org.opensearch.core.tasks.resourcetracker.ResourceStatsType; -import org.opensearch.core.tasks.resourcetracker.ResourceUsageMetric; -import org.opensearch.core.tasks.resourcetracker.ThreadResourceInfo; import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; @@ -97,105 +91,38 @@ public void clearTask() { } @Test - public void tracksResourceUsageWhenTaskSupportsIt() { - TrackingTask task = new TrackingTask(true); + public void accountingScopeSpansExecution() { + AccountingTask task = new AccountingTask(); OpenSearchQueryManager.setCancellableTask(task); - runSubmit(); + runSubmit(() -> task.openDuringExecute.set(task.openScopes.get() == 1), null); - Map> stats = task.getResourceStats(); - assertEquals(1, stats.size()); - List infos = stats.values().iterator().next(); - assertEquals(1, infos.size()); - assertFalse("resource entry should be closed after execution", infos.get(0).isActive()); + assertTrue("scope must be open while the plan executes", task.openDuringExecute.get()); + assertEquals(1, task.entered.get()); + assertEquals("scope must close after execution", 0, task.openScopes.get()); } @Test - public void skipsResourceTrackingWhenTaskDoesNotSupportIt() { - TrackingTask task = new TrackingTask(false); + public void accountingScopeClosesWhenExecutionThrows() { + AccountingTask task = new AccountingTask(); OpenSearchQueryManager.setCancellableTask(task); + RuntimeException boom = new RuntimeException("execute failed"); - runSubmit(); + assertThrows(RuntimeException.class, () -> runSubmit(null, boom)); - assertTrue("no resource entries expected", task.getResourceStats().isEmpty()); + assertEquals(1, task.entered.get()); + assertEquals(0, task.openScopes.get()); } @Test public void propagatesExecutionException() { RuntimeException boom = new RuntimeException("execute failed"); - RuntimeException thrown = - org.junit.jupiter.api.Assertions.assertThrows( - RuntimeException.class, () -> runSubmitThatThrows(boom)); + RuntimeException thrown = assertThrows(RuntimeException.class, () -> runSubmit(null, boom)); assertEquals(boom, thrown); } - /** Runs a trivial query through the manager with the schedule hook executing tasks inline. */ - private void runSubmit() { - NodeClient nodeClient = mock(NodeClient.class); - ThreadPool threadPool = mock(ThreadPool.class); - Settings settings = mock(Settings.class); - Scheduler.ScheduledCancellable mockScheduledTask = mock(Scheduler.ScheduledCancellable.class); - - when(nodeClient.threadPool()).thenReturn(threadPool); - - when(settings.getSettingValue(Settings.Key.PPL_QUERY_TIMEOUT)) - .thenReturn(TimeValue.timeValueSeconds(60)); - - AbstractPlan queryPlan = - new QueryPlan(queryId, queryType, plan, queryService, listener) { - @Override - public void execute() {} - }; - - doAnswer( - invocation -> { - Runnable task = invocation.getArgument(0); - task.run(); - return mockScheduledTask; - }) - .when(threadPool) - .schedule(any(), any(), any()); - - new OpenSearchQueryManager(nodeClient, settings).submit(queryPlan); - } - - @Test - public void trackingIsSkippedWhenJvmDoesNotSupportIt() { - // Core refuses to track unless per-thread allocation is supported AND enabled - // (TaskResourceTrackingService.isTaskResourceTrackingSupported). We mirror that, so a PPL task - // never records stats on a JVM where core would record none. - TrackingTask task = new TrackingTask(true); - boolean started = OpenSearchQueryManager.startThreadResourceTracking(task, 1L); - if (OpenSearchQueryManager.isThreadResourceTrackingSupported()) { - assertTrue("supported JVM should track", started); - } else { - assertFalse("unsupported JVM must not track", started); - assertTrue("no entries recorded when unsupported", task.getResourceStats().isEmpty()); - } - } - - @Test - public void startThreadResourceTrackingReturnsFalseOnFailure() { - CancellableTask task = mock(CancellableTask.class); - doThrow(new RuntimeException("bean failure")) - .when(task) - .startThreadResourceTracking( - anyLong(), any(ResourceStatsType.class), any(ResourceUsageMetric[].class)); - assertFalse(OpenSearchQueryManager.startThreadResourceTracking(task, 1L)); - } - - @Test - public void stopThreadResourceTrackingSwallowsFailure() { - CancellableTask task = mock(CancellableTask.class); - doThrow(new RuntimeException("bean failure")) - .when(task) - .stopThreadResourceTracking( - anyLong(), any(ResourceStatsType.class), any(ResourceUsageMetric[].class)); - OpenSearchQueryManager.stopThreadResourceTracking(task, 1L); - } - - /** Runs a query whose execute() throws, so the wrapped task's catch/finally path is exercised. */ - private void runSubmitThatThrows(RuntimeException toThrow) { + /** Runs a query through the manager, executing only the worker task inline. */ + private void runSubmit(Runnable onExecute, RuntimeException toThrow) { NodeClient nodeClient = mock(NodeClient.class); ThreadPool threadPool = mock(ThreadPool.class); Settings settings = mock(Settings.class); @@ -210,18 +137,19 @@ private void runSubmitThatThrows(RuntimeException toThrow) { new QueryPlan(queryId, queryType, plan, queryService, listener) { @Override public void execute() { - throw toThrow; + if (onExecute != null) { + onExecute.run(); + } + if (toThrow != null) { + throw toThrow; + } } }; - // Run only the worker task inline; do NOT run the timeout task (running it would interrupt the - // thread and turn the thrown exception into an OpenSearchTimeoutException). doAnswer( invocation -> { - Object executor = invocation.getArgument(2); - if ("sql-worker".equals(executor)) { - Runnable task = invocation.getArgument(0); - task.run(); + if ("sql-worker".equals(invocation.getArgument(2))) { + invocation.getArgument(0).run(); } return mockScheduledTask; }) @@ -231,18 +159,21 @@ public void execute() { new OpenSearchQueryManager(nodeClient, settings).submit(queryPlan); } - /** Minimal CancellableTask whose resource-tracking support is configurable. */ - private static class TrackingTask extends CancellableTask { - private final boolean supportsTracking; + /** CancellableTask that counts the accounting scopes the engine opens on it. */ + private static class AccountingTask extends CancellableTask implements ThreadResourceAccounting { + final AtomicInteger entered = new AtomicInteger(); + final AtomicInteger openScopes = new AtomicInteger(); + final AtomicBoolean openDuringExecute = new AtomicBoolean(); - TrackingTask(boolean supportsTracking) { + AccountingTask() { super(1L, "ppl", "action", "desc", TaskId.EMPTY_TASK_ID, Collections.emptyMap()); - this.supportsTracking = supportsTracking; } @Override - public boolean supportsResourceTracking() { - return supportsTracking; + public Scope enterThread() { + entered.incrementAndGet(); + openScopes.incrementAndGet(); + return openScopes::decrementAndGet; } @Override diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java index 510214cda4a..8a54ccd81fe 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java @@ -13,7 +13,6 @@ import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doThrow; @@ -21,6 +20,7 @@ import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; import static org.opensearch.common.settings.Settings.EMPTY; import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_COMPLEX_WORKER_THREAD_POOL_NAME; @@ -40,8 +40,6 @@ import org.mockito.quality.Strictness; import org.opensearch.common.unit.TimeValue; import org.opensearch.common.util.concurrent.ThreadContext; -import org.opensearch.core.tasks.resourcetracker.ResourceStatsType; -import org.opensearch.core.tasks.resourcetracker.ResourceUsageMetric; import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; @@ -378,14 +376,16 @@ void cancellableTaskAvailableDuringExecution() { } @Test - void bracketsResourceTrackingOnComplexPool() { + void opensAccountingScopeOnComplexPool() { when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) .thenReturn(true); when(settings.getSettingValue(Settings.Key.PPL_QUERY_TIMEOUT)) .thenReturn(new TimeValue(60000)); - CancellableTask trackingTask = mock(CancellableTask.class); - when(trackingTask.supportsResourceTracking()).thenReturn(true); - OpenSearchQueryManager.setCancellableTask(trackingTask); + CancellableTask accountingTask = + mock(CancellableTask.class, withSettings().extraInterfaces(ThreadResourceAccounting.class)); + ThreadResourceAccounting.Scope scope = mock(ThreadResourceAccounting.Scope.class); + when(((ThreadResourceAccounting) accountingTask).enterThread()).thenReturn(scope); + OpenSearchQueryManager.setCancellableTask(accountingTask); doAnswer( invocation -> { @@ -402,14 +402,8 @@ void bracketsResourceTrackingOnComplexPool() { AbstractCalciteIndexScan scan = createMockScanWithScripts(); dispatcher.dispatch(scan, context, listener, engine); - // start/stopThreadResourceTracking take (threadId, statsType, ResourceUsageMetric...) varargs. - verify(trackingTask).supportsResourceTracking(); - verify(trackingTask) - .startThreadResourceTracking( - anyLong(), any(ResourceStatsType.class), any(ResourceUsageMetric[].class)); - verify(trackingTask) - .stopThreadResourceTracking( - anyLong(), any(ResourceStatsType.class), any(ResourceUsageMetric[].class)); + verify((ThreadResourceAccounting) accountingTask).enterThread(); + verify(scope).close(); } @Test diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java index a463376d59e..e933fc8bc3e 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java @@ -5,12 +5,21 @@ package org.opensearch.sql.plugin.transport; +import com.sun.management.ThreadMXBean; +import java.lang.management.ManagementFactory; import java.util.List; import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.LongAdder; import org.opensearch.core.tasks.TaskId; +import org.opensearch.sql.opensearch.executor.ThreadResourceAccounting; import org.opensearch.tasks.CancellableTask; -public class PPLQueryTask extends CancellableTask { +public class PPLQueryTask extends CancellableTask implements ThreadResourceAccounting { + + /** Null when the JVM can't report per-thread CPU and allocation; accounting is then a no-op. */ + private static final ThreadMXBean THREAD_MX_BEAN = resolveThreadMXBean(); // The fields below are captured on the request thread and read later by the report listener // (which runs on a different thread), for the Query Insights record. @@ -30,12 +39,14 @@ public class PPLQueryTask extends CancellableTask { /** Whether the statement is an {@code explain}, which plans but runs no search. */ private volatile boolean queryInsightsExplain = false; - /** - * Whether per-thread CPU/memory accounting is enabled for this task. Off by default so a normal - * PPL query behaves like any {@link CancellableTask}; turned on only when Query Insights - * recording and core's task resource tracking are both enabled. - */ - private volatile boolean resourceTrackingEnabled = false; + /** Whether this query's thread usage is measured; off unless Query Insights will report it. */ + private volatile boolean resourceAccountingEnabled = false; + + private final LongAdder cpuNanos = new LongAdder(); + private final LongAdder allocatedBytes = new LongAdder(); + + /** Threads with an open scope, so a nested scope on the same thread isn't counted twice. */ + private final Set accountedThreads = ConcurrentHashMap.newKeySet(); public PPLQueryTask( long id, @@ -87,22 +98,74 @@ public boolean isQueryInsightsFailed() { return queryInsightsFailed; } - public void setResourceTrackingEnabled(boolean enabled) { - this.resourceTrackingEnabled = enabled; + public void setResourceAccountingEnabled(boolean enabled) { + this.resourceAccountingEnabled = enabled && THREAD_MX_BEAN != null; } - @Override - public boolean shouldCancelChildrenOnCancellation() { - return true; + /** CPU time used on the SQL plugin's threads; child searches are measured by core. */ + public long getCpuNanos() { + return cpuNanos.sum(); + } + + /** Bytes allocated on the SQL plugin's threads; child searches are measured by core. */ + public long getAllocatedBytes() { + return allocatedBytes.sum(); } /** - * Per-thread CPU/memory accounting for the coordinator task, enabled only when Query Insights - * recording and core's {@code task_resource_tracking.enabled} are both on (see {@code - * TransportPPLQueryAction}). {@link CancellableTask} defaults to {@code false}. + * Measures the current thread until the scope closes. The scope also holds the task's + * resource-tracking thread count, so the completion listener the Query Insights report hangs off + * fires only after every scope has recorded its usage. Nothing is written to the task's {@code + * resource_stats}, so {@code _tasks} output and {@code task_resource_tracking.enabled} are + * untouched. */ @Override - public boolean supportsResourceTracking() { - return resourceTrackingEnabled; + public Scope enterThread() { + if (!resourceAccountingEnabled) { + return Scope.NOOP; + } + final Thread thread = Thread.currentThread(); + final long threadId = thread.threadId(); + if (!accountedThreads.add(threadId)) { + return Scope.NOOP; + } + incrementResourceTrackingThreads(); + final long cpuStart = THREAD_MX_BEAN.getCurrentThreadCpuTime(); + final long allocatedStart = THREAD_MX_BEAN.getCurrentThreadAllocatedBytes(); + return () -> { + assert Thread.currentThread() == thread : "scope closed on a different thread"; + try { + addDelta(cpuNanos, cpuStart, THREAD_MX_BEAN.getCurrentThreadCpuTime()); + addDelta(allocatedBytes, allocatedStart, THREAD_MX_BEAN.getCurrentThreadAllocatedBytes()); + } finally { + accountedThreads.remove(threadId); + decrementResourceTrackingThreads(); + } + }; + } + + /** The bean returns -1 when the measurement is unsupported or disabled; record nothing then. */ + private static void addDelta(LongAdder total, long start, long end) { + if (start >= 0 && end >= start) { + total.add(end - start); + } + } + + private static ThreadMXBean resolveThreadMXBean() { + try { + if (ManagementFactory.getThreadMXBean() instanceof ThreadMXBean bean + && bean.isCurrentThreadCpuTimeSupported() + && bean.isThreadAllocatedMemorySupported()) { + return bean; + } + } catch (Exception e) { + // Fall through: accounting stays off. + } + return null; + } + + @Override + public boolean shouldCancelChildrenOnCancellation() { + return true; } } 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 7fd468559c8..ef51f51feb5 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 @@ -34,7 +34,6 @@ import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.ActionListener; import org.opensearch.core.action.NotifyOnceListener; -import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.common.utils.QueryContext; @@ -214,9 +213,9 @@ private boolean isQueryInsightsRecordingEnabled() { } /** - * Mirrors core's {@code task_resource_tracking.enabled} gate so PPL resource tracking is not - * enabled when an operator has disabled task resource tracking cluster-wide. Defaults to the core - * setting's own default ({@code true}) if the value can't be read. + * Whether to measure this query's thread usage. Follows core's {@code + * task_resource_tracking.enabled}: with it off, core doesn't measure the child searches either, + * so measuring only the PPL threads would produce a partial total. Treated as off if unreadable. */ private boolean isTaskResourceTrackingEnabled() { try { @@ -224,12 +223,15 @@ private boolean isTaskResourceTrackingEnabled() { .getClusterSettings() .get(TaskResourceTrackingService.TASK_RESOURCE_TRACKING_ENABLED); } catch (Exception e) { - LOG.debug("Failed to read task_resource_tracking.enabled; defaulting to enabled", e); - return true; + LOG.debug("Failed to read task_resource_tracking.enabled; not measuring PPL threads", e); + return false; } } - /** Report the finished PPL query once its resource tracking completes (stats are final). */ + /** + * Report the finished PPL query once the task is unregistered and every accounting scope has + * closed, so the measured usage is final. + */ private void registerQueryInsightsReport(PPLQueryTask pplQueryTask) { try { boolean registered = @@ -252,7 +254,6 @@ protected void innerOnFailure(Exception e) {} } } - /** Serialize a completed PPL query and send it to Query Insights; errors are ignored. */ /** * Whether a finished task should produce a Top N record. Decided from what the metadata sink * captured rather than from the route, because the route cannot see either case: {@code explain} @@ -264,6 +265,7 @@ static boolean shouldReportToQueryInsights(PPLQueryTask pplQueryTask) { && !pplQueryTask.isQueryInsightsExplain(); } + /** Serialize a completed PPL query and send it to Query Insights; errors are ignored. */ private void writeQueryInsightsRecord(PPLQueryTask pplQueryTask) { if (!shouldReportToQueryInsights(pplQueryTask)) { return; @@ -272,9 +274,8 @@ private void writeQueryInsightsRecord(PPLQueryTask pplQueryTask) { String nodeId = clusterServiceRef.localNode().getId(); String queryText = pplQueryTask.getQueryInsightsAnonymizedQuery(); - TaskResourceUsage usage = pplQueryTask.getTotalResourceStats(); - long cpuNanos = usage == null ? 0L : usage.getCpuTimeInNanos(); - long memoryBytes = usage == null ? 0L : usage.getMemoryInBytes(); + long cpuNanos = pplQueryTask.getCpuNanos(); + long memoryBytes = pplQueryTask.getAllocatedBytes(); // Wall-clock latency (monotonic); the profile total collapses to ~0 across the finish thread. long latencyMillis = @@ -435,12 +436,9 @@ protected void doExecute( // plain DSL search) with this query's marker. Child searches still see it: the executor // captures the thread context when the work is scheduled, which happens inside this block. try (ThreadContext.StoredContext ignored = threadContext.newStoredContext(true)) { - // Enable per-thread accounting only here, where a record will actually consume it, and - // only when core's task_resource_tracking is on so we never populate _tasks - // resource_stats behind the operator's back. The engine's brackets gate on - // supportsResourceTracking(), which stays false for every other route and when the - // feature is off, so they cost nothing then. - pplQueryTask.setResourceTrackingEnabled(isTaskResourceTrackingEnabled()); + // Measure thread usage only on this route, where a record consumes it. Every other route + // leaves accounting off, so the engine's scopes are no-ops there. + pplQueryTask.setResourceAccountingEnabled(isTaskResourceTrackingEnabled()); stampQueryInsightsParentHeader(pplQueryTask); registerQueryInsightsReport(pplQueryTask); executePplQuery(pplService, transformedRequest, callerListener, queryInsightsSink); diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java index 33178445deb..fa0ecf56a7d 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java @@ -8,8 +8,12 @@ import static org.junit.Assert.*; import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.Test; +import org.opensearch.core.action.NotifyOnceListener; import org.opensearch.core.tasks.TaskId; +import org.opensearch.sql.opensearch.executor.ThreadResourceAccounting; +import org.opensearch.tasks.Task; public class PPLQueryTaskTest { @@ -74,17 +78,68 @@ private PPLQueryTask newTask() { } @Test - public void testSupportsResourceTrackingOffByDefault() { - // Off by default so a PPL query adds no resource-tracking overhead unless Query Insights - // recording (and core's task resource tracking) are enabled. - assertFalse(newTask().supportsResourceTracking()); + public void testAccountingOffByDefault() { + PPLQueryTask task = newTask(); + assertSame(ThreadResourceAccounting.Scope.NOOP, task.enterThread()); + // Never touches core's tracking, so _tasks resource_stats stays empty. + assertFalse(task.supportsResourceTracking()); + assertTrue(task.getResourceStats().isEmpty()); + } + + @Test + public void testAccountingRecordsUsageWithoutTouchingResourceStats() { + PPLQueryTask task = newTask(); + task.setResourceAccountingEnabled(true); + try (ThreadResourceAccounting.Scope ignored = task.enterThread()) { + burnCpu(); + } + assertTrue(task.getCpuNanos() > 0); + assertTrue(task.getAllocatedBytes() > 0); + assertFalse(task.supportsResourceTracking()); + assertTrue(task.getResourceStats().isEmpty()); } @Test - public void testSupportsResourceTrackingWhenEnabled() { + public void testNestedScopeOnSameThreadIsNotCountedTwice() { PPLQueryTask task = newTask(); - task.setResourceTrackingEnabled(true); - assertTrue(task.supportsResourceTracking()); + task.setResourceAccountingEnabled(true); + try (ThreadResourceAccounting.Scope outer = task.enterThread()) { + assertSame(ThreadResourceAccounting.Scope.NOOP, task.enterThread()); + } + } + + @Test + public void testReportWaitsForOpenScopes() { + PPLQueryTask task = newTask(); + task.setResourceAccountingEnabled(true); + AtomicInteger fired = new AtomicInteger(); + assertTrue( + task.addResourceTrackingCompletionListener( + new NotifyOnceListener<>() { + @Override + protected void innerOnResponse(Task t) { + fired.incrementAndGet(); + } + + @Override + protected void innerOnFailure(Exception e) {} + })); + + ThreadResourceAccounting.Scope scope = task.enterThread(); + // What TaskManager.unregister does once the response is sent. + task.decrementResourceTrackingThreads(); + assertEquals("must not report while a thread is still accounting", 0, fired.get()); + + scope.close(); + assertEquals(1, fired.get()); + } + + private static void burnCpu() { + long x = 0; + for (int i = 0; i < 1_000_000; i++) { + x += Integer.toString(i).hashCode(); + } + assertNotEquals(42, x); } @Test From eb58f19e9a1889242909a2433a0d6479bd5c9d17 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 6 Oct 2026 15:43:27 -0700 Subject: [PATCH 11/26] fix(ppl): Stamp Top N records with query start time Query Insights stamps a native search record with the request's start time (SearchRequest#getOrCreateAbsoluteStartMillis), but the PPL record used the report time. A query that crossed a Top N window boundary was therefore counted in a later window than the child searches it spawned, so the parent-marker roll-up could not find them. Use the task's start time, which core already tracks on Task. Signed-off-by: Kishore Kumaar Natarajan --- .../sql/plugin/transport/TransportPPLQueryAction.java | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) 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 ef51f51feb5..9b5ad81ed0a 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 @@ -281,6 +281,12 @@ private void writeQueryInsightsRecord(PPLQueryTask pplQueryTask) { long latencyMillis = Math.max(0L, (System.nanoTime() - pplQueryTask.getStartTimeNanos()) / 1_000_000L); + // Start time, not report time: Query Insights stamps a search record with the request's + // start (SearchRequest#getOrCreateAbsoluteStartMillis). Stamping completion here would put a + // query that crosses a Top N window boundary in a later window than its own child searches, + // and the marker roll-up would then miss them. + long timestampMillis = pplQueryTask.getStartTime(); + // Same marker stamped on child DSL tasks, so their cpu/memory rolls up into this parent. String parentMarker = QueryInsightsMarker.value("PPL", nodeId, pplQueryTask.getId()); @@ -295,7 +301,7 @@ private void writeQueryInsightsRecord(PPLQueryTask pplQueryTask) { parentMarker, nodeId, queryText, - System.currentTimeMillis(), + timestampMillis, latencyMillis, cpuNanos, memoryBytes, From 46062e017e4676be4f85e7606e0c67dcf2f169fa Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 6 Oct 2026 15:47:17 -0700 Subject: [PATCH 12/26] fix(ppl): Skip the Top N record for a nested PPL query The parent marker already in the thread context is kept rather than replaced, because putHeader rejects a duplicate key and clearing the context first would drop the authenticated-user headers the child searches need. The record, however, was always built from this task's own id, so it advertised a marker that no child search references while the children's cost rolled up into the outer query's record instead. Skip the record in that case. No in-process caller reaches it today -- the only callers are RestPPLQueryAction and RestPPLGrammarAction, both on a fresh REST thread, and the header is registered via getTaskHeaders rather than getRestHeaders so a client cannot inject it -- but the two sides now agree instead of silently disagreeing. Signed-off-by: Kishore Kumaar Natarajan --- .../sql/plugin/transport/PPLQueryTask.java | 15 +++++++++++++++ .../transport/TransportPPLQueryAction.java | 13 ++++++++++++- .../plugin/transport/PPLQueryTaskTest.java | 19 +++++++++++++++++++ 3 files changed, 46 insertions(+), 1 deletion(-) diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java index e933fc8bc3e..4225487afab 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java @@ -39,6 +39,13 @@ public class PPLQueryTask extends CancellableTask implements ThreadResourceAccou /** Whether the statement is an {@code explain}, which plans but runs no search. */ private volatile boolean queryInsightsExplain = false; + /** + * Whether an outer query's parent marker was already in the thread context, so this query's child + * searches are attributed to that query rather than to this one. No in-process caller does this + * today; the flag keeps the record honest if one is ever added. + */ + private volatile boolean queryInsightsNested = false; + /** Whether this query's thread usage is measured; off unless Query Insights will report it. */ private volatile boolean resourceAccountingEnabled = false; @@ -90,6 +97,14 @@ public boolean isQueryInsightsExplain() { return queryInsightsExplain; } + public void setQueryInsightsNested(boolean nested) { + this.queryInsightsNested = nested; + } + + public boolean isQueryInsightsNested() { + return queryInsightsNested; + } + public void setQueryInsightsFailed(boolean failed) { this.queryInsightsFailed = failed; } 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 9b5ad81ed0a..62e702b3ffc 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 @@ -176,8 +176,14 @@ private void stampQueryInsightsParentHeader(PPLQueryTask pplQueryTask) { String value = QueryInsightsMarker.value( "PPL", clusterServiceRef.localNode().getId(), pplQueryTask.getId()); + // An outer marker is kept rather than replaced: putHeader rejects a duplicate key, and + // clearing the context first would drop the authenticated-user headers the child searches + // need. The child searches then belong to that outer query, so this one does not report a + // record -- see shouldReportToQueryInsights. if (threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER) == null) { threadContext.putHeader(QueryInsightsMarker.PARENT_HEADER, value); + } else { + pplQueryTask.setQueryInsightsNested(true); } // Capture the user here; the report listener runs on a thread without the transient. pplQueryTask.setQueryInsightsUserInfo( @@ -259,10 +265,15 @@ protected void innerOnFailure(Exception e) {} * captured rather than from the route, because the route cannot see either case: {@code explain} * is valid inside the query text (not just on {@code /_explain}), and a query that fails to parse * never reaches the sink at all, so it has no query text or indices to report. + * + *

A nested query is also skipped: its child searches carry the outer query's marker, so a + * record built from this task's own id would advertise a marker no child references, and the + * children's cost would be counted against both records. */ static boolean shouldReportToQueryInsights(PPLQueryTask pplQueryTask) { return pplQueryTask.getQueryInsightsAnonymizedQuery() != null - && !pplQueryTask.isQueryInsightsExplain(); + && !pplQueryTask.isQueryInsightsExplain() + && !pplQueryTask.isQueryInsightsNested(); } /** Serialize a completed PPL query and send it to Query Insights; errors are ignored. */ diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java index fa0ecf56a7d..6b86ca5ea36 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java @@ -179,6 +179,25 @@ public void testRuntimeFailureAfterPlanningIsReported() { assertTrue(TransportPPLQueryAction.shouldReportToQueryInsights(task)); } + @Test + public void testNestedQueryIsNotReported() { + // An outer marker was already in the context, so the child searches rolled up into that query. + // Reporting here would advertise a marker no child references and double count their cost. + PPLQueryTask task = newTask(); + task.setQueryInsightsAnonymizedQuery("source=table | where identifier > ***"); + task.setQueryInsightsNested(true); + assertFalse(TransportPPLQueryAction.shouldReportToQueryInsights(task)); + } + + @Test + public void testNonNestedQueryIsReportedByDefault() { + // Guards the default: a plain query must not be mistaken for a nested one. + PPLQueryTask task = newTask(); + task.setQueryInsightsAnonymizedQuery("source=table | where identifier > ***"); + assertFalse("nested must default to false", task.isQueryInsightsNested()); + assertTrue(TransportPPLQueryAction.shouldReportToQueryInsights(task)); + } + @Test public void testQueryInsightsParentHeaderName() { assertEquals("X-Query-Insights-Parent", QueryInsightsMarker.PARENT_HEADER); From c49b29c91e21fb326d6085116f123c92e93fa526 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Wed, 7 Oct 2026 00:13:24 -0700 Subject: [PATCH 13/26] test(ppl): Cover the analyze sink and exists subquery The analyze route builds its AST on a path of its own rather than sharing execute's, and nothing exercised it: the query-insights sink it feeds, the no-sink overload, and the parse-failure path were all untested. Add those three, asserting the sink stays untouched when the query never parses. Also cover the exists subquery, which reaches a different visit method than the already-tested in subquery. Signed-off-by: Kishore Kumaar Natarajan --- .../opensearch/sql/ppl/PPLServiceTest.java | 64 +++++++++++++++++++ .../ppl/utils/PPLQueryIndexExtractorTest.java | 8 +++ 2 files changed, 72 insertions(+) diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java index 68f88c6acd4..d8cb56706e2 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java @@ -21,6 +21,7 @@ import org.mockito.junit.MockitoJUnitRunner; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.executor.AnalyzeResponse; import org.opensearch.sql.executor.DefaultQueryManager; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.ExecutionEngine.ExplainResponse; @@ -81,6 +82,20 @@ public void onFailure(Exception e) { }; } + /** + * Tolerant of either outcome: these tests assert on the query-insights sink, which is fed before + * the analyze plan is submitted, and the plan itself runs against a mocked query service. + */ + private ResponseListener getAnalyzeListener() { + return new ResponseListener() { + @Override + public void onResponse(AnalyzeResponse response) {} + + @Override + public void onFailure(Exception e) {} + }; + } + private ResponseListener getExplainListener(boolean fail) { return new ResponseListener() { @Override @@ -156,6 +171,55 @@ public void testExplainPassesAnonymizedQueryToSink() { Assert.assertEquals(Collections.singletonList("t"), metadata.get().indices()); } + @Test + public void testAnalyzePassesAnonymizedQueryToSink() { + AtomicReference metadata = new AtomicReference<>(); + pplService.analyze( + new PPLQueryRequest("search source=t a=42", null, QUERY), + getAnalyzeListener(), + metadata::set); + + // analyze builds the AST on its own path rather than sharing execute's, so the sink has to be + // fed there too; without this the analyze route would report an empty record. + Assert.assertNotNull(metadata.get()); + Assert.assertTrue(metadata.get().anonymizedQuery().contains("***")); + Assert.assertFalse(metadata.get().anonymizedQuery().contains("42")); + Assert.assertEquals(Collections.singletonList("t"), metadata.get().indices()); + Assert.assertFalse(metadata.get().explain()); + } + + @Test + public void testAnalyzeWithoutSinkShouldPass() { + // The two-arg overload delegates with the no-op sink; it must not throw. + pplService.analyze( + new PPLQueryRequest("search source=t a=42", null, QUERY), getAnalyzeListener()); + } + + @Test + public void testAnalyzeWithIllegalQueryShouldBeCaughtByHandler() { + AtomicReference failure = new AtomicReference<>(); + AtomicReference metadata = new AtomicReference<>(); + pplService.analyze( + new PPLQueryRequest("search", null, QUERY), + new ResponseListener() { + @Override + public void onResponse(AnalyzeResponse response) { + Assert.fail("a query that fails to parse must not produce an analyze response"); + } + + @Override + public void onFailure(Exception e) { + failure.set(e); + } + }, + metadata::set); + + // The parse failure must reach the listener rather than escape analyze(), and the sink must + // stay untouched: there is no anonymized query to report for a query that never parsed. + Assert.assertNotNull(failure.get()); + Assert.assertNull(metadata.get()); + } + @Test public void testExecuteCsvFormatShouldPass() { doAnswer( diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java index d36faee2ce5..49cd9ecff94 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java @@ -94,6 +94,14 @@ public void evalScalarSubqueryCapturesInnerSource() { assertTrue(result.contains("inner")); } + @Test + public void existsSubqueryCapturesInnerSource() { + // exists takes a different AST node than `in`, so it needs its own visit to be reached. + List result = indices("source=outer | where exists [ source=inner | where a = b ]"); + assertTrue(result.contains("outer")); + assertTrue(result.contains("inner")); + } + @Test public void graphLookupCapturesFromTable() { List result = From 8ae79ba4806740b6c69412929e621c51278bbf45 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Wed, 7 Oct 2026 00:39:33 -0700 Subject: [PATCH 14/26] refactor(ppl): Collect query sources in one prepare step Both plan() and analyze() parsed the request and collected the source indices from the resulting AST, in two copies of the same block that had to stay in sync. Fold them into a single prepare() that parses once and returns the statement together with the metadata taken from that same AST, so the recorded indices are provably the ones that execute. Each route still supplies its own statement-builder context, which is the only part that differed. The repeated explain-format ternary becomes a helper. Signed-off-by: Kishore Kumaar Natarajan --- .../org/opensearch/sql/ppl/PPLService.java | 109 +++++++++--------- 1 file changed, 52 insertions(+), 57 deletions(-) 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 f8a2483a94f..0cdc457256f 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java @@ -130,41 +130,18 @@ public void analyze( ResponseListener listener, Consumer sink) { try { - String queryText = request.getRequest(); - Statement statement; - String anonymized; - // Transport-thread work — parse, AST build, anonymize. Trace-only: QueryProfiling isn't - // active yet on this thread. Cold-start ANTLR grammar init dominates this region. - try (ProfileScope preparePhase = ProfileScope.openTraceOnly("prepare")) { - try { - ParseTree cst = parser.parse(queryText); - statement = - cst.accept( - new AstStatementBuilder( - new AstBuilder(queryText, settings, request.getTimeBounds()), - AstStatementBuilder.StatementBuilderContext.builder() - .isExplain(false) - .fetchSize(request.getFetchSize()) - .highlightConfig(request.getHighlightConfig()) - .format( - request.getFormat() != null && !request.getFormat().isEmpty() - ? org.opensearch.sql.protocol.response.format.Format.ofExplain( - request.getFormat()) - .orElse(null) - : null) - .build())); - anonymized = anonymizer.anonymizeStatement(statement); - } catch (Exception e) { - preparePhase.setError(e); - throw e; - } - } - log.info("[{}] Incoming request {}", QueryContext.getRequestId(), anonymized); - - UnresolvedPlan unresolvedPlan = ((Query) statement).getPlan(); - sink.accept( - new QueryInsightsMetadata( - anonymized, extractIndexNames(statement), statement instanceof Explain)); + Prepared prepared = + prepare( + request, + AstStatementBuilder.StatementBuilderContext.builder() + .isExplain(false) + .fetchSize(request.getFetchSize()) + .highlightConfig(request.getHighlightConfig()) + .format(explainFormat(request)) + .build()); + sink.accept(prepared.metadata()); + + UnresolvedPlan unresolvedPlan = ((Query) prepared.statement()).getPlan(); AbstractPlan analyzePlan = queryExecutionFactory.createAnalyzePlan(unresolvedPlan, PPL_QUERY, listener); queryManager.submit(analyzePlan); @@ -178,6 +155,33 @@ private AbstractPlan plan( ResponseListener queryListener, ResponseListener explainListener, Consumer sink) { + Prepared prepared = + prepare( + request, + AstStatementBuilder.StatementBuilderContext.builder() + .isExplain(request.isExplainRequest()) + .fetchSize(request.getFetchSize()) + .highlightConfig(request.getHighlightConfig()) + .format(explainFormat(request)) + .explainMode(request.getExplainMode()) + .includeMetadata(request.getIncludeMetadata()) + .build()); + sink.accept(prepared.metadata()); + + return queryExecutionFactory.create(prepared.statement(), queryListener, explainListener); + } + + /** A parsed statement and the query-insights metadata collected from that same AST. */ + private record Prepared(Statement statement, QueryInsightsMetadata metadata) {} + + /** + * Parses the request once and collects the query-insights metadata from the AST it produced, so + * the recorded query text and source indices are the ones that execute. The caller supplies the + * statement-builder context because each route configures it differently. + */ + private Prepared prepare( + PPLQueryRequest request, AstStatementBuilder.StatementBuilderContext context) { + String queryText = request.getRequest(); Statement statement; String anonymized; // Transport-thread work — parse, AST build, anonymize. Trace-only: QueryProfiling isn't @@ -185,25 +189,11 @@ private AbstractPlan plan( try (ProfileScope preparePhase = ProfileScope.openTraceOnly("prepare")) { try { // 1. Parse query and convert parse tree (CST) to abstract syntax tree (AST) - ParseTree cst = parser.parse(request.getRequest()); - boolean includeMetadata = request.getIncludeMetadata(); + ParseTree cst = parser.parse(queryText); statement = cst.accept( new AstStatementBuilder( - new AstBuilder(request.getRequest(), settings, request.getTimeBounds()), - AstStatementBuilder.StatementBuilderContext.builder() - .isExplain(request.isExplainRequest()) - .fetchSize(request.getFetchSize()) - .highlightConfig(request.getHighlightConfig()) - .format( - request.getFormat() != null && !request.getFormat().isEmpty() - ? org.opensearch.sql.protocol.response.format.Format.ofExplain( - request.getFormat()) - .orElse(null) - : null) - .explainMode(request.getExplainMode()) - .includeMetadata(includeMetadata) - .build())); + new AstBuilder(queryText, settings, request.getTimeBounds()), context)); anonymized = anonymizer.anonymizeStatement(statement); } catch (RuntimeException e) { preparePhase.setError(e); @@ -211,15 +201,20 @@ private AbstractPlan plan( } } log.info("[{}] Incoming request {}", QueryContext.getRequestId(), anonymized); - - // Extract indices from the same AST that will execute, so recorded metadata matches the - // executed query and the query is parsed only once. - sink.accept( + return new Prepared( + statement, new QueryInsightsMetadata( anonymized, extractIndexNames(statement), statement instanceof Explain)); + } - AbstractPlan plan = queryExecutionFactory.create(statement, queryListener, explainListener); - return plan; + /** Explain format requested on the URL, or null when none was given. */ + private static org.opensearch.sql.protocol.response.format.Format explainFormat( + PPLQueryRequest request) { + if (request.getFormat() == null || request.getFormat().isEmpty()) { + return null; + } + return org.opensearch.sql.protocol.response.format.Format.ofExplain(request.getFormat()) + .orElse(null); } /** From 8ed8fa4e3eb6737a3269eac4720b1130ee8e4e21 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Wed, 7 Oct 2026 11:20:25 -0700 Subject: [PATCH 15/26] test(ppl): Cover the index-extraction fallbacks Neither fallback is reachable through the public API, since the parser only yields a Query or an Explain wrapping one, so cover them directly: make extractIndexNames package-private and assert a non-Query statement and a plan that throws mid-walk both yield no indices. The second is the one that matters: index extraction is best-effort metadata, and a plan that blows up during the walk must not fail the query that is about to execute. Signed-off-by: Kishore Kumaar Natarajan --- .../org/opensearch/sql/ppl/PPLService.java | 5 ++- .../opensearch/sql/ppl/PPLServiceTest.java | 35 +++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) 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 0cdc457256f..a7262d6aee9 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java @@ -220,8 +220,11 @@ private static org.opensearch.sql.protocol.response.format.Format explainFormat( /** * Best-effort source index names from an already-parsed statement (unwrapping an {@code explain} * to its inner query); empty on any failure. + * + *

Package-private so the fallbacks can be tested: neither is reachable through the public API, + * because the parser only ever yields a {@link Query} or an {@link Explain} wrapping one. */ - private static List extractIndexNames(Statement statement) { + static List extractIndexNames(Statement statement) { try { Statement inner = statement instanceof Explain ? ((Explain) statement).getStatement() : statement; diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java index d8cb56706e2..86c254a008a 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java @@ -19,6 +19,10 @@ import org.junit.runner.RunWith; import org.mockito.Mock; import org.mockito.junit.MockitoJUnitRunner; +import org.opensearch.sql.ast.AbstractNodeVisitor; +import org.opensearch.sql.ast.statement.Query; +import org.opensearch.sql.ast.statement.Statement; +import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.executor.AnalyzeResponse; @@ -27,6 +31,7 @@ import org.opensearch.sql.executor.ExecutionEngine.ExplainResponse; import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; import org.opensearch.sql.executor.QueryService; +import org.opensearch.sql.executor.QueryType; import org.opensearch.sql.executor.execution.QueryPlanFactory; import org.opensearch.sql.executor.pagination.Cursor; import org.opensearch.sql.ppl.antlr.PPLSyntaxParser; @@ -220,6 +225,36 @@ public void onFailure(Exception e) { Assert.assertNull(metadata.get()); } + @Test + public void testExtractIndexNamesOnNonQueryStatementYieldsEmpty() { + // Defensive fallback: the parser only produces Query or Explain, so this cannot happen through + // the public API. Pinned so the fallback stays a fallback and does not become a null return. + Statement notAQuery = new Statement() {}; + Assert.assertEquals(Collections.emptyList(), PPLService.extractIndexNames(notAQuery)); + } + + @Test + public void testExtractIndexNamesSwallowsWalkFailure() { + // Index extraction is best-effort metadata: a plan that blows up during the walk must yield no + // indices rather than failing the query that is about to execute. + UnresolvedPlan exploding = + new UnresolvedPlan() { + @Override + public UnresolvedPlan attach(UnresolvedPlan child) { + return this; + } + + @Override + public R accept(AbstractNodeVisitor visitor, C context) { + throw new IllegalStateException("walk blew up"); + } + }; + + Assert.assertEquals( + Collections.emptyList(), + PPLService.extractIndexNames(new Query(exploding, 0, QueryType.PPL))); + } + @Test public void testExecuteCsvFormatShouldPass() { doAnswer( From 0125171c781b2f2def73239eb4557a1bcf784f65 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Wed, 7 Oct 2026 11:26:33 -0700 Subject: [PATCH 16/26] test(ppl): Cover the extractor's missing-shape guards The walk guards against an absent sub-plan, filter condition, eval list and eval expression, because it produces best-effort metadata: a shape it did not expect must cost the executing query nothing. None are produced by the parser, so build them by hand and assert the walk still reaches the child relation. Also collapse the explain-format check to Strings.isNullOrEmpty. The null half was unreachable -- the REST layer always supplies a string, and the three-argument request constructor defaults it to empty -- so one condition says the same thing without a branch no test can take. Signed-off-by: Kishore Kumaar Natarajan --- .../org/opensearch/sql/ppl/PPLService.java | 3 +- .../ppl/utils/PPLQueryIndexExtractorTest.java | 52 +++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) 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 a7262d6aee9..93235459642 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java @@ -8,6 +8,7 @@ import static org.opensearch.sql.executor.ExecutionEngine.QueryResponse; import static org.opensearch.sql.executor.execution.QueryPlanFactory.NO_CONSUMER_RESPONSE_LISTENER; +import com.google.common.base.Strings; import java.util.List; import java.util.function.Consumer; import lombok.extern.log4j.Log4j2; @@ -210,7 +211,7 @@ private Prepared prepare( /** Explain format requested on the URL, or null when none was given. */ private static org.opensearch.sql.protocol.response.format.Format explainFormat( PPLQueryRequest request) { - if (request.getFormat() == null || request.getFormat().isEmpty()) { + if (Strings.isNullOrEmpty(request.getFormat())) { return null; } return org.opensearch.sql.protocol.response.format.Format.ofExplain(request.getFormat()) diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java index 49cd9ecff94..009353da608 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java @@ -8,13 +8,18 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; +import java.util.Arrays; +import java.util.Collections; import java.util.List; import org.junit.Test; import org.opensearch.sql.ast.expression.DataType; +import org.opensearch.sql.ast.expression.Field; +import org.opensearch.sql.ast.expression.Let; import org.opensearch.sql.ast.expression.Literal; import org.opensearch.sql.ast.expression.QualifiedName; import org.opensearch.sql.ast.tree.AppendCol; import org.opensearch.sql.ast.tree.AppendPipe; +import org.opensearch.sql.ast.tree.Eval; import org.opensearch.sql.ast.tree.Filter; import org.opensearch.sql.ast.tree.Relation; import org.opensearch.sql.ast.tree.UnresolvedPlan; @@ -145,4 +150,51 @@ private UnresolvedPlan subPipelineOver(String index) { public void nullPlanYieldsEmpty() { assertTrue(PPLQueryIndexExtractor.extractIndexNames(null).isEmpty()); } + + // The walk guards against absent sub-plans and sub-expressions because it runs on best-effort + // metadata: a shape it did not expect must cost the query nothing. The parser does not produce + // these shapes, so they are built by hand. + + @Test + public void nullSourceInGetSourcesIsSkipped() { + UnresolvedPlan declaresNullSource = + new UnresolvedPlan() { + @Override + public List getSources() { + return Collections.singletonList(null); + } + + @Override + public List getChild() { + return List.of(); + } + + @Override + public UnresolvedPlan attach(UnresolvedPlan child) { + return this; + } + }; + assertEquals(List.of(), PPLQueryIndexExtractor.extractIndexNames(declaresNullSource)); + } + + @Test + public void filterWithoutConditionStillWalksItsChild() { + UnresolvedPlan plan = new Filter(null).attach(baseRelation()); + assertEquals(List.of("base"), PPLQueryIndexExtractor.extractIndexNames(plan)); + } + + @Test + public void evalWithoutExpressionsStillWalksItsChild() { + UnresolvedPlan plan = new Eval(null).attach(baseRelation()); + assertEquals(List.of("base"), PPLQueryIndexExtractor.extractIndexNames(plan)); + } + + @Test + public void evalWithNullAssignmentOrExpressionStillWalksItsChild() { + // Covers both halves of the guard: a null Let, and a Let carrying no expression. + Let noExpression = new Let(new Field(new QualifiedName("x")), null); + List lets = Arrays.asList(null, noExpression); + UnresolvedPlan plan = new Eval(lets).attach(baseRelation()); + assertEquals(List.of("base"), PPLQueryIndexExtractor.extractIndexNames(plan)); + } } From 3d0ab55c8ac3b8348e4494e3a8015c360958cecf Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Thu, 8 Oct 2026 02:41:10 -0700 Subject: [PATCH 17/26] fix(ppl): Send no parent marker when stamping it failed The record rebuilt the marker from the task id regardless of whether the stamp had succeeded, so a failed stamp produced a record advertising a marker no child search carried. Store the marker on the task only once it is in the thread context and send whatever was stored; the reporter already writes a null marker as empty, which Query Insights treats as a parent with no children. This also removes the duplicate computation. Review cleanup in the same pass: drop the metadataTarget alias of an already-final local; in the dispatcher test remove an unused thread- context stub that only passed under LENIENT and revert the log4j static-import churn it forced; drop two tests of pre-existing dispatcher behaviour and make the third assert the accounting scope closes on the exception path, which is the line this change added; remove a stray blank line; correct the settings doc to 3.10; and reword the security IT so it claims only what it exercises -- Query Insights is not installed there, so the send never runs. Signed-off-by: Kishore Kumaar Natarajan --- docs/user/ppl/admin/settings.md | 2 +- .../QueryInsightsReportingSecurityIT.java | 26 +++--- .../executor/OpenSearchQueryManagerTest.java | 1 - .../ThreadPoolExecutionDispatcherTest.java | 93 ++++--------------- .../sql/plugin/transport/PPLQueryTask.java | 14 +++ .../transport/TransportPPLQueryAction.java | 22 +++-- .../plugin/transport/PPLQueryTaskTest.java | 10 ++ 7 files changed, 68 insertions(+), 100 deletions(-) diff --git a/docs/user/ppl/admin/settings.md b/docs/user/ppl/admin/settings.md index 28986cd732d..c7551c64bae 100644 --- a/docs/user/ppl/admin/settings.md +++ b/docs/user/ppl/admin/settings.md @@ -478,7 +478,7 @@ Expected output: ### Version -3.9 +3.10 ### Description diff --git a/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsReportingSecurityIT.java b/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsReportingSecurityIT.java index 0ec7903d63c..5ea3f66479a 100644 --- a/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsReportingSecurityIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsReportingSecurityIT.java @@ -14,19 +14,20 @@ import org.junit.Test; /** - * Security tests for the PPL → Query Insights reporting path on an FGAC (security-enabled) cluster. + * Runs PPL as a low-privilege user on the FGAC (security-enabled) cluster with the Query Insights + * reporting code present. * - *

Reporting a completed PPL query is an internal {@code cluster:admin/...} action the end user - * never invokes; the send is wrapped in {@code stashContext()} so it runs user-less rather than - * being authorized against the querying user (see {@code QueryInsightsReporter}). These tests pin - * that contract: a low-privilege user who holds only PPL + index-read permissions — and explicitly - * NOT the report action's {@code cluster:admin} permission — can still run PPL queries, including - * the multi-scan (join) path that stamps the parent marker for child sub-query tagging. + *

Query Insights is not installed in this cluster, so the reporting gate ({@code + * search.insights.top_queries.ppl.enabled}) is unregistered and the record is never sent. What + * these tests do pin is the non-interference contract on a secured cluster: a user holding only PPL + * + index-read permissions, and explicitly not the report action's {@code cluster:admin} + * permission, can run PPL queries unchanged. * - *

Read-side visibility of the recorded user identity is governed by the Query Insights plugin's - * own RBAC filter ({@code search.insights.top_queries.filter_by_mode}); that plugin is not - * installed in the SQL integ-test cluster, so it is covered by Query Insights' own tests rather - * than here. + *

The send itself is wrapped in {@code stashContext()} so it runs user-less rather than being + * authorized against the querying user (see {@code QueryInsightsReporter}); proving that holds for + * a low-privilege user needs Query Insights co-installed and belongs with that plugin's tests, as + * does read-side visibility of the recorded identity ({@code + * search.insights.top_queries.filter_by_mode}). */ public class QueryInsightsReportingSecurityIT extends SecurityTestBase { @@ -62,8 +63,7 @@ protected void init() throws Exception { @Test public void lowPrivilegeUserCanRunPplQueryWithoutReportPermission() throws IOException { - // The report send is stashed to Origin.LOCAL, so lacking the report cluster:admin permission - // must not cause the user's own query to fail. + // Lacking the report action's cluster:admin permission must not affect the user's own query. JSONObject result = executeQueryAsUser( String.format("source=%s | where age > 30 | fields firstname", TEST_INDEX_BANK), diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java index 396bf29f239..a0d9fa1f945 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java @@ -58,7 +58,6 @@ public void submitQuery() { Scheduler.ScheduledCancellable mockScheduledTask = mock(Scheduler.ScheduledCancellable.class); when(nodeClient.threadPool()).thenReturn(threadPool); - when(settings.getSettingValue(Settings.Key.PPL_QUERY_TIMEOUT)) .thenReturn(TimeValue.timeValueSeconds(60)); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java index 8a54ccd81fe..665e0bf6fcb 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ThreadPoolExecutionDispatcherTest.java @@ -5,13 +5,9 @@ package org.opensearch.sql.opensearch.executor; -import static org.apache.logging.log4j.ThreadContext.clearAll; -import static org.apache.logging.log4j.ThreadContext.get; -import static org.apache.logging.log4j.ThreadContext.put; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doAnswer; @@ -21,7 +17,6 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.mockito.Mockito.withSettings; -import static org.opensearch.common.settings.Settings.EMPTY; import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_COMPLEX_WORKER_THREAD_POOL_NAME; import java.util.List; @@ -30,6 +25,7 @@ import org.apache.calcite.rel.RelVisitor; import org.apache.calcite.rel.metadata.JaninoRelMetadataProvider; import org.apache.calcite.rel.metadata.RelMetadataQueryBase; +import org.apache.logging.log4j.ThreadContext; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -39,7 +35,6 @@ import org.mockito.junit.jupiter.MockitoSettings; import org.mockito.quality.Strictness; import org.opensearch.common.unit.TimeValue; -import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.sql.calcite.CalcitePlanContext; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; @@ -66,7 +61,6 @@ class ThreadPoolExecutionDispatcherTest { @BeforeEach void setUp() { dispatcher = new ThreadPoolExecutionDispatcher(threadPool, settings); - when(threadPool.getThreadContext()).thenReturn(new ThreadContext(EMPTY)); // Mock schedule calls to return non-null cancellables (for both outer dispatch and inner // timeout) when(threadPool.schedule(any(Runnable.class), any(TimeValue.class), any())) @@ -77,7 +71,7 @@ void setUp() { @AfterEach void tearDown() { - clearAll(); + ThreadContext.clearAll(); OpenSearchQueryManager.clearCancellableTask(); RelMetadataQueryBase.THREAD_PROVIDERS.remove(); CalcitePlanContext.clearTimewrapSignals(); @@ -182,8 +176,8 @@ void propagatesLog4jThreadContextToSlowPool() { .thenReturn(true); when(settings.getSettingValue(Settings.Key.PPL_QUERY_TIMEOUT)) .thenReturn(new TimeValue(60000)); - put("request.id", "test-123"); - put("user", "admin"); + ThreadContext.put("request.id", "test-123"); + ThreadContext.put("user", "admin"); AtomicReference requestIdOnSlowPool = new AtomicReference<>(); AtomicReference userOnSlowPool = new AtomicReference<>(); @@ -191,10 +185,10 @@ void propagatesLog4jThreadContextToSlowPool() { invocation -> { Runnable task = invocation.getArgument(0); // Simulate a different thread — clear MDC - clearAll(); + ThreadContext.clearAll(); task.run(); - requestIdOnSlowPool.set(get("request.id")); - userOnSlowPool.set(get("user")); + requestIdOnSlowPool.set(ThreadContext.get("request.id")); + userOnSlowPool.set(ThreadContext.get("user")); return mock(ScheduledCancellable.class); }) .when(threadPool) @@ -407,73 +401,17 @@ void opensAccountingScopeOnComplexPool() { } @Test - void cancellationPollerInterruptsWhenTaskCancelled() { + void accountingScopeClosesOnComplexPoolWhenTaskThrows() { when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) .thenReturn(true); when(settings.getSettingValue(Settings.Key.PPL_QUERY_TIMEOUT)) .thenReturn(new TimeValue(60000)); - CancellableTask cancelledTask = mock(CancellableTask.class); - when(cancelledTask.isCancelled()).thenReturn(true); - OpenSearchQueryManager.setCancellableTask(cancelledTask); - - // Capture the poller runnable so we can fire it directly and assert it interrupts the thread. - AtomicReference poller = new AtomicReference<>(); - when(threadPool.scheduleWithFixedDelay(any(Runnable.class), any(TimeValue.class), any())) - .thenAnswer( - invocation -> { - poller.set(invocation.getArgument(0)); - return mock(Cancellable.class); - }); - - AtomicReference interruptedInExecute = new AtomicReference<>(false); - doAnswer( - invocation -> { - Object executor = invocation.getArgument(2); - if (SQL_COMPLEX_WORKER_THREAD_POOL_NAME.equals(executor)) { - invocation.getArgument(0).run(); - } - return mock(ScheduledCancellable.class); - }) - .when(threadPool) - .schedule(any(Runnable.class), any(TimeValue.class), any()); - doAnswer( - invocation -> { - // Fire the poller on this (execution) thread; it should interrupt us. - poller.get().run(); - interruptedInExecute.set(Thread.currentThread().isInterrupted()); - Thread.interrupted(); // clear so we don't leak the interrupt into the pool thread - return null; - }) - .when(engine) - .execute(any(RelNode.class), any(), any(ResponseListener.class)); - - AbstractCalciteIndexScan scan = createMockScanWithScripts(); - dispatcher.dispatch(scan, context, listener, engine); - - assertTrue( - interruptedInExecute.get(), - "cancellation poller should interrupt the execution thread when the task is cancelled"); - } - - @Test - void dispatchTaskRunsInlineWhenNoScripts() { - when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) - .thenReturn(true); - RelNode plan = createMockNode(); - AtomicReference ran = new AtomicReference<>(false); - - dispatcher.dispatchTask(plan, context, () -> ran.set(true)); - - assertTrue(ran.get()); - verify(threadPool, never()).schedule(any(), any(TimeValue.class), any()); - } + CancellableTask accountingTask = + mock(CancellableTask.class, withSettings().extraInterfaces(ThreadResourceAccounting.class)); + ThreadResourceAccounting.Scope scope = mock(ThreadResourceAccounting.Scope.class); + when(((ThreadResourceAccounting) accountingTask).enterThread()).thenReturn(scope); + OpenSearchQueryManager.setCancellableTask(accountingTask); - @Test - void dispatchTaskSwallowsFailureWhenNoListener() { - when(settings.getSettingValue(Settings.Key.SQL_COMPLEX_WORKER_POOL_ENABLED)) - .thenReturn(true); - when(settings.getSettingValue(Settings.Key.PPL_QUERY_TIMEOUT)) - .thenReturn(new TimeValue(60000)); doAnswer( invocation -> { Object executor = invocation.getArgument(2); @@ -486,7 +424,6 @@ void dispatchTaskSwallowsFailureWhenNoListener() { .schedule(any(Runnable.class), any(TimeValue.class), any()); AbstractCalciteIndexScan scan = createMockScanWithScripts(); - dispatcher.dispatchTask( scan, context, @@ -494,7 +431,9 @@ void dispatchTaskSwallowsFailureWhenNoListener() { throw new RuntimeException("boom"); }); - // No exception escaped and the task was cleared: the null-listener catch + finally ran. + // The scope is closed from the finally, so a failing plan still records its usage, and the + // null-listener catch keeps the exception from escaping. + verify(scope).close(); assertNull(OpenSearchQueryManager.getCancellableTask()); } diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java index 4225487afab..7644cd574af 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java @@ -46,6 +46,12 @@ public class PPLQueryTask extends CancellableTask implements ThreadResourceAccou */ private volatile boolean queryInsightsNested = false; + /** + * The {@code ::} marker stamped on this query's child searches; null when + * it was never stamped, so the record does not advertise a marker no child carries. + */ + private volatile String queryInsightsParentMarker; + /** Whether this query's thread usage is measured; off unless Query Insights will report it. */ private volatile boolean resourceAccountingEnabled = false; @@ -105,6 +111,14 @@ public boolean isQueryInsightsNested() { return queryInsightsNested; } + public void setQueryInsightsParentMarker(String parentMarker) { + this.queryInsightsParentMarker = parentMarker; + } + + public String getQueryInsightsParentMarker() { + return queryInsightsParentMarker; + } + public void setQueryInsightsFailed(boolean failed) { this.queryInsightsFailed = failed; } 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 bf44b48e746..84eab99d835 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 @@ -188,7 +188,12 @@ private void buildUnifiedQueryHandlerIfReady() { } } - /** Stamp the parent marker so child DSL searches carry it back to this PPL query. */ + /** + * Stamp the parent marker so child DSL searches carry it back to this PPL query. The marker is + * stored on the task only once it is actually in the thread context, and the record later sends + * whatever was stored: if stamping fails the record goes out with no marker rather than one that + * no child search carries. + */ private void stampQueryInsightsParentHeader(PPLQueryTask pplQueryTask) { try { ThreadContext threadContext = clientRef.threadPool().getThreadContext(); @@ -201,6 +206,7 @@ private void stampQueryInsightsParentHeader(PPLQueryTask pplQueryTask) { // record -- see shouldReportToQueryInsights. if (threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER) == null) { threadContext.putHeader(QueryInsightsMarker.PARENT_HEADER, value); + pplQueryTask.setQueryInsightsParentMarker(value); } else { pplQueryTask.setQueryInsightsNested(true); } @@ -317,8 +323,9 @@ private void writeQueryInsightsRecord(PPLQueryTask pplQueryTask) { // and the marker roll-up would then miss them. long timestampMillis = pplQueryTask.getStartTime(); - // Same marker stamped on child DSL tasks, so their cpu/memory rolls up into this parent. - String parentMarker = QueryInsightsMarker.value("PPL", nodeId, pplQueryTask.getId()); + // The marker stamped on child DSL tasks, so their cpu/memory rolls up into this parent; null + // when stamping failed, which the reporter sends as an empty marker. + String parentMarker = pplQueryTask.getQueryInsightsParentMarker(); List indices = pplQueryTask.getQueryInsightsIndices(); String userInfo = pplQueryTask.getQueryInsightsUserInfo(); @@ -436,14 +443,13 @@ protected void doExecute( return; } - final PPLQueryTask metadataTarget = pplQueryTask; Consumer queryInsightsSink = metadata -> { rootSpan.addAttribute("db.query.text", metadata.anonymizedQuery()); - if (metadataTarget != null) { - metadataTarget.setQueryInsightsAnonymizedQuery(metadata.anonymizedQuery()); - metadataTarget.setQueryInsightsIndices(metadata.indices()); - metadataTarget.setQueryInsightsExplain(metadata.explain()); + if (pplQueryTask != null) { + pplQueryTask.setQueryInsightsAnonymizedQuery(metadata.anonymizedQuery()); + pplQueryTask.setQueryInsightsIndices(metadata.indices()); + pplQueryTask.setQueryInsightsExplain(metadata.explain()); } }; PPLService pplService = injector.getInstance(PPLService.class); diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java index 6b86ca5ea36..2ab513ea482 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java @@ -189,6 +189,16 @@ public void testNestedQueryIsNotReported() { assertFalse(TransportPPLQueryAction.shouldReportToQueryInsights(task)); } + @Test + public void testParentMarkerIsAbsentUntilStamped() { + // The record sends whatever marker was stamped. A task whose stamp never happened must carry + // none, so the record does not advertise a marker no child search references. + PPLQueryTask task = newTask(); + assertNull(task.getQueryInsightsParentMarker()); + task.setQueryInsightsParentMarker("PPL:node-1:42"); + assertEquals("PPL:node-1:42", task.getQueryInsightsParentMarker()); + } + @Test public void testNonNestedQueryIsReportedByDefault() { // Guards the default: a plain query must not be mistaken for a nested one. From f12f52f23e01a6301bd82e3a10e2f3199ab7b99c Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Thu, 8 Oct 2026 02:41:10 -0700 Subject: [PATCH 18/26] test(ppl): Pin the Query Insights wire format The reader lives in a different plugin and the two share no classes, so the byte layout is the contract, and nothing on this side exercised the writer. Pin it three ways: a golden payload that fails on any one-sided layout change, a field-by-field decode in the consumer's read order that documents each field and rejects trailing bytes, and the null/negative normalisation. Also pin the send itself: it runs with the caller's context stashed so the internal cluster:admin action is not authorised against the end user, the context is restored afterwards, and a transport failure never reaches the query that just finished. Signed-off-by: Kishore Kumaar Natarajan --- .../transport/QueryInsightsReporterTest.java | 210 ++++++++++++++++++ 1 file changed, 210 insertions(+) create mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsReporterTest.java diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsReporterTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsReporterTest.java new file mode 100644 index 00000000000..28df747d915 --- /dev/null +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsReporterTest.java @@ -0,0 +1,210 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Base64; +import java.util.List; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.Before; +import org.junit.Test; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.common.bytes.BytesReference; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.BytesTransportRequest; +import org.opensearch.transport.TransportService; + +/** + * Pins the wire format {@link QueryInsightsReporter} sends to Query Insights. The reader lives in a + * different plugin ({@code ReportQueryBytesRequestHandler}) and the two share no classes, so the + * byte layout is the contract. The golden-bytes test fails on any one-sided change to the layout; + * the consumer-order test documents what each field is and would catch a reorder. + */ +public class QueryInsightsReporterTest { + + private static final String ACTION = "cluster:admin/opensearch/query_insights/report_query_bytes"; + + /** + * Payload for the inputs in {@link #reportSample}; recompute only on a deliberate format bump. + */ + private static final String GOLDEN_BASE64 = + "AQNQUEwQUFBMOm5vZGUtMToxMjM0NQZub2RlLTElc291cmNlPXRhYmxlIHwgd2hlcmUgaWRlbnRpZmllciA+ICoqKo" + + "DQlf+8MYgIwMjmxgLwzaunAQIGZW50aXR5CWVtcGxveWVlcxphZG1pbnx8YWxsX2FjY2Vzc3xfX3VzZXJfXwE="; + + private TransportService transportService; + private DiscoveryNode localNode; + private ThreadContext threadContext; + private final AtomicReference captured = new AtomicReference<>(); + private final AtomicReference actionSeen = new AtomicReference<>(); + private final AtomicReference headerSeenDuringSend = new AtomicReference<>("unset"); + + @Before + public void setUp() { + transportService = mock(TransportService.class); + localNode = mock(DiscoveryNode.class); + threadContext = new ThreadContext(Settings.EMPTY); + ThreadPool threadPool = mock(ThreadPool.class); + when(threadPool.getThreadContext()).thenReturn(threadContext); + when(transportService.getThreadPool()).thenReturn(threadPool); + doAnswer( + invocation -> { + actionSeen.set(invocation.getArgument(1)); + captured.set(invocation.getArgument(2)); + headerSeenDuringSend.set(threadContext.getHeader("caller-header")); + return null; + }) + .when(transportService) + .sendRequest(eq(localNode), any(String.class), any(), any()); + } + + /** Fixed inputs covering every field, including a multi-index list and a user string. */ + private void reportSample() { + QueryInsightsReporter.report( + transportService, + localNode, + "PPL", + "PPL:node-1:12345", + "node-1", + "source=table | where identifier > ***", + 1_700_000_000_000L, + 1032L, + 685_352_000L, + 350_938_864L, + Arrays.asList("entity", "employees"), + "admin||all_access|__user__", + true); + } + + @Test + public void goldenBytes() { + reportSample(); + assertEquals(ACTION, actionSeen.get()); + byte[] bytes = BytesReference.toBytes(captured.get().bytes()); + assertEquals( + "wire layout changed; if deliberate, bump FORMAT_VERSION on both plugins and update the" + + " golden value", + GOLDEN_BASE64, + Base64.getEncoder().encodeToString(bytes)); + } + + @Test + public void decodesInConsumerOrder() throws IOException { + reportSample(); + // Mirrors ReportQueryBytesRequestHandler#deserialize field for field. + try (StreamInput in = captured.get().bytes().streamInput()) { + assertEquals(QueryInsightsReporter.FORMAT_VERSION, in.readVInt()); + assertEquals("PPL", in.readString()); // querySource + assertEquals("PPL:node-1:12345", in.readString()); // coordinatorId / parent marker + assertEquals("node-1", in.readString()); // nodeId + assertEquals("source=table | where identifier > ***", in.readString()); // queryText + assertEquals(1_700_000_000_000L, in.readVLong()); // timestampMillis + assertEquals(1032L, in.readVLong()); // latencyMillis + assertEquals(685_352_000L, in.readVLong()); // cpuNanos + assertEquals(350_938_864L, in.readVLong()); // memoryBytes + int indexCount = in.readVInt(); + String[] indices = new String[indexCount]; + for (int i = 0; i < indexCount; i++) { + indices[i] = in.readString(); + } + assertArrayEquals(new String[] {"entity", "employees"}, indices); + assertEquals("admin||all_access|__user__", in.readString()); // userInfo + assertTrue(in.readBoolean()); // failed + assertEquals("trailing bytes would be misread by the consumer", 0, in.available()); + } + } + + @Test + public void formatVersionIsOne() { + // The consumer rejects any other value; bumping this is a cross-plugin change. + assertEquals(1, QueryInsightsReporter.FORMAT_VERSION); + } + + @Test + public void nullsBecomeEmptyStringsAndNegativesClampToZero() throws IOException { + QueryInsightsReporter.report( + transportService, localNode, null, null, null, null, 5L, -1L, -2L, -3L, null, null, false); + try (StreamInput in = captured.get().bytes().streamInput()) { + in.readVInt(); + assertEquals("", in.readString()); + assertEquals("", in.readString()); + assertEquals("", in.readString()); + assertEquals("", in.readString()); + assertEquals(5L, in.readVLong()); + assertEquals(0L, in.readVLong()); + assertEquals(0L, in.readVLong()); + assertEquals(0L, in.readVLong()); + assertEquals(0, in.readVInt()); // null index list -> empty + assertEquals("", in.readString()); + assertFalse(in.readBoolean()); + } + } + + @Test + public void nullIndexEntryBecomesEmptyString() throws IOException { + QueryInsightsReporter.report( + transportService, + localNode, + "PPL", + "m", + "n", + "q", + 1L, + 1L, + 1L, + 1L, + Arrays.asList("a", null, "b"), + "", + false); + try (StreamInput in = captured.get().bytes().streamInput()) { + in.readVInt(); + for (int i = 0; i < 4; i++) { + in.readString(); + } + for (int i = 0; i < 4; i++) { + in.readVLong(); + } + assertEquals(3, in.readVInt()); + assertEquals( + List.of("a", "", "b"), List.of(in.readString(), in.readString(), in.readString())); + } + } + + @Test + public void sendRunsWithCallerContextStashed() { + // The send is an internal cluster:admin action; it must not be authorized against the end + // user, so the caller's context is stashed for the duration of the send and restored after. + threadContext.putHeader("caller-header", "present"); + reportSample(); + assertNull("send must not see the caller's context", headerSeenDuringSend.get()); + assertEquals( + "caller's context must be restored", "present", threadContext.getHeader("caller-header")); + } + + @Test + public void sendFailureDoesNotPropagate() { + doThrow(new RuntimeException("transport down")) + .when(transportService) + .sendRequest(eq(localNode), any(String.class), any(), any()); + // Reporting is best-effort; a failure here must never surface to the query that just finished. + reportSample(); + } +} From 1ccbe68c441253d87c7b1595e82b814550858e17 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Thu, 8 Oct 2026 02:57:38 -0700 Subject: [PATCH 19/26] fix(ppl): Collect subquery sources from every command that evaluates one Only where and eval conditions were bridged into expressions, so a subquery anywhere else was dropped from the recorded indices. Probing each expression-bearing command against a live cluster showed the gap is real: stats, chart, timechart, fillnull and foreach all parse and execute a scalar subquery in their argument, and the documented join ON form (on l.a = r.a AND r.a in [ source = inner ... ]) was missed as well. Bridge all of them, plus eventstats/streamstats for symmetry. Two walk fixes surfaced doing this. Node.getChild() defaults to null and Alias keeps that default, so descending into any aggregation argument NPEd and the function beneath the alias was never reached; the walk now treats null as no children and visits an alias's delegate. Null-safety is concentrated in one visit() helper so each bridge is a single call. Every plan node holding a plan outside its child list was re-checked: all either override getSources() or expose it through getChild(). The class comment previously claimed join ON was unsupported. Signed-off-by: Kishore Kumaar Natarajan --- .../sql/ppl/utils/PPLQueryIndexExtractor.java | 139 ++++++++++++-- .../ppl/utils/PPLQueryIndexExtractorTest.java | 175 ++++++++++++++++++ 2 files changed, 302 insertions(+), 12 deletions(-) diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java b/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java index 44c1af0fd09..298eaeb8e9a 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java @@ -9,18 +9,28 @@ import java.util.LinkedHashSet; import java.util.List; import java.util.Set; +import org.apache.commons.lang3.tuple.Pair; import org.opensearch.sql.ast.AbstractNodeVisitor; import org.opensearch.sql.ast.Node; +import org.opensearch.sql.ast.expression.Alias; +import org.opensearch.sql.ast.expression.Field; import org.opensearch.sql.ast.expression.Let; import org.opensearch.sql.ast.expression.QualifiedName; +import org.opensearch.sql.ast.expression.UnresolvedExpression; import org.opensearch.sql.ast.expression.subquery.ExistsSubquery; import org.opensearch.sql.ast.expression.subquery.InSubquery; import org.opensearch.sql.ast.expression.subquery.ScalarSubquery; +import org.opensearch.sql.ast.tree.Aggregation; +import org.opensearch.sql.ast.tree.Chart; import org.opensearch.sql.ast.tree.Eval; +import org.opensearch.sql.ast.tree.FillNull; import org.opensearch.sql.ast.tree.Filter; +import org.opensearch.sql.ast.tree.Foreach; import org.opensearch.sql.ast.tree.Join; import org.opensearch.sql.ast.tree.Relation; +import org.opensearch.sql.ast.tree.StreamWindow; import org.opensearch.sql.ast.tree.UnresolvedPlan; +import org.opensearch.sql.ast.tree.Window; /** * Collects the source index name(s) a PPL query reads from by walking the parsed AST (via {@link @@ -32,9 +42,13 @@ * sub-search, a lookup table, a join's right branch). A new command therefore only has to declare * its sources on the node itself to be covered here. * - *

Subqueries carried in expressions are not reachable that way, so {@link Filter} and - * {@link Eval} conditions are bridged explicitly. A subquery in a {@link Join}'s ON condition is - * still not walked and is dropped. + *

Subqueries carried in expressions are not reachable that way, so every command whose + * expression the engine evaluates is bridged explicitly: {@link Filter} and {@link Join} ON + * conditions, {@link Eval}, {@link Foreach} and {@link FillNull} assignments, and {@link + * Aggregation}, {@link Window}, {@link StreamWindow} and {@link Chart} function arguments. A + * subquery in any of these positions parses and, except under {@code eventstats}, executes, so its + * inner source is a real read and belongs in the metadata. Commands that take only field names, + * literals or patterns cannot carry one. * *

Returns distinct names in encounter order. This is best-effort metadata; a source it cannot * resolve is simply omitted and query execution is unaffected. @@ -65,10 +79,20 @@ public Void visitRelation(Relation node, Set names) { * Walks {@code getChild()} as usual, then the node's declared {@link * UnresolvedPlan#getSources()}. Commands whose source lives in a field are covered without this * collector knowing each one. + * + *

{@code Node.getChild()} defaults to null and several expression nodes keep that default, + * so the walk treats null as "no children" rather than relying on the base visitor. */ @Override public Void visitChildren(Node node, Set names) { - super.visitChildren(node, names); + List children = node.getChild(); + if (children != null) { + for (Node child : children) { + if (child != null) { + child.accept(this, names); + } + } + } if (node instanceof UnresolvedPlan plan) { for (UnresolvedPlan source : plan.getSources()) { if (source != null) { @@ -79,14 +103,15 @@ public Void visitChildren(Node node, Set names) { return null; } + // Expression bridges. Each command below evaluates an expression that may hold a subquery; the + // plan walk does not descend into expressions, so hand each one to visit(). Null-safety lives + // in + // visit()/visitAll() alone so a hand-built AST with a missing part costs nothing. + @Override public Void visitFilter(Filter node, Set names) { - // The condition can hold a subquery (where id in [ source=b ... ]); the plan walk doesn't - // descend into expressions, so bridge into it here. visitChildren(node, names); - if (node.getCondition() != null) { - node.getCondition().accept(this, names); - } + visit(node.getCondition(), names); return null; } @@ -95,14 +120,104 @@ public Void visitEval(Eval node, Set names) { visitChildren(node, names); if (node.getExpressionList() != null) { for (Let let : node.getExpressionList()) { - if (let != null && let.getExpression() != null) { - let.getExpression().accept(this, names); - } + visit(let == null ? null : let.getExpression(), names); + } + } + return null; + } + + @Override + public Void visitAlias(Alias node, Set names) { + // Alias keeps Node's null getChild(), and every aggregation argument and most eval + // assignments arrive wrapped in one, so a subquery beneath it would otherwise be invisible. + visit(node.getDelegated(), names); + return null; + } + + @Override + public Void visitAggregation(Aggregation node, Set names) { + visitChildren(node, names); + visitAll(node.getAggExprList(), names); + visitAll(node.getGroupExprList(), names); + visitAll(node.getSortExprList(), names); + visit(node.getSpan(), names); + return null; + } + + @Override + public Void visitWindow(Window node, Set names) { + visitChildren(node, names); + visitAll(node.getWindowFunctionList(), names); + visitAll(node.getGroupList(), names); + return null; + } + + @Override + public Void visitStreamWindow(StreamWindow node, Set names) { + visitChildren(node, names); + visitAll(node.getWindowFunctionList(), names); + visitAll(node.getGroupList(), names); + visit(node.getResetBefore(), names); + visit(node.getResetAfter(), names); + return null; + } + + @Override + public Void visitChart(Chart node, Set names) { + visitChildren(node, names); + visit(node.getAggregationFunction(), names); + visit(node.getRowSplit(), names); + visit(node.getColumnSplit(), names); + return null; + } + + @Override + public Void visitFillNull(FillNull node, Set names) { + visitChildren(node, names); + visit(node.getReplacementForAll().orElse(null), names); + if (node.getReplacementPairs() != null) { + for (Pair pair : node.getReplacementPairs()) { + visit(pair == null ? null : pair.getRight(), names); + } + } + return null; + } + + @Override + public Void visitForeach(Foreach node, Set names) { + visitChildren(node, names); + visit(node.getCollectionExpression(), names); + if (node.getEvalClauses() != null) { + for (Foreach.ForeachEvalClause clause : node.getEvalClauses()) { + visit(clause == null ? null : clause.getExpression(), names); } } return null; } + @Override + public Void visitJoin(Join node, Set names) { + // Both sides come through getChild() + getSources(); the ON condition is an expression and + // can carry a subquery (on l.a = r.a AND r.a in [ source=inner | ... ]). + visitChildren(node, names); + visit(node.getJoinCondition() == null ? null : node.getJoinCondition().orElse(null), names); + return null; + } + + private void visit(UnresolvedExpression expr, Set names) { + if (expr != null) { + expr.accept(this, names); + } + } + + private void visitAll(List exprs, Set names) { + if (exprs != null) { + for (UnresolvedExpression expr : exprs) { + visit(expr, names); + } + } + } + @Override public Void visitScalarSubquery(ScalarSubquery node, Set names) { node.getQuery().accept(this, names); diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java index 009353da608..8641fb8f914 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java @@ -11,16 +11,23 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Optional; +import org.apache.commons.lang3.tuple.Pair; import org.junit.Test; import org.opensearch.sql.ast.expression.DataType; import org.opensearch.sql.ast.expression.Field; import org.opensearch.sql.ast.expression.Let; import org.opensearch.sql.ast.expression.Literal; import org.opensearch.sql.ast.expression.QualifiedName; +import org.opensearch.sql.ast.expression.UnresolvedExpression; +import org.opensearch.sql.ast.tree.Aggregation; import org.opensearch.sql.ast.tree.AppendCol; import org.opensearch.sql.ast.tree.AppendPipe; import org.opensearch.sql.ast.tree.Eval; +import org.opensearch.sql.ast.tree.FillNull; import org.opensearch.sql.ast.tree.Filter; +import org.opensearch.sql.ast.tree.Foreach; +import org.opensearch.sql.ast.tree.Join; import org.opensearch.sql.ast.tree.Relation; import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.ppl.AstPlanningTestBase; @@ -99,6 +106,89 @@ public void evalScalarSubqueryCapturesInnerSource() { assertTrue(result.contains("inner")); } + // Every command whose expression the engine evaluates can carry a subquery. These were found by + // probing each expression-bearing command; all but eventstats/streamstats execute on a live + // cluster, so the inner source is a real read. + + @Test + public void joinOnConditionSubqueryCapturesInnerSource() { + // Documented form: `on l.a = r.a AND r.a in [ source = inner | fields d ]` (subquery.md). + List result = + indices( + "source = t1 | inner join left = l right = r" + + " on l.a = r.a AND r.a in [ source = inner | fields d ] t2"); + assertEquals(List.of("t1", "t2", "inner"), result); + } + + @Test + public void joinOnConditionAndRightSubsearchBothCaptured() { + List result = + indices( + "source = t1 | inner join left = l right = r on l.a = r.a AND r.a in [ source = inner |" + + " fields d ] [ source = t2 | fields a ]"); + assertEquals(List.of("t1", "t2", "inner"), result); + } + + @Test + public void statsArgumentSubqueryCapturesInnerSource() { + List result = indices("source=outer | stats max([ source=inner | stats max(b) ])"); + assertEquals(List.of("outer", "inner"), result); + } + + @Test + public void eventstatsArgumentSubqueryCapturesInnerSource() { + List result = indices("source=outer | eventstats max([ source=inner | stats max(b) ])"); + assertEquals(List.of("outer", "inner"), result); + } + + @Test + public void streamstatsArgumentSubqueryCapturesInnerSource() { + List result = + indices("source=outer | streamstats max([ source=inner | stats max(b) ])"); + assertEquals(List.of("outer", "inner"), result); + } + + @Test + public void chartArgumentSubqueryCapturesInnerSource() { + List result = indices("source=outer | chart max([ source=inner | stats max(b) ]) by c"); + assertEquals(List.of("outer", "inner"), result); + } + + @Test + public void timechartArgumentSubqueryCapturesInnerSource() { + List result = indices("source=outer | timechart max([ source=inner | stats max(b) ])"); + assertEquals(List.of("outer", "inner"), result); + } + + @Test + public void fillnullReplacementSubqueryCapturesInnerSource() { + assertEquals( + List.of("outer", "inner"), + indices("source=outer | fillnull with [ source=inner | stats max(b) ] in a")); + assertEquals( + List.of("outer", "inner"), + indices("source=outer | fillnull value=[ source=inner | stats max(b) ]")); + assertEquals( + List.of("outer", "inner"), + indices("source=outer | fillnull using a = [ source=inner | stats max(b) ]")); + } + + @Test + public void foreachAssignmentSubqueryCapturesInnerSource() { + List result = + indices("source=outer | foreach a [ eval x = [ source=inner | stats max(b) ] ]"); + assertEquals(List.of("outer", "inner"), result); + } + + @Test + public void aliasedExpressionIsWalkedThrough() { + // Alias keeps Node's null getChild(); without an explicit bridge the aggregate beneath it, and + // any subquery in its argument, is invisible and the walk itself NPEs. + List result = + indices("source=outer | stats max([ source=inner | stats max(b) ]) as m by c"); + assertEquals(List.of("outer", "inner"), result); + } + @Test public void existsSubqueryCapturesInnerSource() { // exists takes a different AST node than `in`, so it needs its own visit to be reached. @@ -189,6 +279,91 @@ public void evalWithoutExpressionsStillWalksItsChild() { assertEquals(List.of("base"), PPLQueryIndexExtractor.extractIndexNames(plan)); } + @Test + public void planWithNullChildListIsWalkedAsLeaf() { + // Node.getChild() defaults to null; a plan that keeps the default must read as "no children", + // not NPE, and still contribute its declared sources. + UnresolvedPlan nullChildren = + new UnresolvedPlan() { + @Override + public List getSources() { + return List.of(baseRelation()); + } + + @Override + public UnresolvedPlan attach(UnresolvedPlan child) { + return this; + } + }; + assertEquals(List.of("base"), PPLQueryIndexExtractor.extractIndexNames(nullChildren)); + } + + @Test + public void nullEntryInChildListIsSkipped() { + UnresolvedPlan holeyChildren = + new UnresolvedPlan() { + @Override + public List getChild() { + return Arrays.asList(null, baseRelation()); + } + + @Override + public UnresolvedPlan attach(UnresolvedPlan child) { + return this; + } + }; + assertEquals(List.of("base"), PPLQueryIndexExtractor.extractIndexNames(holeyChildren)); + } + + @Test + public void aggregationWithNullExpressionListsStillWalksItsChild() { + UnresolvedPlan plan = new Aggregation(null, null, null).attach(baseRelation()); + assertEquals(List.of("base"), PPLQueryIndexExtractor.extractIndexNames(plan)); + } + + @Test + public void fillnullWithNullOrMissingReplacementsStillWalksItsChild() { + assertEquals( + List.of("base"), + PPLQueryIndexExtractor.extractIndexNames( + FillNull.ofVariousValue(null).attach(baseRelation()))); + List> holey = Arrays.asList(null, Pair.of(null, null)); + assertEquals( + List.of("base"), + PPLQueryIndexExtractor.extractIndexNames( + FillNull.ofVariousValue(holey).attach(baseRelation()))); + } + + @Test + public void foreachWithNullOrMissingClausesStillWalksItsChild() { + assertEquals( + List.of("base"), + PPLQueryIndexExtractor.extractIndexNames( + new Foreach(null, null, null, null, null).attach(baseRelation()))); + List holey = + Arrays.asList(null, new Foreach.ForeachEvalClause("x", null)); + assertEquals( + List.of("base"), + PPLQueryIndexExtractor.extractIndexNames( + new Foreach(null, null, null, null, holey).attach(baseRelation()))); + } + + @Test + public void joinWithNullConditionStillWalksBothSides() { + UnresolvedPlan plan = + new Join( + new Relation(new QualifiedName("right")), + Optional.empty(), + Optional.empty(), + Join.JoinType.INNER, + null, + new Join.JoinHint(), + Optional.empty(), + new org.opensearch.sql.ast.expression.Argument.ArgumentMap(List.of())) + .attach(baseRelation()); + assertEquals(List.of("base", "right"), PPLQueryIndexExtractor.extractIndexNames(plan)); + } + @Test public void evalWithNullAssignmentOrExpressionStillWalksItsChild() { // Covers both halves of the guard: a null Let, and a Let carrying no expression. From 1025fd0595a4bbeb050c79f7cd04ad0f3b32b1ff Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Fri, 16 Oct 2026 00:24:38 -0700 Subject: [PATCH 20/26] test(ppl): Assert PPL queries are recorded into Query Insights The existing ITs only prove PPL still runs when Query Insights is absent; nothing asserted that a record is written. Add a cluster with Query Insights installed and an IT that checks, against the real plugin: - a query is recorded with query_source=PPL, its indices, a :: parent marker and non-zero cpu/memory - a missing index and an unknown field are recorded as failed - explain (endpoint and inline), syntax errors and async submissions write nothing; neither does a query with the gate off, while its search is still recorded as an ordinary top-level query - every search a join spawns is filed as a DSL child of the record, reads the indices the record names, and is not stamped earlier than it Each scenario reads its own index, since recorded text is anonymized, and absence is asserted only after a later record has arrived. Runs via :integ-test:integTestWithQueryInsights and stays out of check until the published Query Insights snapshot includes the PPL consumer; until then pass -PqueryInsightsZip=. Signed-off-by: Kishore Kumaar Natarajan --- integ-test/build.gradle | 56 +++ .../sql/ppl/QueryInsightsRecordingIT.java | 331 ++++++++++++++++++ 2 files changed, 387 insertions(+) create mode 100644 integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsRecordingIT.java diff --git a/integ-test/build.gradle b/integ-test/build.gradle index 8bf61e5ae11..56756bea831 100644 --- a/integ-test/build.gradle +++ b/integ-test/build.gradle @@ -205,6 +205,11 @@ configurations.all { configurations { zipArchive + // Resolved only by integTestWithQueryInsights, so no other task depends on a Query Insights + // snapshot being published. + queryInsightsZip { + transitive = false + } } dependencies { @@ -238,6 +243,9 @@ dependencies { // For GeoIP PPL functions zipArchive group: 'org.opensearch.plugin', name:'geospatial', version: "${opensearch_build}" + // Installed alongside the plugin by integTestWithQueryInsights + queryInsightsZip group: 'org.opensearch.plugin', name:'query-insights', version: "${opensearch_build}" + } java { @@ -281,6 +289,17 @@ def getGeoSpatialPlugin() { } } +// The published Query Insights snapshot, or a local build via -PqueryInsightsZip=/path/to.zip. +def getQueryInsightsPlugin() { + provider { (RegularFile) (() -> + project.findProperty('queryInsightsZip') + ? file(project.findProperty('queryInsightsZip')) + : configurations.queryInsightsZip.asFileTree.matching { + include '**/query-insights*' + }.singleFile ) + } +} + // fetch from the feature-build artifact for now (linux/x64 only; for local dev pass -PanalyticsEngineZip=/path instead). ext.pluginVersion = opensearch_version.tokenize('-')[0] ext.featureBuildBase = "https://ci.opensearch.org/ci/dbc/feature-build-opensearch/feature-datafusion/latest/linux/x64/tar/builds/opensearch/plugins" @@ -540,6 +559,14 @@ testClusters { plugin(getJobSchedulerPlugin()) plugin ":opensearch-sql-plugin" } + // Query Insights installed, so the PPL -> Query Insights reporting path actually runs. + integTestWithQueryInsights { + testDistribution = 'archive' + plugin(getJobSchedulerPlugin()) + plugin(getQueryInsightsPlugin()) + plugin ":opensearch-sql-plugin" + setting "plugins.query.datasources.encryption.masterkey", "1234567812345678" + } // Smoke test: verify sql loads cleanly alongside analytics-engine. analyticsEngineCompat { testDistribution = 'archive' @@ -795,6 +822,29 @@ task asyncMultiNodeIT(type: RestIntegTestTask) { // (needs a 2-node cluster) and would never actually run. check.dependsOn asyncMultiNodeIT +// Asserts that PPL queries are recorded into Query Insights. Not part of `check` yet: the published +// Query Insights snapshot does not include the PPL consumer until that change merges there, so run +// it against a local build until then: +// ./gradlew :integ-test:integTestWithQueryInsights -PqueryInsightsZip=/path/to/query-insights.zip +task integTestWithQueryInsights(type: RestIntegTestTask) { + useCluster testClusters.integTestWithQueryInsights + + dependsOn ':opensearch-sql-plugin:bundlePlugin' + + systemProperty 'tests.security.manager', 'false' + systemProperty 'project.root', project.projectDir.absolutePath + systemProperty 'defaultQuerySizeLimit', '10000' + systemProperty 'defaultQueryBucketSize', '1000' + systemProperty 'defaultMaxResultWindow', '10000' + doFirst { + systemProperty 'cluster.debug', getDebug() + } + + filter { + includeTestsMatching 'org.opensearch.sql.ppl.QueryInsightsRecordingIT' + } +} + task integTestWithSecurity(type: RestIntegTestTask) { useCluster testClusters.integTestWithSecurity useCluster testClusters.remoteIntegTestWithSecurity @@ -949,6 +999,9 @@ integTest { // Run on dedicated two-node cluster via :integ-test:asyncMultiNodeIT. exclude 'org/opensearch/sql/ppl/AsyncPPLMultiNodeRoutingIT.class' + // Needs Query Insights installed; runs via :integ-test:integTestWithQueryInsights. + exclude 'org/opensearch/sql/ppl/QueryInsightsRecordingIT.class' + // Exclude tests that require the analytics engine plugin stack (run separately via dedicated tasks) exclude 'org/opensearch/sql/analytics/AnalyticsEngineProfileIT.class' @@ -1300,6 +1353,9 @@ task integTestRemote(type: RestIntegTestTask) { excludeTestsMatching "org.opensearch.sql.bwc.*IT" } + // Needs Query Insights installed; runs via :integ-test:integTestWithQueryInsights. + excludeTestsMatching 'org.opensearch.sql.ppl.QueryInsightsRecordingIT' + if (analyticsEnabled) { // Full-text search (search-filter syntax -> Lucene query_string) needs an inverted-index // reader that parquet-backed analytics indices lack, so it's unsupported on this route. diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsRecordingIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsRecordingIT.java new file mode 100644 index 00000000000..aa0cab7e097 --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsRecordingIT.java @@ -0,0 +1,331 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ppl; + +import static org.opensearch.sql.legacy.TestUtils.getResponseBody; +import static org.opensearch.sql.plugin.rest.RestPPLQueryAction.QUERY_API_ENDPOINT; + +import java.io.IOException; +import java.time.Instant; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Predicate; +import org.json.JSONArray; +import org.json.JSONObject; +import org.junit.Test; +import org.opensearch.client.Request; +import org.opensearch.client.ResponseException; + +/** + * End to end: PPL queries are recorded into Query Insights Top N on a cluster that has Query + * Insights installed. + * + *

Runs only in {@code :integ-test:integTestWithQueryInsights}, whose cluster installs Query + * Insights; the default integTest cluster does not, and {@link QueryInsightsReportingIT} covers + * that case. + * + *

Each scenario reads its own index, so a record is identified by the index it reports rather + * than by its query text, which is anonymized. Absence is only asserted after a later query's + * record has arrived, so a missing record means it was never sent rather than not yet processed. + */ +public class QueryInsightsRecordingIT extends PPLIntegTestCase { + + private static final String PPL_ENABLED = "search.insights.top_queries.ppl.enabled"; + + /** Records are stamped with their query's start, so everything this class records is later. */ + private static final long CLASS_START_MILLIS = System.currentTimeMillis(); + + private static final AtomicInteger BARRIER_SEQ = new AtomicInteger(); + + @Override + public void init() throws Exception { + super.init(); + enableCalcite(); // join runs on the Calcite path + assertQueryInsightsSupportsPpl(); + updateClusterSettings(new ClusterSetting(PERSISTENT, PPL_ENABLED, "true")); + updateClusterSettings( + new ClusterSetting(PERSISTENT, "search.insights.top_queries.latency.enabled", "true")); + // Room for every query this class runs, so none is evicted from the window it is read from. + updateClusterSettings( + new ClusterSetting(PERSISTENT, "search.insights.top_queries.latency.top_n_size", "100")); + // The smallest window, so records reach the exported store within a minute. + updateClusterSettings( + new ClusterSetting(PERSISTENT, "search.insights.top_queries.latency.window_size", "1m")); + } + + @Test + public void successfulQueryIsRecorded() throws Exception { + indexWithDocs("qi_it_success"); + executeQuery("source=qi_it_success | where v > 1 | fields v"); + + JSONObject record = awaitPplRecord("qi_it_success"); + assertFalse(record.getBoolean("failed")); + assertEquals(Set.of("qi_it_success"), indicesOf(record)); + assertTrue( + "parent marker must be ::, was " + + record.optString("parent_marker"), + record.optString("parent_marker").matches("PPL:[^:]+:\\d+")); + assertFalse("query text must be recorded", record.optString("source").isEmpty()); + assertTrue(record.getLong("timestamp") >= CLASS_START_MILLIS); + // The plugin measures its own threads; zero here means the accounting never ran. + assertTrue(measurement(record, "cpu") > 0); + assertTrue(measurement(record, "memory") > 0); + } + + @Test + public void queryOnMissingIndexIsRecordedAsFailed() throws Exception { + expectThrows(ResponseException.class, () -> executeQuery("source=qi_it_missing | fields v")); + + // Fails in analysis, before any search runs; the record still names the index it targeted. + JSONObject record = awaitPplRecord("qi_it_missing"); + assertTrue(record.getBoolean("failed")); + } + + @Test + public void queryOnUnknownFieldIsRecordedAsFailed() throws Exception { + indexWithDocs("qi_it_bad_field"); + expectThrows( + ResponseException.class, + () -> executeQuery("source=qi_it_bad_field | where no_such_field > 1 | fields v")); + + JSONObject record = awaitPplRecord("qi_it_bad_field"); + assertTrue(record.getBoolean("failed")); + } + + @Test + public void explainIsNotRecorded() throws Exception { + indexWithDocs("qi_it_explain"); + explainQueryToString("source=qi_it_explain | fields v"); + // explain is also valid inside the query text, which reaches the execute route. + executeQuery("explain source=qi_it_explain | fields v"); + + awaitBarrier(); + assertNoPplRecord("qi_it_explain"); + } + + @Test + public void syntaxErrorIsNotRecorded() throws Exception { + expectThrows(ResponseException.class, () -> executeQuery("source=qi_it_syntax | wherr v > 1")); + + awaitBarrier(); + // A query that never parses has no text or indices to report. Every other query in this class + // names an index, so a PPL record without one could only be the syntax error. + for (JSONObject record : topQueries(QueryInsightsRecordingIT::isPpl)) { + assertFalse("unparsed query was recorded: " + record, indicesOf(record).isEmpty()); + } + } + + @Test + public void nothingIsRecordedWhenTheGateIsOff() throws Exception { + indexWithDocs("qi_it_gated"); + updateClusterSettings(new ClusterSetting(PERSISTENT, PPL_ENABLED, "false")); + try { + executeQuery("source=qi_it_gated | where v > 1 | fields v"); + } finally { + updateClusterSettings(new ClusterSetting(PERSISTENT, PPL_ENABLED, "true")); + } + + awaitBarrier(); + assertNoPplRecord("qi_it_gated"); + // The search it ran is still recorded as an ordinary query: with the gate off no parent marker + // is stamped, so it is nobody's child and stays in the Top N overview. + JSONObject search = + awaitRecord(r -> !isPpl(r) && indicesOf(r).contains("qi_it_gated"), "native search"); + assertFalse(search.optBoolean("is_child", false)); + } + + @Test + public void asyncSubmissionIsNotRecorded() throws Exception { + indexWithDocs("qi_it_async"); + Request request = new Request("POST", QUERY_API_ENDPOINT); + request.setJsonEntity( + "{\"query\": \"source=qi_it_async | fields v\", \"wait_for_completion_timeout\": \"30s\"}"); + client().performRequest(request); + + // The query outlives its task on this route, so a record built when the task ends would carry + // near-zero latency; the route is left out of Top N. + awaitBarrier(); + assertNoPplRecord("qi_it_async"); + } + + @Test + public void childSearchesCarryTheParentMarker() throws Exception { + indexWithDocs("qi_it_join_left"); + indexWithDocs("qi_it_join_right"); + executeQuery( + "source=qi_it_join_left | join left=l right=r on l.k = r.k qi_it_join_right" + + " | fields l.v, r.v"); + + JSONObject parent = awaitPplRecord("qi_it_join_left"); + assertEquals(Set.of("qi_it_join_left", "qi_it_join_right"), indicesOf(parent)); + + // Children are kept out of the Top N overview, so read them from the exported store. + List children = awaitChildren(parent.getString("parent_marker"), 2); + Set childIndices = new HashSet<>(); + for (JSONObject child : children) { + assertTrue(child.getBoolean("is_child")); + assertEquals("DSL", child.getString("query_source")); + // Stamped with the query's start, the record never sorts after the searches it spawned; a + // completion-time stamp could put it in a later Top N window than its children. + assertTrue( + "parent " + parent.getLong("timestamp") + " after child " + child.getLong("timestamp"), + parent.getLong("timestamp") <= child.getLong("timestamp")); + childIndices.addAll(indicesOf(child)); + } + // The searches read exactly the indices the record names. + assertEquals(indicesOf(parent), childIndices); + } + + // --------------------------------------------------------------------------------------------- + + /** + * Fails fast, naming the cause, when the installed Query Insights predates the PPL consumer: + * without it the gate setting is unregistered and no record would ever arrive. + */ + private static void assertQueryInsightsSupportsPpl() throws IOException { + Request request = new Request("GET", "/_cluster/settings"); + request.addParameter("include_defaults", "true"); + request.addParameter("flat_settings", "true"); + JSONObject settings = new JSONObject(getResponseBody(client().performRequest(request))); + // Listed under defaults until a test sets it, then under persistent. + boolean registered = false; + for (String section : List.of("defaults", "persistent", "transient")) { + JSONObject values = settings.optJSONObject(section); + registered |= values != null && values.has(PPL_ENABLED); + } + assertTrue( + "Query Insights in this cluster does not register " + + PPL_ENABLED + + ", so it predates the PPL consumer. Run with -PqueryInsightsZip=.", + registered); + } + + private static void indexWithDocs(String index) throws IOException { + Request create = new Request("PUT", "/" + index); + create.setJsonEntity( + "{\"settings\": {\"number_of_shards\": 1, \"number_of_replicas\": 0}," + + " \"mappings\": {\"properties\": {\"k\": {\"type\": \"keyword\"}," + + " \"v\": {\"type\": \"integer\"}}}}"); + client().performRequest(create); + Request bulk = new Request("POST", "/" + index + "/_bulk"); + bulk.addParameter("refresh", "true"); + bulk.setJsonEntity( + "{\"index\": {}}\n{\"k\": \"a\", \"v\": 1}\n{\"index\": {}}\n{\"k\": \"b\", \"v\": 2}\n"); + client().performRequest(bulk); + } + + /** Runs a recorded query and waits for its record, so everything sent before it has arrived. */ + private void awaitBarrier() throws Exception { + String index = "qi_it_barrier_" + BARRIER_SEQ.incrementAndGet(); + indexWithDocs(index); + executeQuery("source=" + index + " | fields v"); + awaitPplRecord(index); + } + + private static JSONObject awaitPplRecord(String index) throws Exception { + return awaitRecord(r -> isPpl(r) && indicesOf(r).contains(index), "PPL record for " + index); + } + + private static void assertNoPplRecord(String index) throws IOException { + List found = topQueries(r -> isPpl(r) && indicesOf(r).contains(index)); + assertTrue("unexpected PPL record for " + index + ": " + found, found.isEmpty()); + } + + private static JSONObject awaitRecord(Predicate match, String what) throws Exception { + AtomicReference found = new AtomicReference<>(); + assertBusy( + () -> { + List records = topQueries(match); + assertFalse("no " + what + " yet", records.isEmpty()); + found.set(records.get(0)); + }, + 60, + TimeUnit.SECONDS); + return found.get(); + } + + /** + * Top N overview over this class's run. The time range makes the API merge the in-memory window + * with the exported store, so a record is found even after its window has rolled over. + */ + private static List topQueries(Predicate match) throws IOException { + Request request = new Request("GET", "/_insights/top_queries"); + request.addParameter("type", "latency"); + request.addParameter("from", Instant.ofEpochMilli(CLASS_START_MILLIS - 60_000).toString()); + request.addParameter("to", Instant.now().plusSeconds(60).toString()); + JSONArray records = + new JSONObject(getResponseBody(client().performRequest(request))) + .optJSONArray("top_queries"); + return select(records, match); + } + + private static List awaitChildren(String parentMarker, int atLeast) throws Exception { + AtomicReference> found = new AtomicReference<>(); + // The store is written when the window rolls over, so allow a few windows. + assertBusy( + () -> { + Request search = new Request("POST", "/top_queries-*/_search"); + search.setJsonEntity( + "{\"size\": 50, \"query\": {\"term\": {\"derived_from\": \"" + + parentMarker + + "\"}}}"); + JSONArray hits = + new JSONObject(getResponseBody(client().performRequest(search))) + .getJSONObject("hits") + .getJSONArray("hits"); + List children = new ArrayList<>(); + for (int i = 0; i < hits.length(); i++) { + children.add(hits.getJSONObject(i).getJSONObject("_source")); + } + assertTrue( + "children of " + parentMarker + " so far: " + children.size(), + children.size() >= atLeast); + found.set(children); + }, + 3, + TimeUnit.MINUTES); + return found.get(); + } + + private static List select(JSONArray records, Predicate match) { + List selected = new ArrayList<>(); + if (records != null) { + for (int i = 0; i < records.length(); i++) { + JSONObject record = records.getJSONObject(i); + if (match.test(record)) { + selected.add(record); + } + } + } + return selected; + } + + private static boolean isPpl(JSONObject record) { + return "PPL".equals(record.optString("query_source")); + } + + private static Set indicesOf(JSONObject record) { + Set indices = new HashSet<>(); + JSONArray array = record.optJSONArray("indices"); + if (array != null) { + for (int i = 0; i < array.length(); i++) { + indices.add(array.getString(i)); + } + } + return indices; + } + + private static long measurement(JSONObject record, String metric) { + JSONObject value = record.getJSONObject("measurements").optJSONObject(metric); + return value == null ? 0L : value.getLong("number"); + } +} From 01794d82da1be3ba411186da6d9556b8f4a4dea0 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Fri, 16 Oct 2026 01:22:31 -0700 Subject: [PATCH 21/26] test(ppl): Assert Query Insights records on a secured cluster Add a cluster with Query Insights and the security plugin, and an IT that checks the identity a record carries and who can read it back: - a user with PPL access but no Query Insights permission at all, including the cluster:admin report action, still gets a record, and it carries their username and roles - that user gets 403 on the Top N API - with filter_by_mode=username, two users each see their own queries and not the other's Removing the identity capture fails the first and third tests. Move the record-reading helpers the two Query Insights ITs share into QueryInsightsRecords, and the PPL index permissions into a constant on SecurityTestBase so new roles grant what the existing ones do. Runs via :integ-test:integTestWithQueryInsightsAndSecurity, outside check for the same reason as integTestWithQueryInsights. Signed-off-by: Kishore Kumaar Natarajan --- integ-test/build.gradle | 43 +++- .../sql/ppl/QueryInsightsRecordingIT.java | 171 ++------------ .../QueryInsightsRecordingSecurityIT.java | 153 ++++++++++++ .../sql/security/SecurityTestBase.java | 23 +- .../sql/util/QueryInsightsRecords.java | 218 ++++++++++++++++++ 5 files changed, 451 insertions(+), 157 deletions(-) create mode 100644 integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsRecordingSecurityIT.java create mode 100644 integ-test/src/test/java/org/opensearch/sql/util/QueryInsightsRecords.java diff --git a/integ-test/build.gradle b/integ-test/build.gradle index 56756bea831..54ebb15c7fc 100644 --- a/integ-test/build.gradle +++ b/integ-test/build.gradle @@ -567,6 +567,15 @@ testClusters { plugin ":opensearch-sql-plugin" setting "plugins.query.datasources.encryption.masterkey", "1234567812345678" } + // The same, plus the security plugin (added by configureSecurityPlugin below), so records carry + // the requesting user and Top N reads go through Query Insights' RBAC filter. + integTestWithQueryInsightsAndSecurity { + testDistribution = 'archive' + plugin(getJobSchedulerPlugin()) + plugin(getQueryInsightsPlugin()) + plugin ":opensearch-sql-plugin" + setting "plugins.query.datasources.encryption.masterkey", "1234567812345678" + } // Smoke test: verify sql loads cleanly alongside analytics-engine. analyticsEngineCompat { testDistribution = 'archive' @@ -845,6 +854,35 @@ task integTestWithQueryInsights(type: RestIntegTestTask) { } } +// The reporting path on a secured cluster: whose identity a record carries and who can read it +// back. Not part of `check` yet, for the same reason as integTestWithQueryInsights: +// ./gradlew :integ-test:integTestWithQueryInsightsAndSecurity -PqueryInsightsZip=/path/to.zip +task integTestWithQueryInsightsAndSecurity(type: RestIntegTestTask) { + useCluster testClusters.integTestWithQueryInsightsAndSecurity + + dependsOn ':opensearch-sql-plugin:bundlePlugin' + + // This task installs the security plugin, so its absence is a breakage rather than a skip. + systemProperty 'tests.required.plugins', 'opensearch-security' + systemProperty 'tests.security.manager', 'false' + systemProperty 'project.root', project.projectDir.absolutePath + systemProperty 'defaultQuerySizeLimit', '10000' + systemProperty 'defaultQueryBucketSize', '1000' + systemProperty 'defaultMaxResultWindow', '10000' + doFirst { + systemProperty 'cluster.debug', getDebug() + systemProperty "https", "false" + systemProperty "user", "admin" + systemProperty "password", "admin" + } + + filter { + includeTestsMatching 'org.opensearch.sql.security.QueryInsightsRecordingSecurityIT' + } +} + +configureSecurityPlugin(testClusters.integTestWithQueryInsightsAndSecurity) + task integTestWithSecurity(type: RestIntegTestTask) { useCluster testClusters.integTestWithSecurity useCluster testClusters.remoteIntegTestWithSecurity @@ -874,6 +912,8 @@ task integTestWithSecurity(type: RestIntegTestTask) { // Exclude tests that require the analytics engine plugin stack (run separately via analyticsEngineSecurityIT) filter { excludeTestsMatching 'org.opensearch.sql.security.AnalyticsEngineSecurityIT' + // Needs Query Insights installed; runs via integTestWithQueryInsightsAndSecurity. + excludeTestsMatching 'org.opensearch.sql.security.QueryInsightsRecordingSecurityIT' } systemProperty 'tests.security.manager', 'false' @@ -1353,8 +1393,9 @@ task integTestRemote(type: RestIntegTestTask) { excludeTestsMatching "org.opensearch.sql.bwc.*IT" } - // Needs Query Insights installed; runs via :integ-test:integTestWithQueryInsights. + // Need Query Insights installed; run via the integTestWithQueryInsights* tasks. excludeTestsMatching 'org.opensearch.sql.ppl.QueryInsightsRecordingIT' + excludeTestsMatching 'org.opensearch.sql.security.QueryInsightsRecordingSecurityIT' if (analyticsEnabled) { // Full-text search (search-filter syntax -> Lucene query_string) needs an inverted-index diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsRecordingIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsRecordingIT.java index aa0cab7e097..826a77dbc6d 100644 --- a/integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsRecordingIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsRecordingIT.java @@ -5,24 +5,26 @@ package org.opensearch.sql.ppl; -import static org.opensearch.sql.legacy.TestUtils.getResponseBody; import static org.opensearch.sql.plugin.rest.RestPPLQueryAction.QUERY_API_ENDPOINT; +import static org.opensearch.sql.util.QueryInsightsRecords.PPL_ENABLED; +import static org.opensearch.sql.util.QueryInsightsRecords.awaitChildren; +import static org.opensearch.sql.util.QueryInsightsRecords.indicesOf; +import static org.opensearch.sql.util.QueryInsightsRecords.isPpl; +import static org.opensearch.sql.util.QueryInsightsRecords.measurement; +import static org.opensearch.sql.util.QueryInsightsRecords.pplRecordFor; -import java.io.IOException; -import java.time.Instant; -import java.util.ArrayList; import java.util.HashSet; import java.util.List; +import java.util.Map; import java.util.Set; -import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicReference; import java.util.function.Predicate; -import org.json.JSONArray; import org.json.JSONObject; import org.junit.Test; import org.opensearch.client.Request; +import org.opensearch.client.RequestOptions; import org.opensearch.client.ResponseException; +import org.opensearch.sql.util.QueryInsightsRecords; /** * End to end: PPL queries are recorded into Query Insights Top N on a cluster that has Query @@ -38,8 +40,6 @@ */ public class QueryInsightsRecordingIT extends PPLIntegTestCase { - private static final String PPL_ENABLED = "search.insights.top_queries.ppl.enabled"; - /** Records are stamped with their query's start, so everything this class records is later. */ private static final long CLASS_START_MILLIS = System.currentTimeMillis(); @@ -49,16 +49,8 @@ public class QueryInsightsRecordingIT extends PPLIntegTestCase { public void init() throws Exception { super.init(); enableCalcite(); // join runs on the Calcite path - assertQueryInsightsSupportsPpl(); - updateClusterSettings(new ClusterSetting(PERSISTENT, PPL_ENABLED, "true")); - updateClusterSettings( - new ClusterSetting(PERSISTENT, "search.insights.top_queries.latency.enabled", "true")); - // Room for every query this class runs, so none is evicted from the window it is read from. - updateClusterSettings( - new ClusterSetting(PERSISTENT, "search.insights.top_queries.latency.top_n_size", "100")); - // The smallest window, so records reach the exported store within a minute. - updateClusterSettings( - new ClusterSetting(PERSISTENT, "search.insights.top_queries.latency.window_size", "1m")); + QueryInsightsRecords.assertSupportsPpl(client()); + QueryInsightsRecords.enableRecording(client()); } @Test @@ -118,7 +110,7 @@ public void syntaxErrorIsNotRecorded() throws Exception { awaitBarrier(); // A query that never parses has no text or indices to report. Every other query in this class // names an index, so a PPL record without one could only be the syntax error. - for (JSONObject record : topQueries(QueryInsightsRecordingIT::isPpl)) { + for (JSONObject record : topQueries(QueryInsightsRecords::isPpl)) { assertFalse("unparsed query was recorded: " + record, indicesOf(record).isEmpty()); } } @@ -126,11 +118,11 @@ public void syntaxErrorIsNotRecorded() throws Exception { @Test public void nothingIsRecordedWhenTheGateIsOff() throws Exception { indexWithDocs("qi_it_gated"); - updateClusterSettings(new ClusterSetting(PERSISTENT, PPL_ENABLED, "false")); + QueryInsightsRecords.putPersistentSettings(client(), Map.of(PPL_ENABLED, "false")); try { executeQuery("source=qi_it_gated | where v > 1 | fields v"); } finally { - updateClusterSettings(new ClusterSetting(PERSISTENT, PPL_ENABLED, "true")); + QueryInsightsRecords.putPersistentSettings(client(), Map.of(PPL_ENABLED, "true")); } awaitBarrier(); @@ -167,8 +159,7 @@ public void childSearchesCarryTheParentMarker() throws Exception { JSONObject parent = awaitPplRecord("qi_it_join_left"); assertEquals(Set.of("qi_it_join_left", "qi_it_join_right"), indicesOf(parent)); - // Children are kept out of the Top N overview, so read them from the exported store. - List children = awaitChildren(parent.getString("parent_marker"), 2); + List children = awaitChildren(client(), parent.getString("parent_marker"), 2); Set childIndices = new HashSet<>(); for (JSONObject child : children) { assertTrue(child.getBoolean("is_child")); @@ -186,41 +177,8 @@ public void childSearchesCarryTheParentMarker() throws Exception { // --------------------------------------------------------------------------------------------- - /** - * Fails fast, naming the cause, when the installed Query Insights predates the PPL consumer: - * without it the gate setting is unregistered and no record would ever arrive. - */ - private static void assertQueryInsightsSupportsPpl() throws IOException { - Request request = new Request("GET", "/_cluster/settings"); - request.addParameter("include_defaults", "true"); - request.addParameter("flat_settings", "true"); - JSONObject settings = new JSONObject(getResponseBody(client().performRequest(request))); - // Listed under defaults until a test sets it, then under persistent. - boolean registered = false; - for (String section : List.of("defaults", "persistent", "transient")) { - JSONObject values = settings.optJSONObject(section); - registered |= values != null && values.has(PPL_ENABLED); - } - assertTrue( - "Query Insights in this cluster does not register " - + PPL_ENABLED - + ", so it predates the PPL consumer. Run with -PqueryInsightsZip=.", - registered); - } - - private static void indexWithDocs(String index) throws IOException { - Request create = new Request("PUT", "/" + index); - create.setJsonEntity( - "{\"settings\": {\"number_of_shards\": 1, \"number_of_replicas\": 0}," - + " \"mappings\": {\"properties\": {\"k\": {\"type\": \"keyword\"}," - + " \"v\": {\"type\": \"integer\"}}}}"); - client().performRequest(create); - Request bulk = new Request("POST", "/" + index + "/_bulk"); - bulk.addParameter("refresh", "true"); - bulk.setJsonEntity( - "{\"index\": {}}\n{\"k\": \"a\", \"v\": 1}\n{\"index\": {}}\n{\"k\": \"b\", \"v\": 2}\n"); - client().performRequest(bulk); + private static void indexWithDocs(String index) throws Exception { + QueryInsightsRecords.indexWithDocs(client(), index); } /** Runs a recorded query and waits for its record, so everything sent before it has arrived. */ @@ -232,100 +190,21 @@ private void awaitBarrier() throws Exception { } private static JSONObject awaitPplRecord(String index) throws Exception { - return awaitRecord(r -> isPpl(r) && indicesOf(r).contains(index), "PPL record for " + index); + return awaitRecord(pplRecordFor(index), "PPL record for " + index); } - private static void assertNoPplRecord(String index) throws IOException { - List found = topQueries(r -> isPpl(r) && indicesOf(r).contains(index)); + private static void assertNoPplRecord(String index) throws Exception { + List found = topQueries(pplRecordFor(index)); assertTrue("unexpected PPL record for " + index + ": " + found, found.isEmpty()); } private static JSONObject awaitRecord(Predicate match, String what) throws Exception { - AtomicReference found = new AtomicReference<>(); - assertBusy( - () -> { - List records = topQueries(match); - assertFalse("no " + what + " yet", records.isEmpty()); - found.set(records.get(0)); - }, - 60, - TimeUnit.SECONDS); - return found.get(); - } - - /** - * Top N overview over this class's run. The time range makes the API merge the in-memory window - * with the exported store, so a record is found even after its window has rolled over. - */ - private static List topQueries(Predicate match) throws IOException { - Request request = new Request("GET", "/_insights/top_queries"); - request.addParameter("type", "latency"); - request.addParameter("from", Instant.ofEpochMilli(CLASS_START_MILLIS - 60_000).toString()); - request.addParameter("to", Instant.now().plusSeconds(60).toString()); - JSONArray records = - new JSONObject(getResponseBody(client().performRequest(request))) - .optJSONArray("top_queries"); - return select(records, match); - } - - private static List awaitChildren(String parentMarker, int atLeast) throws Exception { - AtomicReference> found = new AtomicReference<>(); - // The store is written when the window rolls over, so allow a few windows. - assertBusy( - () -> { - Request search = new Request("POST", "/top_queries-*/_search"); - search.setJsonEntity( - "{\"size\": 50, \"query\": {\"term\": {\"derived_from\": \"" - + parentMarker - + "\"}}}"); - JSONArray hits = - new JSONObject(getResponseBody(client().performRequest(search))) - .getJSONObject("hits") - .getJSONArray("hits"); - List children = new ArrayList<>(); - for (int i = 0; i < hits.length(); i++) { - children.add(hits.getJSONObject(i).getJSONObject("_source")); - } - assertTrue( - "children of " + parentMarker + " so far: " + children.size(), - children.size() >= atLeast); - found.set(children); - }, - 3, - TimeUnit.MINUTES); - return found.get(); - } - - private static List select(JSONArray records, Predicate match) { - List selected = new ArrayList<>(); - if (records != null) { - for (int i = 0; i < records.length(); i++) { - JSONObject record = records.getJSONObject(i); - if (match.test(record)) { - selected.add(record); - } - } - } - return selected; - } - - private static boolean isPpl(JSONObject record) { - return "PPL".equals(record.optString("query_source")); - } - - private static Set indicesOf(JSONObject record) { - Set indices = new HashSet<>(); - JSONArray array = record.optJSONArray("indices"); - if (array != null) { - for (int i = 0; i < array.length(); i++) { - indices.add(array.getString(i)); - } - } - return indices; + return QueryInsightsRecords.awaitRecord( + client(), RequestOptions.DEFAULT, CLASS_START_MILLIS, match, what); } - private static long measurement(JSONObject record, String metric) { - JSONObject value = record.getJSONObject("measurements").optJSONObject(metric); - return value == null ? 0L : value.getLong("number"); + private static List topQueries(Predicate match) throws Exception { + return QueryInsightsRecords.topQueries( + client(), RequestOptions.DEFAULT, CLASS_START_MILLIS, match); } } diff --git a/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsRecordingSecurityIT.java b/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsRecordingSecurityIT.java new file mode 100644 index 00000000000..3b2b33ebcfd --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsRecordingSecurityIT.java @@ -0,0 +1,153 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.security; + +import static org.opensearch.sql.util.QueryInsightsRecords.awaitRecord; +import static org.opensearch.sql.util.QueryInsightsRecords.pplRecordFor; +import static org.opensearch.sql.util.QueryInsightsRecords.stringsOf; +import static org.opensearch.sql.util.QueryInsightsRecords.topQueries; + +import java.util.List; +import java.util.Map; +import org.json.JSONArray; +import org.json.JSONObject; +import org.junit.Test; +import org.opensearch.client.RequestOptions; +import org.opensearch.client.ResponseException; +import org.opensearch.sql.util.ClusterPlugins; +import org.opensearch.sql.util.QueryInsightsRecords; + +/** + * Query Insights recording on a security-enabled cluster: whose identity a PPL record carries, and + * who can read it back. + * + *

Runs only in {@code :integ-test:integTestWithQueryInsightsAndSecurity}, whose cluster installs + * both Query Insights and the security plugin. {@link QueryInsightsReportingSecurityIT} covers a + * secured cluster without Query Insights. + * + *

Each user queries an index named after them, so a record is identified by its index. + */ +public class QueryInsightsRecordingSecurityIT extends SecurityTestBase { + + /** + * Holds PPL and read access to its index but no Query Insights permission at all, including the + * {@code cluster:admin} report action the plugin sends when the query finishes. + */ + private static final String PPL_ONLY_USER = "qi_sec_carol"; + + /** Can also read Top N; with filter_by_mode=username each should see only their own queries. */ + private static final String ALICE = "qi_sec_alice"; + + private static final String BOB = "qi_sec_bob"; + + private static final String TOP_QUERIES_ACTION = "cluster:admin/opensearch/insights/top_queries"; + private static final String FILTER_BY_MODE = "search.insights.top_queries.filter_by_mode"; + + private static final long CLASS_START_MILLIS = System.currentTimeMillis(); + + private static boolean usersCreated = false; + + @Override + protected void init() throws Exception { + super.init(); + enableCalcite(); + ClusterPlugins.requirePluginOrAssume( + client(), ClusterPlugins.SECURITY_PLUGIN, "security plugin is not installed"); + QueryInsightsRecords.assertSupportsPpl(client()); + QueryInsightsRecords.enableRecording(client()); + if (!usersCreated) { + for (String user : List.of(PPL_ONLY_USER, ALICE, BOB)) { + QueryInsightsRecords.indexWithDocs(client(), user); + } + createRoleWithIndexAccess(roleOf(PPL_ONLY_USER), "qi_sec_*"); + createUser(PPL_ONLY_USER, roleOf(PPL_ONLY_USER)); + for (String user : List.of(ALICE, BOB)) { + createRoleWithPermissions( + roleOf(user), + "qi_sec_*", + new String[] {"cluster:admin/opensearch/ppl", TOP_QUERIES_ACTION}, + PPL_INDEX_PERMISSIONS); + createUser(user, roleOf(user)); + } + usersCreated = true; + } + } + + @Test + public void recordCarriesTheIdentityOfAUserWithoutAnyInsightsPermission() throws Exception { + executeQueryAsUser("source=" + PPL_ONLY_USER + " | where v > 1 | fields v", PPL_ONLY_USER); + + // Read as admin. The record exists although the user holds no Query Insights permission, so + // the report is not authorized against the user who ran the query. + JSONObject record = awaitRecordFor(PPL_ONLY_USER); + assertEquals(PPL_ONLY_USER, record.optString("username")); + JSONArray roles = record.optJSONArray("user_roles"); + assertTrue( + "record must carry the user's roles, was " + roles, + stringsOf(roles).contains(roleOf(PPL_ONLY_USER))); + } + + @Test + public void userWithoutInsightsPermissionCannotReadTopN() throws Exception { + ResponseException e = + expectThrows( + ResponseException.class, + () -> topQueries(client(), asUser(PPL_ONLY_USER), CLASS_START_MILLIS, r -> true)); + assertEquals(403, e.getResponse().getStatusLine().getStatusCode()); + } + + @Test + public void withFilterByUsernameEachUserSeesOnlyTheirOwnQueries() throws Exception { + QueryInsightsRecords.putPersistentSettings(client(), Map.of(FILTER_BY_MODE, "username")); + try { + executeQueryAsUser("source=" + ALICE + " | where v > 1 | fields v", ALICE); + executeQueryAsUser("source=" + BOB + " | where v > 1 | fields v", BOB); + // Admin sees everything, so both records must exist before checking what each user sees. + awaitRecordFor(ALICE); + awaitRecordFor(BOB); + + assertSeesOnlyOwnQueries(ALICE, BOB); + assertSeesOnlyOwnQueries(BOB, ALICE); + } finally { + QueryInsightsRecords.putPersistentSettings(client(), Map.of(FILTER_BY_MODE, "none")); + } + } + + // --------------------------------------------------------------------------------------------- + + private void assertSeesOnlyOwnQueries(String user, String other) throws Exception { + List visible = + topQueries(client(), asUser(user), CLASS_START_MILLIS, QueryInsightsRecords::isPpl); + assertTrue( + user + " sees none of their own queries", visible.stream().anyMatch(pplRecordFor(user))); + assertTrue( + user + " sees " + other + "'s queries: " + visible, + visible.stream().noneMatch(pplRecordFor(other))); + for (JSONObject record : visible) { + assertEquals( + user + " sees another user's record: " + record, user, record.optString("username")); + } + } + + private static JSONObject awaitRecordFor(String user) throws Exception { + return awaitRecord( + client(), + RequestOptions.DEFAULT, + CLASS_START_MILLIS, + pplRecordFor(user), + "PPL record for " + user); + } + + private RequestOptions asUser(String user) { + return RequestOptions.DEFAULT.toBuilder() + .addHeader("Authorization", createBasicAuthHeader(user, STRONG_PASSWORD)) + .build(); + } + + private static String roleOf(String user) { + return user + "_role"; + } +} diff --git a/integ-test/src/test/java/org/opensearch/sql/security/SecurityTestBase.java b/integ-test/src/test/java/org/opensearch/sql/security/SecurityTestBase.java index 750adcc4693..50ceccc3856 100644 --- a/integ-test/src/test/java/org/opensearch/sql/security/SecurityTestBase.java +++ b/integ-test/src/test/java/org/opensearch/sql/security/SecurityTestBase.java @@ -24,6 +24,18 @@ public abstract class SecurityTestBase extends PPLIntegTestCase { protected static final String STRONG_PASSWORD = "correcthorsebatterystaple"; + /** Index permissions a PPL query needs on the indices it reads. */ + protected static final String[] PPL_INDEX_PERMISSIONS = { + "indices:data/read/search*", + "indices:admin/mappings/get", + "indices:monitor/settings/get", + "indices:data/read/point_in_time/create", + "indices:data/read/point_in_time/delete", + // Index pruning probes with these; ppl_full_access grants both since 3.9. + "indices:admin/resolve/index", + "indices:data/read/field_caps*" + }; + /** * Creates a role with access to a specific index pattern and standard permissions. * @@ -36,16 +48,7 @@ protected void createRoleWithIndexAccess(String roleName, String indexPattern) roleName, indexPattern, new String[] {"cluster:admin/opensearch/ppl"}, - new String[] { - "indices:data/read/search*", - "indices:admin/mappings/get", - "indices:monitor/settings/get", - "indices:data/read/point_in_time/create", - "indices:data/read/point_in_time/delete", - // Index pruning probes with these; ppl_full_access grants both since 3.9. - "indices:admin/resolve/index", - "indices:data/read/field_caps*" - }); + PPL_INDEX_PERMISSIONS); } /** diff --git a/integ-test/src/test/java/org/opensearch/sql/util/QueryInsightsRecords.java b/integ-test/src/test/java/org/opensearch/sql/util/QueryInsightsRecords.java new file mode 100644 index 00000000000..a8149b80e1d --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/util/QueryInsightsRecords.java @@ -0,0 +1,218 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.util; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.opensearch.sql.legacy.TestUtils.getResponseBody; + +import java.io.IOException; +import java.time.Instant; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Predicate; +import org.json.JSONArray; +import org.json.JSONObject; +import org.opensearch.client.Request; +import org.opensearch.client.RequestOptions; +import org.opensearch.client.RestClient; +import org.opensearch.test.OpenSearchTestCase; + +/** + * Reads Query Insights Top N records for the ITs that run with Query Insights installed. + * + *

Records are matched by the index they report rather than by query text, which is anonymized, + * so each scenario should read an index of its own. + */ +public final class QueryInsightsRecords { + + public static final String PPL_ENABLED = "search.insights.top_queries.ppl.enabled"; + + private QueryInsightsRecords() {} + + /** + * Fails, naming the cause, when the installed Query Insights predates the PPL consumer: without + * it the gate setting is unregistered and no record would ever arrive. + */ + public static void assertSupportsPpl(RestClient client) throws IOException { + Request request = new Request("GET", "/_cluster/settings"); + request.addParameter("include_defaults", "true"); + request.addParameter("flat_settings", "true"); + JSONObject settings = new JSONObject(getResponseBody(client.performRequest(request))); + // Listed under defaults until a test sets it, then under persistent. + boolean registered = false; + for (String section : List.of("defaults", "persistent", "transient")) { + JSONObject values = settings.optJSONObject(section); + registered |= values != null && values.has(PPL_ENABLED); + } + assertTrue( + "Query Insights in this cluster does not register " + + PPL_ENABLED + + ", so it predates the PPL consumer. Run with -PqueryInsightsZip=.", + registered); + } + + /** Records PPL queries, with room for every query a suite runs and the smallest window. */ + public static void enableRecording(RestClient client) throws IOException { + putPersistentSettings( + client, + Map.of( + PPL_ENABLED, + "true", + "search.insights.top_queries.latency.enabled", + "true", + // Room for every query a suite runs, so none is evicted from the window it is read + // from. + "search.insights.top_queries.latency.top_n_size", + "100", + // The smallest window, so records reach the exported store within a minute. + "search.insights.top_queries.latency.window_size", + "1m")); + } + + public static void putPersistentSettings(RestClient client, Map settings) + throws IOException { + Request request = new Request("PUT", "/_cluster/settings"); + request.setJsonEntity(new JSONObject().put("persistent", new JSONObject(settings)).toString()); + client.performRequest(request); + } + + /** Creates {@code index} with two documents over a keyword {@code k} and an integer {@code v}. */ + public static void indexWithDocs(RestClient client, String index) throws IOException { + Request create = new Request("PUT", "/" + index); + create.setJsonEntity( + "{\"settings\": {\"number_of_shards\": 1, \"number_of_replicas\": 0}," + + " \"mappings\": {\"properties\": {\"k\": {\"type\": \"keyword\"}," + + " \"v\": {\"type\": \"integer\"}}}}"); + client.performRequest(create); + Request bulk = new Request("POST", "/" + index + "/_bulk"); + bulk.addParameter("refresh", "true"); + bulk.setJsonEntity( + "{\"index\": {}}\n{\"k\": \"a\", \"v\": 1}\n{\"index\": {}}\n{\"k\": \"b\", \"v\": 2}\n"); + client.performRequest(bulk); + } + + /** + * Top N overview records since {@code sinceMillis}, read with {@code options} (pass an + * Authorization header to read as a particular user). The time range makes the API merge the + * in-memory window with the exported store, so a record is found even after its window has rolled + * over. + */ + public static List topQueries( + RestClient client, RequestOptions options, long sinceMillis, Predicate match) + throws IOException { + Request request = new Request("GET", "/_insights/top_queries"); + request.addParameter("type", "latency"); + request.addParameter("from", Instant.ofEpochMilli(sinceMillis - 60_000).toString()); + request.addParameter("to", Instant.now().plusSeconds(60).toString()); + request.setOptions(options); + JSONArray records = + new JSONObject(getResponseBody(client.performRequest(request))).optJSONArray("top_queries"); + return select(records, match); + } + + public static JSONObject awaitRecord( + RestClient client, + RequestOptions options, + long sinceMillis, + Predicate match, + String what) + throws Exception { + AtomicReference found = new AtomicReference<>(); + OpenSearchTestCase.assertBusy( + () -> { + List records = topQueries(client, options, sinceMillis, match); + assertFalse("no " + what + " yet", records.isEmpty()); + found.set(records.get(0)); + }, + 60, + TimeUnit.SECONDS); + return found.get(); + } + + /** + * The child records filed under {@code parentMarker}. Children are kept out of the Top N + * overview, so they are read from the exported store, which is written when the window rolls + * over. + */ + public static List awaitChildren(RestClient client, String parentMarker, int atLeast) + throws Exception { + AtomicReference> found = new AtomicReference<>(); + OpenSearchTestCase.assertBusy( + () -> { + Request search = new Request("POST", "/top_queries-*/_search"); + search.setJsonEntity( + new JSONObject() + .put("size", 50) + .put( + "query", + new JSONObject() + .put("term", new JSONObject().put("derived_from", parentMarker))) + .toString()); + JSONArray hits = + new JSONObject(getResponseBody(client.performRequest(search))) + .getJSONObject("hits") + .getJSONArray("hits"); + List children = new ArrayList<>(); + for (int i = 0; i < hits.length(); i++) { + children.add(hits.getJSONObject(i).getJSONObject("_source")); + } + assertTrue( + "children of " + parentMarker + " so far: " + children.size(), + children.size() >= atLeast); + found.set(children); + }, + 3, + TimeUnit.MINUTES); + return found.get(); + } + + public static Predicate pplRecordFor(String index) { + return record -> isPpl(record) && indicesOf(record).contains(index); + } + + public static boolean isPpl(JSONObject record) { + return "PPL".equals(record.optString("query_source")); + } + + public static Set indicesOf(JSONObject record) { + return stringsOf(record.optJSONArray("indices")); + } + + public static Set stringsOf(JSONArray array) { + Set values = new HashSet<>(); + if (array != null) { + for (int i = 0; i < array.length(); i++) { + values.add(array.getString(i)); + } + } + return values; + } + + public static long measurement(JSONObject record, String metric) { + JSONObject value = record.getJSONObject("measurements").optJSONObject(metric); + return value == null ? 0L : value.getLong("number"); + } + + public static List select(JSONArray records, Predicate match) { + List selected = new ArrayList<>(); + if (records != null) { + for (int i = 0; i < records.length(); i++) { + JSONObject record = records.getJSONObject(i); + if (match.test(record)) { + selected.add(record); + } + } + } + return selected; + } +} From 50d97714af02a134cf74be4efb9fa33c7b59a875 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 27 Oct 2026 11:42:34 -0700 Subject: [PATCH 22/26] refactor(ppl): Drop indices from the Top N record The record carried the indices a PPL query reads, resolved by walking its AST through a getSources() hook added to Join, Lookup, Append, AppendCol, AppendPipe and GraphLookup. Nothing needs them: the Top N use cases are slow and failed queries. The searches a query runs are already recorded as its children, each with the indices it reads, so Query Insights can show those under the query instead. Remove the extractor, the getSources() hook, and the indices from the record and its wire format. Neither plugin has released the format, so it stays at version 1; the Query Insights handler has to stop reading the field at the same time. Signed-off-by: Kishore Kumaar Natarajan --- .../org/opensearch/sql/ast/tree/Append.java | 6 - .../opensearch/sql/ast/tree/AppendCol.java | 6 - .../opensearch/sql/ast/tree/AppendPipe.java | 6 - .../opensearch/sql/ast/tree/GraphLookup.java | 6 - .../org/opensearch/sql/ast/tree/Join.java | 9 - .../org/opensearch/sql/ast/tree/Lookup.java | 6 - .../sql/ast/tree/UnresolvedPlan.java | 17 - .../ast/tree/UnresolvedPlanSourcesTest.java | 68 ---- docs/user/ppl/admin/settings.md | 2 +- .../sql/plugin/transport/PPLQueryTask.java | 12 - .../transport/QueryInsightsReporter.java | 9 - .../transport/TransportPPLQueryAction.java | 12 +- .../transport/QueryInsightsReporterTest.java | 47 +-- .../opensearch/sql/ppl/PPLQueryRunner.java | 4 +- .../org/opensearch/sql/ppl/PPLService.java | 37 +- .../sql/ppl/QueryInsightsMetadata.java | 10 +- .../sql/ppl/utils/PPLQueryIndexExtractor.java | 239 ----------- .../opensearch/sql/ppl/PPLServiceTest.java | 43 +- .../ppl/utils/PPLQueryIndexExtractorTest.java | 375 ------------------ 19 files changed, 20 insertions(+), 894 deletions(-) delete mode 100644 core/src/test/java/org/opensearch/sql/ast/tree/UnresolvedPlanSourcesTest.java delete mode 100644 ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java delete mode 100644 ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Append.java b/core/src/main/java/org/opensearch/sql/ast/tree/Append.java index edb8b907455..e5d4f54228b 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/Append.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Append.java @@ -33,12 +33,6 @@ public UnresolvedPlan attach(UnresolvedPlan child) { return this; } - /** {@inheritDoc} Adds the appended sub-search, which is not a child. */ - @Override - public List getSources() { - return this.subSearch == null ? ImmutableList.of() : ImmutableList.of(this.subSearch); - } - @Override public List getChild() { return this.child == null ? ImmutableList.of() : ImmutableList.of(this.child); diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/AppendCol.java b/core/src/main/java/org/opensearch/sql/ast/tree/AppendCol.java index ef011ad2627..39bfe255749 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/AppendCol.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/AppendCol.java @@ -37,12 +37,6 @@ public UnresolvedPlan attach(UnresolvedPlan child) { return this; } - /** {@inheritDoc} Adds the appended sub-search, which is not a child. */ - @Override - public List getSources() { - return this.subSearch == null ? ImmutableList.of() : ImmutableList.of(this.subSearch); - } - @Override public List getChild() { return this.child == null ? ImmutableList.of() : ImmutableList.of(this.child); diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/AppendPipe.java b/core/src/main/java/org/opensearch/sql/ast/tree/AppendPipe.java index 6ad7fb4eefd..0ea1cb9b453 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/AppendPipe.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/AppendPipe.java @@ -33,12 +33,6 @@ public AppendPipe attach(UnresolvedPlan child) { return this; } - /** {@inheritDoc} Adds the appended sub-pipeline, which is not a child. */ - @Override - public List getSources() { - return this.subQuery == null ? ImmutableList.of() : ImmutableList.of(this.subQuery); - } - @Override public List getChild() { return this.child == null ? ImmutableList.of() : ImmutableList.of(this.child); diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/GraphLookup.java b/core/src/main/java/org/opensearch/sql/ast/tree/GraphLookup.java index d38ec09fc98..f1457a8a540 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/GraphLookup.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/GraphLookup.java @@ -95,12 +95,6 @@ public UnresolvedPlan attach(UnresolvedPlan child) { return this; } - /** {@inheritDoc} Adds the graph-lookup source table, which is not a child. */ - @Override - public List getSources() { - return this.fromTable == null ? ImmutableList.of() : ImmutableList.of(this.fromTable); - } - @Override public List getChild() { return child == null ? ImmutableList.of() : ImmutableList.of(child); diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Join.java b/core/src/main/java/org/opensearch/sql/ast/tree/Join.java index 3c2e9497521..a0f58d35124 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/Join.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Join.java @@ -73,15 +73,6 @@ public UnresolvedPlan attach(UnresolvedPlan child) { return this; } - /** - * {@inheritDoc} Adds the right branch ({@code getChild()} exposes only the left), which is not a - * child. - */ - @Override - public List getSources() { - return this.right == null ? ImmutableList.of() : ImmutableList.of(this.right); - } - @Override public List getChild() { return this.left == null ? ImmutableList.of() : ImmutableList.of(this.left); diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Lookup.java b/core/src/main/java/org/opensearch/sql/ast/tree/Lookup.java index dc2144232dd..5b0c5d585d8 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/Lookup.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Lookup.java @@ -45,12 +45,6 @@ public UnresolvedPlan attach(UnresolvedPlan child) { return this; } - /** {@inheritDoc} Adds the lookup table, which is not a child. */ - @Override - public List getSources() { - return this.lookupRelation == null ? ImmutableList.of() : ImmutableList.of(this.lookupRelation); - } - @Override public List getChild() { return this.child == null ? ImmutableList.of() : ImmutableList.of(this.child); diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/UnresolvedPlan.java b/core/src/main/java/org/opensearch/sql/ast/tree/UnresolvedPlan.java index d4f616ff3ed..3074303b0d7 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/UnresolvedPlan.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/UnresolvedPlan.java @@ -5,8 +5,6 @@ package org.opensearch.sql.ast.tree; -import com.google.common.collect.ImmutableList; -import java.util.List; import lombok.EqualsAndHashCode; import lombok.ToString; import org.opensearch.sql.ast.AbstractNodeVisitor; @@ -21,20 +19,5 @@ public T accept(AbstractNodeVisitor nodeVisitor, C context) { return nodeVisitor.visitChildren(this, context); } - /** - * Plans this node reads from that {@link #getChild()} does not expose, such as a sub-search or a - * lookup table held in a separate field. - * - *

Override this when adding a command that holds a plan outside its child list. Consumers that - * need every source a query reads from (for example index-name extraction for Query Insights) - * walk {@code getChild()} plus this list, so a plan left out of both is silently invisible to - * them. The returned plans must not duplicate {@code getChild()}. - * - * @return the additional source plans; empty by default - */ - public List getSources() { - return ImmutableList.of(); - } - public abstract UnresolvedPlan attach(UnresolvedPlan child); } diff --git a/core/src/test/java/org/opensearch/sql/ast/tree/UnresolvedPlanSourcesTest.java b/core/src/test/java/org/opensearch/sql/ast/tree/UnresolvedPlanSourcesTest.java deleted file mode 100644 index fddb6cd298e..00000000000 --- a/core/src/test/java/org/opensearch/sql/ast/tree/UnresolvedPlanSourcesTest.java +++ /dev/null @@ -1,68 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.ast.tree; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import java.util.List; -import org.junit.jupiter.api.Test; -import org.opensearch.sql.ast.expression.QualifiedName; - -/** - * Pins the {@link UnresolvedPlan#getSources()} contract: a node declares the plans it reads from - * that {@code getChild()} does not expose, so consumers can find every source generically. A node - * that drops its declaration here silently hides a source from them. - */ -class UnresolvedPlanSourcesTest { - - private static UnresolvedPlan relation(String name) { - return new Relation(new QualifiedName(name)); - } - - @Test - void defaultsToEmpty() { - // A node whose sources are all reachable via getChild() declares none. - assertTrue(relation("t").getSources().isEmpty()); - } - - @Test - void appendDeclaresSubSearch() { - UnresolvedPlan sub = relation("appended"); - assertEquals(List.of(sub), new Append(sub).attach(relation("base")).getSources()); - } - - @Test - void appendColDeclaresSubSearch() { - UnresolvedPlan sub = relation("appended"); - assertEquals(List.of(sub), new AppendCol(false, sub).attach(relation("base")).getSources()); - } - - @Test - void appendPipeDeclaresSubQuery() { - UnresolvedPlan sub = relation("appended"); - assertEquals(List.of(sub), new AppendPipe(sub).attach(relation("base")).getSources()); - } - - @Test - void sourcesAreDisjointFromChild() { - // The contract requires getSources() not to repeat getChild(), or consumers walking both would - // visit the piped input twice. - UnresolvedPlan child = relation("base"); - UnresolvedPlan sub = relation("appended"); - Append append = (Append) new Append(sub).attach(child); - assertEquals(List.of(child), append.getChild()); - assertEquals(List.of(sub), append.getSources()); - } - - @Test - void nullSourceYieldsEmptyRatherThanNpe() { - // Nodes can be inspected before attach()/with an absent sub-search; the walk must not break. - assertTrue(new AppendPipe(null).getSources().isEmpty()); - assertTrue(new Append(null).getSources().isEmpty()); - assertTrue(new AppendCol(false, null).getSources().isEmpty()); - } -} diff --git a/docs/user/ppl/admin/settings.md b/docs/user/ppl/admin/settings.md index c7551c64bae..2eed41f5fe9 100644 --- a/docs/user/ppl/admin/settings.md +++ b/docs/user/ppl/admin/settings.md @@ -482,7 +482,7 @@ Expected output: ### Description -This setting controls whether completed PPL queries are recorded into Query Insights Top N. When enabled, the PPL plugin reports each finished PPL query (with its resolved indices, latency, CPU, memory, and — on a security-enabled cluster — the requesting user) to the Query Insights plugin, and tags the child DSL searches a PPL query spawns so they appear as sub-queries. +This setting controls whether completed PPL queries are recorded into Query Insights Top N. When enabled, the PPL plugin reports each finished PPL query (with its latency, CPU, memory, and — on a security-enabled cluster — the requesting user) to the Query Insights plugin, and tags the child DSL searches a PPL query spawns so they appear as sub-queries. 1. The default value is false (opt-in). 2. This setting is node scope, matching the rest of the `search.insights.top_queries.*` family. diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java index 7644cd574af..668441c8def 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java @@ -7,7 +7,6 @@ import com.sun.management.ThreadMXBean; import java.lang.management.ManagementFactory; -import java.util.List; import java.util.Map; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; @@ -27,9 +26,6 @@ public class PPLQueryTask extends CancellableTask implements ThreadResourceAccou /** Security {@code _opendistro_security_user_info} string; null when unsecured. */ private volatile String queryInsightsUserInfo; - /** Source index name(s) resolved from the AST; empty when unresolved. */ - private volatile List queryInsightsIndices = List.of(); - /** Anonymized query text ({@code PPLQueryDataAnonymizer}); null when anonymization didn't run. */ private volatile String queryInsightsAnonymizedQuery; @@ -79,14 +75,6 @@ public String getQueryInsightsUserInfo() { return queryInsightsUserInfo; } - public void setQueryInsightsIndices(List indices) { - this.queryInsightsIndices = indices == null ? List.of() : indices; - } - - public List getQueryInsightsIndices() { - return queryInsightsIndices; - } - public void setQueryInsightsAnonymizedQuery(String anonymizedQuery) { this.queryInsightsAnonymizedQuery = anonymizedQuery; } diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java index 7adaf64109f..df0362cb131 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java @@ -5,7 +5,6 @@ package org.opensearch.sql.plugin.transport; -import java.util.List; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.Version; @@ -52,7 +51,6 @@ public static void report( long latencyMillis, long cpuNanos, long memoryBytes, - List indices, String userInfo, boolean failed) { try { @@ -66,13 +64,6 @@ public static void report( out.writeVLong(Math.max(0L, latencyMillis)); out.writeVLong(Math.max(0L, cpuNanos)); out.writeVLong(Math.max(0L, memoryBytes)); - - final List safeIndices = indices == null ? List.of() : indices; - out.writeVInt(safeIndices.size()); - for (final String index : safeIndices) { - out.writeString(index == null ? "" : index); - } - out.writeString(userInfo == null ? "" : userInfo); out.writeBoolean(failed); 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 5a0015baf52..e576a76ebff 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 @@ -11,7 +11,6 @@ import static org.opensearch.sql.protocol.response.format.JsonResponseFormatter.Style.PRETTY; import java.time.Clock; -import java.util.List; import java.util.Locale; import java.util.Optional; import java.util.concurrent.CompletionStage; @@ -294,7 +293,7 @@ protected void innerOnFailure(Exception e) {} * Whether a finished task should produce a Top N record. Decided from what the metadata sink * captured rather than from the route, because the route cannot see either case: {@code explain} * is valid inside the query text (not just on {@code /_explain}), and a query that fails to parse - * never reaches the sink at all, so it has no query text or indices to report. + * never reaches the sink at all, so there is no parsed query to report. * *

A nested query is also skipped: its child searches carry the outer query's marker, so a * record built from this task's own id would advertise a marker no child references, and the @@ -332,7 +331,6 @@ private void writeQueryInsightsRecord(PPLQueryTask pplQueryTask) { // when stamping failed, which the reporter sends as an empty marker. String parentMarker = pplQueryTask.getQueryInsightsParentMarker(); - List indices = pplQueryTask.getQueryInsightsIndices(); String userInfo = pplQueryTask.getQueryInsightsUserInfo(); boolean failed = pplQueryTask.isQueryInsightsFailed(); @@ -347,7 +345,6 @@ private void writeQueryInsightsRecord(PPLQueryTask pplQueryTask) { latencyMillis, cpuNanos, memoryBytes, - indices, userInfo, failed); } catch (Exception e) { @@ -453,7 +450,6 @@ protected void doExecute( rootSpan.addAttribute("db.query.text", metadata.anonymizedQuery()); if (pplQueryTask != null) { pplQueryTask.setQueryInsightsAnonymizedQuery(metadata.anonymizedQuery()); - pplQueryTask.setQueryInsightsIndices(metadata.indices()); pplQueryTask.setQueryInsightsExplain(metadata.explain()); } }; @@ -520,11 +516,7 @@ private void submitAsync( Runnable cancelExecution = () -> pplQueryTask.cancel("async PPL query cancelled"); QueryRunner runner = new PPLQueryRunner( - pplService, - transformedRequest, - queryInsightsSink, - Clock.systemUTC(), - cancelExecution); + pplService, transformedRequest, queryInsightsSink, Clock.systemUTC(), cancelExecution); ThreadContext threadContext = clientRef.threadPool().getThreadContext(); ActionListener ctxListener = ContextPreservingActionListener.wrapPreservingContext(listener, threadContext); diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsReporterTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsReporterTest.java index 28df747d915..be31c485215 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsReporterTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsReporterTest.java @@ -5,7 +5,6 @@ package org.opensearch.sql.plugin.transport; -import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNull; @@ -18,9 +17,7 @@ import static org.mockito.Mockito.when; import java.io.IOException; -import java.util.Arrays; import java.util.Base64; -import java.util.List; import java.util.concurrent.atomic.AtomicReference; import org.junit.Before; import org.junit.Test; @@ -48,7 +45,7 @@ public class QueryInsightsReporterTest { */ private static final String GOLDEN_BASE64 = "AQNQUEwQUFBMOm5vZGUtMToxMjM0NQZub2RlLTElc291cmNlPXRhYmxlIHwgd2hlcmUgaWRlbnRpZmllciA+ICoqKo" - + "DQlf+8MYgIwMjmxgLwzaunAQIGZW50aXR5CWVtcGxveWVlcxphZG1pbnx8YWxsX2FjY2Vzc3xfX3VzZXJfXwE="; + + "DQlf+8MYgIwMjmxgLwzaunARphZG1pbnx8YWxsX2FjY2Vzc3xfX3VzZXJfXwE="; private TransportService transportService; private DiscoveryNode localNode; @@ -76,7 +73,7 @@ public void setUp() { .sendRequest(eq(localNode), any(String.class), any(), any()); } - /** Fixed inputs covering every field, including a multi-index list and a user string. */ + /** Fixed inputs covering every field, including a user string. */ private void reportSample() { QueryInsightsReporter.report( transportService, @@ -89,7 +86,6 @@ private void reportSample() { 1032L, 685_352_000L, 350_938_864L, - Arrays.asList("entity", "employees"), "admin||all_access|__user__", true); } @@ -120,12 +116,6 @@ public void decodesInConsumerOrder() throws IOException { assertEquals(1032L, in.readVLong()); // latencyMillis assertEquals(685_352_000L, in.readVLong()); // cpuNanos assertEquals(350_938_864L, in.readVLong()); // memoryBytes - int indexCount = in.readVInt(); - String[] indices = new String[indexCount]; - for (int i = 0; i < indexCount; i++) { - indices[i] = in.readString(); - } - assertArrayEquals(new String[] {"entity", "employees"}, indices); assertEquals("admin||all_access|__user__", in.readString()); // userInfo assertTrue(in.readBoolean()); // failed assertEquals("trailing bytes would be misread by the consumer", 0, in.available()); @@ -141,7 +131,7 @@ public void formatVersionIsOne() { @Test public void nullsBecomeEmptyStringsAndNegativesClampToZero() throws IOException { QueryInsightsReporter.report( - transportService, localNode, null, null, null, null, 5L, -1L, -2L, -3L, null, null, false); + transportService, localNode, null, null, null, null, 5L, -1L, -2L, -3L, null, false); try (StreamInput in = captured.get().bytes().streamInput()) { in.readVInt(); assertEquals("", in.readString()); @@ -152,42 +142,11 @@ public void nullsBecomeEmptyStringsAndNegativesClampToZero() throws IOException assertEquals(0L, in.readVLong()); assertEquals(0L, in.readVLong()); assertEquals(0L, in.readVLong()); - assertEquals(0, in.readVInt()); // null index list -> empty assertEquals("", in.readString()); assertFalse(in.readBoolean()); } } - @Test - public void nullIndexEntryBecomesEmptyString() throws IOException { - QueryInsightsReporter.report( - transportService, - localNode, - "PPL", - "m", - "n", - "q", - 1L, - 1L, - 1L, - 1L, - Arrays.asList("a", null, "b"), - "", - false); - try (StreamInput in = captured.get().bytes().streamInput()) { - in.readVInt(); - for (int i = 0; i < 4; i++) { - in.readString(); - } - for (int i = 0; i < 4; i++) { - in.readVLong(); - } - assertEquals(3, in.readVInt()); - assertEquals( - List.of("a", "", "b"), List.of(in.readString(), in.readString(), in.readString())); - } - } - @Test public void sendRunsWithCallerContextStashed() { // The send is an internal cluster:admin action; it must not be authorized against the end diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/PPLQueryRunner.java b/ppl/src/main/java/org/opensearch/sql/ppl/PPLQueryRunner.java index 2a0a95d8238..0226dcdeb59 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLQueryRunner.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLQueryRunner.java @@ -43,8 +43,8 @@ public final class PPLQueryRunner implements QueryRunner { /** * @param pplService live PPL service; not owned by the runner * @param request rich PPL request; carries include_metadata, time_bounds, format, etc. - * @param queryInsightsSink receives the PII-scrubbed query text and the indices it reads; supply - * {@link PPLService#NO_QUERY_INSIGHTS_SINK} when no telemetry is wanted + * @param queryInsightsSink receives the PII-scrubbed query text and whether it is an explain; + * supply {@link PPLService#NO_QUERY_INSIGHTS_SINK} when no telemetry is wanted * @param clock time source; used to measure {@code tookMillis} * @param cancelExecution stops the engine's in-flight execution; invoked at most once, only when * {@link #cancel()} wins against completion 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 93235459642..b6aa37227ad 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java @@ -9,7 +9,6 @@ import static org.opensearch.sql.executor.execution.QueryPlanFactory.NO_CONSUMER_RESPONSE_LISTENER; import com.google.common.base.Strings; -import java.util.List; import java.util.function.Consumer; import lombok.extern.log4j.Log4j2; import org.antlr.v4.runtime.tree.ParseTree; @@ -32,7 +31,6 @@ import org.opensearch.sql.ppl.parser.AstBuilder; import org.opensearch.sql.ppl.parser.AstStatementBuilder; import org.opensearch.sql.ppl.utils.PPLQueryDataAnonymizer; -import org.opensearch.sql.ppl.utils.PPLQueryIndexExtractor; /** PPLService. */ @Log4j2 @@ -79,7 +77,7 @@ public void execute( execute(request, queryListener, explainListener, NO_QUERY_INSIGHTS_SINK); } - /** Variant that hands query-insights metadata (anonymized query + indices) to {@code sink}. */ + /** Variant that hands query-insights metadata for the parsed query to {@code sink}. */ public void execute( PPLQueryRequest request, ResponseListener queryListener, @@ -103,7 +101,7 @@ public void explain(PPLQueryRequest request, ResponseListener l explain(request, listener, NO_QUERY_INSIGHTS_SINK); } - /** Variant that hands query-insights metadata (anonymized query + indices) to {@code sink}. */ + /** Variant that hands query-insights metadata for the parsed query to {@code sink}. */ public void explain( PPLQueryRequest request, ResponseListener listener, @@ -125,7 +123,7 @@ public void analyze(PPLQueryRequest request, ResponseListener l analyze(request, listener, NO_QUERY_INSIGHTS_SINK); } - /** Variant that hands query-insights metadata (anonymized query + indices) to {@code sink}. */ + /** Variant that hands query-insights metadata for the parsed query to {@code sink}. */ public void analyze( PPLQueryRequest request, ResponseListener listener, @@ -177,8 +175,8 @@ private record Prepared(Statement statement, QueryInsightsMetadata metadata) {} /** * Parses the request once and collects the query-insights metadata from the AST it produced, so - * the recorded query text and source indices are the ones that execute. The caller supplies the - * statement-builder context because each route configures it differently. + * the metadata describes the statement that executes. The caller supplies the statement-builder + * context because each route configures it differently. */ private Prepared prepare( PPLQueryRequest request, AstStatementBuilder.StatementBuilderContext context) { @@ -203,9 +201,7 @@ private Prepared prepare( } log.info("[{}] Incoming request {}", QueryContext.getRequestId(), anonymized); return new Prepared( - statement, - new QueryInsightsMetadata( - anonymized, extractIndexNames(statement), statement instanceof Explain)); + statement, new QueryInsightsMetadata(anonymized, statement instanceof Explain)); } /** Explain format requested on the URL, or null when none was given. */ @@ -217,25 +213,4 @@ private static org.opensearch.sql.protocol.response.format.Format explainFormat( return org.opensearch.sql.protocol.response.format.Format.ofExplain(request.getFormat()) .orElse(null); } - - /** - * Best-effort source index names from an already-parsed statement (unwrapping an {@code explain} - * to its inner query); empty on any failure. - * - *

Package-private so the fallbacks can be tested: neither is reachable through the public API, - * because the parser only ever yields a {@link Query} or an {@link Explain} wrapping one. - */ - static List extractIndexNames(Statement statement) { - try { - Statement inner = - statement instanceof Explain ? ((Explain) statement).getStatement() : statement; - if (inner instanceof Query query) { - return PPLQueryIndexExtractor.extractIndexNames(query.getPlan()); - } - return List.of(); - } catch (Exception e) { - log.debug("[{}] Failed to extract PPL index names", QueryContext.getRequestId(), e); - return List.of(); - } - } } diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/QueryInsightsMetadata.java b/ppl/src/main/java/org/opensearch/sql/ppl/QueryInsightsMetadata.java index c638990bf2c..067cd10c69d 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/QueryInsightsMetadata.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/QueryInsightsMetadata.java @@ -5,16 +5,12 @@ package org.opensearch.sql.ppl; -import java.util.List; - /** - * Best-effort metadata about a PPL query, derived once from the parsed AST that is executed, for - * consumers such as Query Insights. + * Metadata about a PPL query, derived once from the parsed AST that is executed, for consumers such + * as tracing and Query Insights. * * @param anonymizedQuery the query with literals masked by {@code PPLQueryDataAnonymizer} - * @param indices the source index name(s) the query reads from; empty when none are resolvable * @param explain whether the statement is an {@code explain}, which plans but runs no search. The * endpoint alone cannot tell: {@code explain} is also valid inside the query text. */ -public record QueryInsightsMetadata( - String anonymizedQuery, List indices, boolean explain) {} +public record QueryInsightsMetadata(String anonymizedQuery, boolean explain) {} diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java b/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java deleted file mode 100644 index 298eaeb8e9a..00000000000 --- a/ppl/src/main/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractor.java +++ /dev/null @@ -1,239 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.ppl.utils; - -import java.util.ArrayList; -import java.util.LinkedHashSet; -import java.util.List; -import java.util.Set; -import org.apache.commons.lang3.tuple.Pair; -import org.opensearch.sql.ast.AbstractNodeVisitor; -import org.opensearch.sql.ast.Node; -import org.opensearch.sql.ast.expression.Alias; -import org.opensearch.sql.ast.expression.Field; -import org.opensearch.sql.ast.expression.Let; -import org.opensearch.sql.ast.expression.QualifiedName; -import org.opensearch.sql.ast.expression.UnresolvedExpression; -import org.opensearch.sql.ast.expression.subquery.ExistsSubquery; -import org.opensearch.sql.ast.expression.subquery.InSubquery; -import org.opensearch.sql.ast.expression.subquery.ScalarSubquery; -import org.opensearch.sql.ast.tree.Aggregation; -import org.opensearch.sql.ast.tree.Chart; -import org.opensearch.sql.ast.tree.Eval; -import org.opensearch.sql.ast.tree.FillNull; -import org.opensearch.sql.ast.tree.Filter; -import org.opensearch.sql.ast.tree.Foreach; -import org.opensearch.sql.ast.tree.Join; -import org.opensearch.sql.ast.tree.Relation; -import org.opensearch.sql.ast.tree.StreamWindow; -import org.opensearch.sql.ast.tree.UnresolvedPlan; -import org.opensearch.sql.ast.tree.Window; - -/** - * Collects the source index name(s) a PPL query reads from by walking the parsed AST (via {@link - * Relation#getQualifiedNames()}) instead of regex-matching text, which mishandles multi-index and - * datasource-qualified sources. - * - *

Plan sources are found generically: the walk follows {@code getChild()} plus {@link - * UnresolvedPlan#getSources()}, which each node uses to declare a plan it holds in a field (a - * sub-search, a lookup table, a join's right branch). A new command therefore only has to declare - * its sources on the node itself to be covered here. - * - *

Subqueries carried in expressions are not reachable that way, so every command whose - * expression the engine evaluates is bridged explicitly: {@link Filter} and {@link Join} ON - * conditions, {@link Eval}, {@link Foreach} and {@link FillNull} assignments, and {@link - * Aggregation}, {@link Window}, {@link StreamWindow} and {@link Chart} function arguments. A - * subquery in any of these positions parses and, except under {@code eventstats}, executes, so its - * inner source is a real read and belongs in the metadata. Commands that take only field names, - * literals or patterns cannot carry one. - * - *

Returns distinct names in encounter order. This is best-effort metadata; a source it cannot - * resolve is simply omitted and query execution is unaffected. - */ -public final class PPLQueryIndexExtractor { - - private PPLQueryIndexExtractor() {} - - /** Return the distinct source index name(s) referenced by {@code plan}, in encounter order. */ - public static List extractIndexNames(UnresolvedPlan plan) { - Set names = new LinkedHashSet<>(); - if (plan != null) { - plan.accept(new Collector(), names); - } - return new ArrayList<>(names); - } - - private static final class Collector extends AbstractNodeVisitor> { - @Override - public Void visitRelation(Relation node, Set names) { - for (QualifiedName qualifiedName : node.getQualifiedNames()) { - names.add(qualifiedName.toString()); - } - return null; - } - - /** - * Walks {@code getChild()} as usual, then the node's declared {@link - * UnresolvedPlan#getSources()}. Commands whose source lives in a field are covered without this - * collector knowing each one. - * - *

{@code Node.getChild()} defaults to null and several expression nodes keep that default, - * so the walk treats null as "no children" rather than relying on the base visitor. - */ - @Override - public Void visitChildren(Node node, Set names) { - List children = node.getChild(); - if (children != null) { - for (Node child : children) { - if (child != null) { - child.accept(this, names); - } - } - } - if (node instanceof UnresolvedPlan plan) { - for (UnresolvedPlan source : plan.getSources()) { - if (source != null) { - source.accept(this, names); - } - } - } - return null; - } - - // Expression bridges. Each command below evaluates an expression that may hold a subquery; the - // plan walk does not descend into expressions, so hand each one to visit(). Null-safety lives - // in - // visit()/visitAll() alone so a hand-built AST with a missing part costs nothing. - - @Override - public Void visitFilter(Filter node, Set names) { - visitChildren(node, names); - visit(node.getCondition(), names); - return null; - } - - @Override - public Void visitEval(Eval node, Set names) { - visitChildren(node, names); - if (node.getExpressionList() != null) { - for (Let let : node.getExpressionList()) { - visit(let == null ? null : let.getExpression(), names); - } - } - return null; - } - - @Override - public Void visitAlias(Alias node, Set names) { - // Alias keeps Node's null getChild(), and every aggregation argument and most eval - // assignments arrive wrapped in one, so a subquery beneath it would otherwise be invisible. - visit(node.getDelegated(), names); - return null; - } - - @Override - public Void visitAggregation(Aggregation node, Set names) { - visitChildren(node, names); - visitAll(node.getAggExprList(), names); - visitAll(node.getGroupExprList(), names); - visitAll(node.getSortExprList(), names); - visit(node.getSpan(), names); - return null; - } - - @Override - public Void visitWindow(Window node, Set names) { - visitChildren(node, names); - visitAll(node.getWindowFunctionList(), names); - visitAll(node.getGroupList(), names); - return null; - } - - @Override - public Void visitStreamWindow(StreamWindow node, Set names) { - visitChildren(node, names); - visitAll(node.getWindowFunctionList(), names); - visitAll(node.getGroupList(), names); - visit(node.getResetBefore(), names); - visit(node.getResetAfter(), names); - return null; - } - - @Override - public Void visitChart(Chart node, Set names) { - visitChildren(node, names); - visit(node.getAggregationFunction(), names); - visit(node.getRowSplit(), names); - visit(node.getColumnSplit(), names); - return null; - } - - @Override - public Void visitFillNull(FillNull node, Set names) { - visitChildren(node, names); - visit(node.getReplacementForAll().orElse(null), names); - if (node.getReplacementPairs() != null) { - for (Pair pair : node.getReplacementPairs()) { - visit(pair == null ? null : pair.getRight(), names); - } - } - return null; - } - - @Override - public Void visitForeach(Foreach node, Set names) { - visitChildren(node, names); - visit(node.getCollectionExpression(), names); - if (node.getEvalClauses() != null) { - for (Foreach.ForeachEvalClause clause : node.getEvalClauses()) { - visit(clause == null ? null : clause.getExpression(), names); - } - } - return null; - } - - @Override - public Void visitJoin(Join node, Set names) { - // Both sides come through getChild() + getSources(); the ON condition is an expression and - // can carry a subquery (on l.a = r.a AND r.a in [ source=inner | ... ]). - visitChildren(node, names); - visit(node.getJoinCondition() == null ? null : node.getJoinCondition().orElse(null), names); - return null; - } - - private void visit(UnresolvedExpression expr, Set names) { - if (expr != null) { - expr.accept(this, names); - } - } - - private void visitAll(List exprs, Set names) { - if (exprs != null) { - for (UnresolvedExpression expr : exprs) { - visit(expr, names); - } - } - } - - @Override - public Void visitScalarSubquery(ScalarSubquery node, Set names) { - node.getQuery().accept(this, names); - return null; - } - - @Override - public Void visitInSubquery(InSubquery node, Set names) { - node.getQuery().accept(this, names); - return null; - } - - @Override - public Void visitExistsSubquery(ExistsSubquery node, Set names) { - node.getQuery().accept(this, names); - return null; - } - } -} diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java index 86c254a008a..2555167650a 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java @@ -19,10 +19,6 @@ import org.junit.runner.RunWith; import org.mockito.Mock; import org.mockito.junit.MockitoJUnitRunner; -import org.opensearch.sql.ast.AbstractNodeVisitor; -import org.opensearch.sql.ast.statement.Query; -import org.opensearch.sql.ast.statement.Statement; -import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.executor.AnalyzeResponse; @@ -31,7 +27,6 @@ import org.opensearch.sql.executor.ExecutionEngine.ExplainResponse; import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; import org.opensearch.sql.executor.QueryService; -import org.opensearch.sql.executor.QueryType; import org.opensearch.sql.executor.execution.QueryPlanFactory; import org.opensearch.sql.executor.pagination.Cursor; import org.opensearch.sql.ppl.antlr.PPLSyntaxParser; @@ -154,12 +149,11 @@ public void testExecutePassesAnonymizedQueryToSink() { getExplainListener(false), metadata::set); - // The sink receives the anonymized query (literal masked, never the raw value) and the source - // index resolved from the same AST that executes. + // The sink receives the anonymized query: the literal is masked, never the raw value. Assert.assertNotNull(metadata.get()); Assert.assertTrue(metadata.get().anonymizedQuery().contains("***")); Assert.assertFalse(metadata.get().anonymizedQuery().contains("42")); - Assert.assertEquals(Collections.singletonList("t"), metadata.get().indices()); + Assert.assertFalse(metadata.get().explain()); } @Test @@ -173,7 +167,7 @@ public void testExplainPassesAnonymizedQueryToSink() { Assert.assertNotNull(metadata.get()); Assert.assertTrue(metadata.get().anonymizedQuery().contains("***")); Assert.assertFalse(metadata.get().anonymizedQuery().contains("42")); - Assert.assertEquals(Collections.singletonList("t"), metadata.get().indices()); + Assert.assertTrue(metadata.get().explain()); } @Test @@ -189,7 +183,6 @@ public void testAnalyzePassesAnonymizedQueryToSink() { Assert.assertNotNull(metadata.get()); Assert.assertTrue(metadata.get().anonymizedQuery().contains("***")); Assert.assertFalse(metadata.get().anonymizedQuery().contains("42")); - Assert.assertEquals(Collections.singletonList("t"), metadata.get().indices()); Assert.assertFalse(metadata.get().explain()); } @@ -225,36 +218,6 @@ public void onFailure(Exception e) { Assert.assertNull(metadata.get()); } - @Test - public void testExtractIndexNamesOnNonQueryStatementYieldsEmpty() { - // Defensive fallback: the parser only produces Query or Explain, so this cannot happen through - // the public API. Pinned so the fallback stays a fallback and does not become a null return. - Statement notAQuery = new Statement() {}; - Assert.assertEquals(Collections.emptyList(), PPLService.extractIndexNames(notAQuery)); - } - - @Test - public void testExtractIndexNamesSwallowsWalkFailure() { - // Index extraction is best-effort metadata: a plan that blows up during the walk must yield no - // indices rather than failing the query that is about to execute. - UnresolvedPlan exploding = - new UnresolvedPlan() { - @Override - public UnresolvedPlan attach(UnresolvedPlan child) { - return this; - } - - @Override - public R accept(AbstractNodeVisitor visitor, C context) { - throw new IllegalStateException("walk blew up"); - } - }; - - Assert.assertEquals( - Collections.emptyList(), - PPLService.extractIndexNames(new Query(exploding, 0, QueryType.PPL))); - } - @Test public void testExecuteCsvFormatShouldPass() { doAnswer( diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java deleted file mode 100644 index 8641fb8f914..00000000000 --- a/ppl/src/test/java/org/opensearch/sql/ppl/utils/PPLQueryIndexExtractorTest.java +++ /dev/null @@ -1,375 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.ppl.utils; - -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertTrue; - -import java.util.Arrays; -import java.util.Collections; -import java.util.List; -import java.util.Optional; -import org.apache.commons.lang3.tuple.Pair; -import org.junit.Test; -import org.opensearch.sql.ast.expression.DataType; -import org.opensearch.sql.ast.expression.Field; -import org.opensearch.sql.ast.expression.Let; -import org.opensearch.sql.ast.expression.Literal; -import org.opensearch.sql.ast.expression.QualifiedName; -import org.opensearch.sql.ast.expression.UnresolvedExpression; -import org.opensearch.sql.ast.tree.Aggregation; -import org.opensearch.sql.ast.tree.AppendCol; -import org.opensearch.sql.ast.tree.AppendPipe; -import org.opensearch.sql.ast.tree.Eval; -import org.opensearch.sql.ast.tree.FillNull; -import org.opensearch.sql.ast.tree.Filter; -import org.opensearch.sql.ast.tree.Foreach; -import org.opensearch.sql.ast.tree.Join; -import org.opensearch.sql.ast.tree.Relation; -import org.opensearch.sql.ast.tree.UnresolvedPlan; -import org.opensearch.sql.ppl.AstPlanningTestBase; - -/** Unit tests for {@link PPLQueryIndexExtractor}. */ -public class PPLQueryIndexExtractorTest extends AstPlanningTestBase { - - private List indices(String query) { - return PPLQueryIndexExtractor.extractIndexNames((UnresolvedPlan) plan(query)); - } - - @Test - public void singleSource() { - assertEquals(List.of("accounts"), indices("source=accounts | fields firstname")); - } - - @Test - public void indexKeyword() { - assertEquals(List.of("logs"), indices("search index=logs")); - } - - @Test - public void multipleCommaSeparatedSourcesWithSpace() { - // The reviewer's case: "source=accounts, account2" must keep BOTH indices — the old regex - // stopped at the whitespace after the comma and dropped account2. - assertEquals(List.of("accounts", "account2"), indices("source=accounts, account2")); - } - - @Test - public void multipleCommaSeparatedSourcesNoSpace() { - assertEquals(List.of("a", "b", "c"), indices("source=a,b,c | stats count()")); - } - - @Test - public void joinCapturesBothSides() { - List result = indices("source=t1 | join on t1.id = t2.id t2"); - assertTrue(result.contains("t1")); - assertTrue(result.contains("t2")); - } - - @Test - public void lookupCapturesLookupTable() { - List result = indices("source=t1 | lookup t2 id"); - assertTrue(result.contains("t1")); - assertTrue(result.contains("t2")); - } - - @Test - public void datasourceQualifiedName() { - assertEquals(List.of("myds.myindex"), indices("source=myds.myindex | head 1")); - } - - @Test - public void filterAndFieldsDoNotAddIndices() { - assertEquals(List.of("accounts"), indices("source=accounts | where age > 30 | fields name")); - } - - @Test - public void appendCapturesSubSearchSource() { - List result = indices("source=accounts | append [ source=account2 | stats count() ]"); - assertTrue(result.contains("accounts")); - assertTrue(result.contains("account2")); - } - - @Test - public void subqueryInWhereCapturesInnerSource() { - List result = indices("source=outer | where a in [ source=inner | fields b ]"); - assertTrue(result.contains("outer")); - assertTrue(result.contains("inner")); - } - - @Test - public void evalScalarSubqueryCapturesInnerSource() { - List result = indices("source=outer | eval m = [ source=inner | stats max(b) ]"); - assertTrue(result.contains("outer")); - assertTrue(result.contains("inner")); - } - - // Every command whose expression the engine evaluates can carry a subquery. These were found by - // probing each expression-bearing command; all but eventstats/streamstats execute on a live - // cluster, so the inner source is a real read. - - @Test - public void joinOnConditionSubqueryCapturesInnerSource() { - // Documented form: `on l.a = r.a AND r.a in [ source = inner | fields d ]` (subquery.md). - List result = - indices( - "source = t1 | inner join left = l right = r" - + " on l.a = r.a AND r.a in [ source = inner | fields d ] t2"); - assertEquals(List.of("t1", "t2", "inner"), result); - } - - @Test - public void joinOnConditionAndRightSubsearchBothCaptured() { - List result = - indices( - "source = t1 | inner join left = l right = r on l.a = r.a AND r.a in [ source = inner |" - + " fields d ] [ source = t2 | fields a ]"); - assertEquals(List.of("t1", "t2", "inner"), result); - } - - @Test - public void statsArgumentSubqueryCapturesInnerSource() { - List result = indices("source=outer | stats max([ source=inner | stats max(b) ])"); - assertEquals(List.of("outer", "inner"), result); - } - - @Test - public void eventstatsArgumentSubqueryCapturesInnerSource() { - List result = indices("source=outer | eventstats max([ source=inner | stats max(b) ])"); - assertEquals(List.of("outer", "inner"), result); - } - - @Test - public void streamstatsArgumentSubqueryCapturesInnerSource() { - List result = - indices("source=outer | streamstats max([ source=inner | stats max(b) ])"); - assertEquals(List.of("outer", "inner"), result); - } - - @Test - public void chartArgumentSubqueryCapturesInnerSource() { - List result = indices("source=outer | chart max([ source=inner | stats max(b) ]) by c"); - assertEquals(List.of("outer", "inner"), result); - } - - @Test - public void timechartArgumentSubqueryCapturesInnerSource() { - List result = indices("source=outer | timechart max([ source=inner | stats max(b) ])"); - assertEquals(List.of("outer", "inner"), result); - } - - @Test - public void fillnullReplacementSubqueryCapturesInnerSource() { - assertEquals( - List.of("outer", "inner"), - indices("source=outer | fillnull with [ source=inner | stats max(b) ] in a")); - assertEquals( - List.of("outer", "inner"), - indices("source=outer | fillnull value=[ source=inner | stats max(b) ]")); - assertEquals( - List.of("outer", "inner"), - indices("source=outer | fillnull using a = [ source=inner | stats max(b) ]")); - } - - @Test - public void foreachAssignmentSubqueryCapturesInnerSource() { - List result = - indices("source=outer | foreach a [ eval x = [ source=inner | stats max(b) ] ]"); - assertEquals(List.of("outer", "inner"), result); - } - - @Test - public void aliasedExpressionIsWalkedThrough() { - // Alias keeps Node's null getChild(); without an explicit bridge the aggregate beneath it, and - // any subquery in its argument, is invisible and the walk itself NPEs. - List result = - indices("source=outer | stats max([ source=inner | stats max(b) ]) as m by c"); - assertEquals(List.of("outer", "inner"), result); - } - - @Test - public void existsSubqueryCapturesInnerSource() { - // exists takes a different AST node than `in`, so it needs its own visit to be reached. - List result = indices("source=outer | where exists [ source=inner | where a = b ]"); - assertTrue(result.contains("outer")); - assertTrue(result.contains("inner")); - } - - @Test - public void graphLookupCapturesFromTable() { - List result = - indices( - "source=t | graphLookup employees start=reportsTo edge=manager-->name" - + " as reportingHierarchy"); - assertTrue(result.contains("t")); - assertTrue(result.contains("employees")); - } - - @Test - public void appendColCapturesSubSearchSource() { - // Built directly, not parsed: the grammar runs appendcol's sub-search over the piped input, so - // a nested source= can't be expressed in text. - UnresolvedPlan plan = new AppendCol(false, subPipelineOver("appended")).attach(baseRelation()); - List result = PPLQueryIndexExtractor.extractIndexNames(plan); - assertTrue(result.contains("base")); - assertTrue(result.contains("appended")); - } - - @Test - public void appendPipeCapturesSubQuerySource() { - // Built directly for the same reason as appendcol above. - UnresolvedPlan plan = new AppendPipe(subPipelineOver("appended")).attach(baseRelation()); - List result = PPLQueryIndexExtractor.extractIndexNames(plan); - assertTrue(result.contains("base")); - assertTrue(result.contains("appended")); - } - - private UnresolvedPlan baseRelation() { - return new Relation(new QualifiedName("base")); - } - - /** A {@link Relation} nested under a piped command, like real parser output (not a bare root). */ - private UnresolvedPlan subPipelineOver(String index) { - return new Filter(new Literal(true, DataType.BOOLEAN)) - .attach(new Relation(new QualifiedName(index))); - } - - @Test - public void nullPlanYieldsEmpty() { - assertTrue(PPLQueryIndexExtractor.extractIndexNames(null).isEmpty()); - } - - // The walk guards against absent sub-plans and sub-expressions because it runs on best-effort - // metadata: a shape it did not expect must cost the query nothing. The parser does not produce - // these shapes, so they are built by hand. - - @Test - public void nullSourceInGetSourcesIsSkipped() { - UnresolvedPlan declaresNullSource = - new UnresolvedPlan() { - @Override - public List getSources() { - return Collections.singletonList(null); - } - - @Override - public List getChild() { - return List.of(); - } - - @Override - public UnresolvedPlan attach(UnresolvedPlan child) { - return this; - } - }; - assertEquals(List.of(), PPLQueryIndexExtractor.extractIndexNames(declaresNullSource)); - } - - @Test - public void filterWithoutConditionStillWalksItsChild() { - UnresolvedPlan plan = new Filter(null).attach(baseRelation()); - assertEquals(List.of("base"), PPLQueryIndexExtractor.extractIndexNames(plan)); - } - - @Test - public void evalWithoutExpressionsStillWalksItsChild() { - UnresolvedPlan plan = new Eval(null).attach(baseRelation()); - assertEquals(List.of("base"), PPLQueryIndexExtractor.extractIndexNames(plan)); - } - - @Test - public void planWithNullChildListIsWalkedAsLeaf() { - // Node.getChild() defaults to null; a plan that keeps the default must read as "no children", - // not NPE, and still contribute its declared sources. - UnresolvedPlan nullChildren = - new UnresolvedPlan() { - @Override - public List getSources() { - return List.of(baseRelation()); - } - - @Override - public UnresolvedPlan attach(UnresolvedPlan child) { - return this; - } - }; - assertEquals(List.of("base"), PPLQueryIndexExtractor.extractIndexNames(nullChildren)); - } - - @Test - public void nullEntryInChildListIsSkipped() { - UnresolvedPlan holeyChildren = - new UnresolvedPlan() { - @Override - public List getChild() { - return Arrays.asList(null, baseRelation()); - } - - @Override - public UnresolvedPlan attach(UnresolvedPlan child) { - return this; - } - }; - assertEquals(List.of("base"), PPLQueryIndexExtractor.extractIndexNames(holeyChildren)); - } - - @Test - public void aggregationWithNullExpressionListsStillWalksItsChild() { - UnresolvedPlan plan = new Aggregation(null, null, null).attach(baseRelation()); - assertEquals(List.of("base"), PPLQueryIndexExtractor.extractIndexNames(plan)); - } - - @Test - public void fillnullWithNullOrMissingReplacementsStillWalksItsChild() { - assertEquals( - List.of("base"), - PPLQueryIndexExtractor.extractIndexNames( - FillNull.ofVariousValue(null).attach(baseRelation()))); - List> holey = Arrays.asList(null, Pair.of(null, null)); - assertEquals( - List.of("base"), - PPLQueryIndexExtractor.extractIndexNames( - FillNull.ofVariousValue(holey).attach(baseRelation()))); - } - - @Test - public void foreachWithNullOrMissingClausesStillWalksItsChild() { - assertEquals( - List.of("base"), - PPLQueryIndexExtractor.extractIndexNames( - new Foreach(null, null, null, null, null).attach(baseRelation()))); - List holey = - Arrays.asList(null, new Foreach.ForeachEvalClause("x", null)); - assertEquals( - List.of("base"), - PPLQueryIndexExtractor.extractIndexNames( - new Foreach(null, null, null, null, holey).attach(baseRelation()))); - } - - @Test - public void joinWithNullConditionStillWalksBothSides() { - UnresolvedPlan plan = - new Join( - new Relation(new QualifiedName("right")), - Optional.empty(), - Optional.empty(), - Join.JoinType.INNER, - null, - new Join.JoinHint(), - Optional.empty(), - new org.opensearch.sql.ast.expression.Argument.ArgumentMap(List.of())) - .attach(baseRelation()); - assertEquals(List.of("base", "right"), PPLQueryIndexExtractor.extractIndexNames(plan)); - } - - @Test - public void evalWithNullAssignmentOrExpressionStillWalksItsChild() { - // Covers both halves of the guard: a null Let, and a Let carrying no expression. - Let noExpression = new Let(new Field(new QualifiedName("x")), null); - List lets = Arrays.asList(null, noExpression); - UnresolvedPlan plan = new Eval(lets).attach(baseRelation()); - assertEquals(List.of("base"), PPLQueryIndexExtractor.extractIndexNames(plan)); - } -} From c7f3018943f7f39a903d4119bcdb2e5f214f8bd5 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 27 Oct 2026 12:00:30 -0700 Subject: [PATCH 23/26] feat(ppl): Record the PPL query text as submitted The Top N record carried the anonymized query text, with every literal masked, which hides what makes a slow or failed query diagnosable: which source and which filter values. Query Insights stores a DSL search's source as sent, literals included, and its RBAC filter (search.insights.top_queries.filter_by_mode) governs who can read either. Record the query as submitted. The tracing span and the request log line keep the anonymized text. Signed-off-by: Kishore Kumaar Natarajan --- docs/user/ppl/admin/settings.md | 4 ++-- .../sql/plugin/transport/PPLQueryTask.java | 12 ++++++------ .../plugin/transport/TransportPPLQueryAction.java | 8 +++++--- .../sql/plugin/transport/PPLQueryTaskTest.java | 8 ++++---- 4 files changed, 17 insertions(+), 15 deletions(-) diff --git a/docs/user/ppl/admin/settings.md b/docs/user/ppl/admin/settings.md index 2eed41f5fe9..cd34905c518 100644 --- a/docs/user/ppl/admin/settings.md +++ b/docs/user/ppl/admin/settings.md @@ -482,13 +482,13 @@ Expected output: ### Description -This setting controls whether completed PPL queries are recorded into Query Insights Top N. When enabled, the PPL plugin reports each finished PPL query (with its latency, CPU, memory, and — on a security-enabled cluster — the requesting user) to the Query Insights plugin, and tags the child DSL searches a PPL query spawns so they appear as sub-queries. +This setting controls whether completed PPL queries are recorded into Query Insights Top N. When enabled, the PPL plugin reports each finished PPL query (with its query text, latency, CPU, memory, and — on a security-enabled cluster — the requesting user) to the Query Insights plugin, and tags the child DSL searches a PPL query spawns so they appear as sub-queries. 1. The default value is false (opt-in). 2. This setting is node scope, matching the rest of the `search.insights.top_queries.*` family. 3. It is dynamic: update it through the cluster settings API (as in the example below), and the value applies per node. -Note on user identity: when a security plugin is installed, the recorded record includes the requesting user (username, roles, backend roles). Who can see that identity in Top N results is governed by the Query Insights plugin's own RBAC filter, `search.insights.top_queries.filter_by_mode` — default `none` (no per-user filtering), `username` (a user sees only their own queries; admins see all), or `backend_roles`. Set it to `username` or `backend_roles` if user identities should not be visible to every Top N reader. +Note on query text and user identity: the record stores the query as it was submitted, including any literal values in it, the same way Query Insights stores the source of a DSL search. When a security plugin is installed, the record also includes the requesting user (username, roles, backend roles). Who can see these in Top N results is governed by the Query Insights plugin's own RBAC filter, `search.insights.top_queries.filter_by_mode` — default `none` (no per-user filtering), `username` (a user sees only their own queries; admins see all), or `backend_roles`. Set it to `username` or `backend_roles` if other users' queries and identities should not be visible to every Top N reader. Note: This setting is **registered by the Query Insights plugin** (it belongs to the `search.insights.top_queries.*` family), not by the PPL plugin. The PPL plugin reads its value to decide whether to report. It therefore only takes effect when the Query Insights plugin is installed; when Query Insights is absent the setting is unregistered and PPL reporting is a no-op. diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java index 668441c8def..90adc4ba45e 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java @@ -26,8 +26,8 @@ public class PPLQueryTask extends CancellableTask implements ThreadResourceAccou /** Security {@code _opendistro_security_user_info} string; null when unsecured. */ private volatile String queryInsightsUserInfo; - /** Anonymized query text ({@code PPLQueryDataAnonymizer}); null when anonymization didn't run. */ - private volatile String queryInsightsAnonymizedQuery; + /** The query text as submitted; null until the query has parsed. */ + private volatile String queryInsightsQueryText; /** Whether the query failed; recorded so a failed PPL query is flagged in Top N. */ private volatile boolean queryInsightsFailed = false; @@ -75,12 +75,12 @@ public String getQueryInsightsUserInfo() { return queryInsightsUserInfo; } - public void setQueryInsightsAnonymizedQuery(String anonymizedQuery) { - this.queryInsightsAnonymizedQuery = anonymizedQuery; + public void setQueryInsightsQueryText(String queryText) { + this.queryInsightsQueryText = queryText; } - public String getQueryInsightsAnonymizedQuery() { - return queryInsightsAnonymizedQuery; + public String getQueryInsightsQueryText() { + return queryInsightsQueryText; } public void setQueryInsightsExplain(boolean explain) { 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 e576a76ebff..9f36629fe8a 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 @@ -300,7 +300,7 @@ protected void innerOnFailure(Exception e) {} * children's cost would be counted against both records. */ static boolean shouldReportToQueryInsights(PPLQueryTask pplQueryTask) { - return pplQueryTask.getQueryInsightsAnonymizedQuery() != null + return pplQueryTask.getQueryInsightsQueryText() != null && !pplQueryTask.isQueryInsightsExplain() && !pplQueryTask.isQueryInsightsNested(); } @@ -312,7 +312,7 @@ private void writeQueryInsightsRecord(PPLQueryTask pplQueryTask) { } try { String nodeId = clusterServiceRef.localNode().getId(); - String queryText = pplQueryTask.getQueryInsightsAnonymizedQuery(); + String queryText = pplQueryTask.getQueryInsightsQueryText(); long cpuNanos = pplQueryTask.getCpuNanos(); long memoryBytes = pplQueryTask.getAllocatedBytes(); @@ -449,7 +449,9 @@ protected void doExecute( metadata -> { rootSpan.addAttribute("db.query.text", metadata.anonymizedQuery()); if (pplQueryTask != null) { - pplQueryTask.setQueryInsightsAnonymizedQuery(metadata.anonymizedQuery()); + // The record keeps the query as written, literals included, as Query Insights does + // for a DSL search's source; the span above stays anonymized. + pplQueryTask.setQueryInsightsQueryText(transformedRequest.getRequest()); pplQueryTask.setQueryInsightsExplain(metadata.explain()); } }; diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java index 2ab513ea482..626eb621bbe 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java @@ -159,7 +159,7 @@ public void testInlineExplainIsNotReported() { // "explain source=t | ..." reaches the execute route because isExplainRequest() only looks at // the path, so the sink's flag is what keeps it out of Top N. PPLQueryTask task = newTask(); - task.setQueryInsightsAnonymizedQuery("explain source=table | fields identifier"); + task.setQueryInsightsQueryText("explain source=table | fields identifier"); task.setQueryInsightsExplain(true); assertFalse(TransportPPLQueryAction.shouldReportToQueryInsights(task)); } @@ -174,7 +174,7 @@ public void testSyntaxErrorIsNotReported() { public void testRuntimeFailureAfterPlanningIsReported() { // Planning succeeded, so the metadata is present; a later execution failure must still report. PPLQueryTask task = newTask(); - task.setQueryInsightsAnonymizedQuery("source=table | where identifier > ***"); + task.setQueryInsightsQueryText("source=table | where identifier > 42"); task.setQueryInsightsFailed(true); assertTrue(TransportPPLQueryAction.shouldReportToQueryInsights(task)); } @@ -184,7 +184,7 @@ public void testNestedQueryIsNotReported() { // An outer marker was already in the context, so the child searches rolled up into that query. // Reporting here would advertise a marker no child references and double count their cost. PPLQueryTask task = newTask(); - task.setQueryInsightsAnonymizedQuery("source=table | where identifier > ***"); + task.setQueryInsightsQueryText("source=table | where identifier > 42"); task.setQueryInsightsNested(true); assertFalse(TransportPPLQueryAction.shouldReportToQueryInsights(task)); } @@ -203,7 +203,7 @@ public void testParentMarkerIsAbsentUntilStamped() { public void testNonNestedQueryIsReportedByDefault() { // Guards the default: a plain query must not be mistaken for a nested one. PPLQueryTask task = newTask(); - task.setQueryInsightsAnonymizedQuery("source=table | where identifier > ***"); + task.setQueryInsightsQueryText("source=table | where identifier > 42"); assertFalse("nested must default to false", task.isQueryInsightsNested()); assertTrue(TransportPPLQueryAction.shouldReportToQueryInsights(task)); } From e9fa479738f28f1fb1cc391f315b02d91a28cb0b Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 27 Oct 2026 12:16:31 -0700 Subject: [PATCH 24/26] test(ppl): Match Top N records by query text The PPL record no longer carries indices and now keeps the query as written, so the ITs find a record by the index its query text reads, assert the recorded text exactly, and check that a join's searches, rather than the record, carry the join's two indices. Signed-off-by: Kishore Kumaar Natarajan --- .../sql/ppl/QueryInsightsRecordingIT.java | 24 ++++++++----------- .../QueryInsightsRecordingSecurityIT.java | 3 ++- .../sql/util/QueryInsightsRecords.java | 10 +++++--- 3 files changed, 19 insertions(+), 18 deletions(-) diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsRecordingIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsRecordingIT.java index 826a77dbc6d..4614ad7649e 100644 --- a/integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsRecordingIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsRecordingIT.java @@ -34,9 +34,9 @@ * Insights; the default integTest cluster does not, and {@link QueryInsightsReportingIT} covers * that case. * - *

Each scenario reads its own index, so a record is identified by the index it reports rather - * than by its query text, which is anonymized. Absence is only asserted after a later query's - * record has arrived, so a missing record means it was never sent rather than not yet processed. + *

Each scenario reads its own index, so a record is identified by the index its query text + * reads. Absence is only asserted after a later query's record has arrived, so a missing record + * means it was never sent rather than not yet processed. */ public class QueryInsightsRecordingIT extends PPLIntegTestCase { @@ -60,12 +60,12 @@ public void successfulQueryIsRecorded() throws Exception { JSONObject record = awaitPplRecord("qi_it_success"); assertFalse(record.getBoolean("failed")); - assertEquals(Set.of("qi_it_success"), indicesOf(record)); + // As written, literal included, the way a DSL search's source is recorded. + assertEquals("source=qi_it_success | where v > 1 | fields v", record.getString("source")); assertTrue( "parent marker must be ::, was " + record.optString("parent_marker"), record.optString("parent_marker").matches("PPL:[^:]+:\\d+")); - assertFalse("query text must be recorded", record.optString("source").isEmpty()); assertTrue(record.getLong("timestamp") >= CLASS_START_MILLIS); // The plugin measures its own threads; zero here means the accounting never ran. assertTrue(measurement(record, "cpu") > 0); @@ -76,7 +76,7 @@ public void successfulQueryIsRecorded() throws Exception { public void queryOnMissingIndexIsRecordedAsFailed() throws Exception { expectThrows(ResponseException.class, () -> executeQuery("source=qi_it_missing | fields v")); - // Fails in analysis, before any search runs; the record still names the index it targeted. + // Fails in analysis, before any search runs; the query parsed, so it is still recorded. JSONObject record = awaitPplRecord("qi_it_missing"); assertTrue(record.getBoolean("failed")); } @@ -108,11 +108,8 @@ public void syntaxErrorIsNotRecorded() throws Exception { expectThrows(ResponseException.class, () -> executeQuery("source=qi_it_syntax | wherr v > 1")); awaitBarrier(); - // A query that never parses has no text or indices to report. Every other query in this class - // names an index, so a PPL record without one could only be the syntax error. - for (JSONObject record : topQueries(QueryInsightsRecords::isPpl)) { - assertFalse("unparsed query was recorded: " + record, indicesOf(record).isEmpty()); - } + // A query that never parses is not reported. + assertNoPplRecord("qi_it_syntax"); } @Test @@ -157,7 +154,6 @@ public void childSearchesCarryTheParentMarker() throws Exception { + " | fields l.v, r.v"); JSONObject parent = awaitPplRecord("qi_it_join_left"); - assertEquals(Set.of("qi_it_join_left", "qi_it_join_right"), indicesOf(parent)); List children = awaitChildren(client(), parent.getString("parent_marker"), 2); Set childIndices = new HashSet<>(); @@ -171,8 +167,8 @@ public void childSearchesCarryTheParentMarker() throws Exception { parent.getLong("timestamp") <= child.getLong("timestamp")); childIndices.addAll(indicesOf(child)); } - // The searches read exactly the indices the record names. - assertEquals(indicesOf(parent), childIndices); + // The record names no indices; its searches carry the join's two sides. + assertEquals(Set.of("qi_it_join_left", "qi_it_join_right"), childIndices); } // --------------------------------------------------------------------------------------------- diff --git a/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsRecordingSecurityIT.java b/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsRecordingSecurityIT.java index 3b2b33ebcfd..58f4e74a512 100644 --- a/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsRecordingSecurityIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsRecordingSecurityIT.java @@ -28,7 +28,8 @@ * both Query Insights and the security plugin. {@link QueryInsightsReportingSecurityIT} covers a * secured cluster without Query Insights. * - *

Each user queries an index named after them, so a record is identified by its index. + *

Each user queries an index named after them, so a record is identified by the index its query + * text reads. */ public class QueryInsightsRecordingSecurityIT extends SecurityTestBase { diff --git a/integ-test/src/test/java/org/opensearch/sql/util/QueryInsightsRecords.java b/integ-test/src/test/java/org/opensearch/sql/util/QueryInsightsRecords.java index a8149b80e1d..118d5b3fed9 100644 --- a/integ-test/src/test/java/org/opensearch/sql/util/QueryInsightsRecords.java +++ b/integ-test/src/test/java/org/opensearch/sql/util/QueryInsightsRecords.java @@ -19,6 +19,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Predicate; +import java.util.regex.Pattern; import org.json.JSONArray; import org.json.JSONObject; import org.opensearch.client.Request; @@ -29,8 +30,8 @@ /** * Reads Query Insights Top N records for the ITs that run with Query Insights installed. * - *

Records are matched by the index they report rather than by query text, which is anonymized, - * so each scenario should read an index of its own. + *

A PPL record is matched by the index its query text reads, since the record keeps the query as + * written, so each scenario should read an index of its own. */ public final class QueryInsightsRecords { @@ -176,8 +177,11 @@ public static List awaitChildren(RestClient client, String parentMar return found.get(); } + /** The PPL record of a query that reads {@code index} with {@code source=}. */ public static Predicate pplRecordFor(String index) { - return record -> isPpl(record) && indicesOf(record).contains(index); + // Bounded on both sides, so qi_it_barrier_1 does not match source=qi_it_barrier_10. + Pattern source = Pattern.compile("(? isPpl(record) && source.matcher(record.optString("source")).find(); } public static boolean isPpl(JSONObject record) { From bb4690eee077a2c9162b91266dc8da15d0729703 Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 27 Oct 2026 14:55:12 -0700 Subject: [PATCH 25/26] feat(ppl): Report engine CPU in PPL task stats The engine measured its own threads privately for the Top N record, so the PPL task's resource_stats stayed at zero, and _tasks and Live Queries showed no CPU or memory for work done inside the engine. Drive the engine's thread scopes through core's TaskResourceTrackingService instead, and mark PPLQueryTask as supporting resource tracking. Core then registers the task when task_resource_tracking.enabled is on, records the engine threads in its resource_stats, and refreshes threads still running whenever tasks are listed. The Top N record reads the same totals. Scopes now run on every execute route, not only when Query Insights recording is on, since _tasks and Live Queries read the task either way. The background scan thread is scoped too, so its work is no longer missed. The record's latency is now taken when the response is sent. The record goes out once the last scope closes, and a lookahead search still running on the background pool after the query has answered would otherwise count toward it. Signed-off-by: Kishore Kumaar Natarajan --- docs/user/ppl/admin/settings.md | 2 + .../executor/ThreadResourceAccounting.java | 5 +- .../storage/scan/BackgroundSearchScanner.java | 6 +- .../scan/BackgroundSearchScannerTest.java | 34 ++++++ .../sql/plugin/transport/PPLQueryTask.java | 100 +++++++++--------- .../transport/TransportPPLQueryAction.java | 58 ++++++---- .../plugin/transport/PPLQueryTaskTest.java | 83 ++++++++++++--- 7 files changed, 199 insertions(+), 89 deletions(-) diff --git a/docs/user/ppl/admin/settings.md b/docs/user/ppl/admin/settings.md index cd34905c518..56aeffe46d7 100644 --- a/docs/user/ppl/admin/settings.md +++ b/docs/user/ppl/admin/settings.md @@ -490,6 +490,8 @@ This setting controls whether completed PPL queries are recorded into Query Insi Note on query text and user identity: the record stores the query as it was submitted, including any literal values in it, the same way Query Insights stores the source of a DSL search. When a security plugin is installed, the record also includes the requesting user (username, roles, backend roles). Who can see these in Top N results is governed by the Query Insights plugin's own RBAC filter, `search.insights.top_queries.filter_by_mode` — default `none` (no per-user filtering), `username` (a user sees only their own queries; admins see all), or `backend_roles`. Set it to `username` or `backend_roles` if other users' queries and identities should not be visible to every Top N reader. +Note on CPU and memory: the record reads them from the PPL task's resource stats (the same values `_tasks` shows), which OpenSearch records only while the cluster setting `task_resource_tracking.enabled` is on (the default). With it off, a PPL record reports 0 CPU and 0 memory. + Note: This setting is **registered by the Query Insights plugin** (it belongs to the `search.insights.top_queries.*` family), not by the PPL plugin. The PPL plugin reads its value to decide whether to report. It therefore only takes effect when the Query Insights plugin is installed; when Query Insights is absent the setting is unregistered and PPL reporting is a no-op. ### Example diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadResourceAccounting.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadResourceAccounting.java index a7ebb88f03d..5e9706fb6e9 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadResourceAccounting.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadResourceAccounting.java @@ -10,9 +10,8 @@ * *

Core's task resource tracking only brackets threads on its task-aware pools, which a plugin * pool cannot opt into, so the engine brackets its own threads instead. It does so through this - * interface rather than {@code Task.startThreadResourceTracking}, so nothing is written into the - * task's {@code resource_stats}: whether and how to measure is the task's decision, not the - * engine's. + * interface rather than calling core directly: whether and how a stretch of work is measured is the + * task's decision, not the engine's. */ public interface ThreadResourceAccounting { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScanner.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScanner.java index 2faf599d99e..c5ac951a48a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScanner.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScanner.java @@ -23,6 +23,7 @@ import org.opensearch.sql.monitor.profile.QueryProfiling; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.executor.OpenSearchQueryManager; +import org.opensearch.sql.opensearch.executor.ThreadResourceAccounting; import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.response.OpenSearchResponse; import org.opensearch.tasks.CancellableTask; @@ -129,7 +130,8 @@ public void startScanning(OpenSearchRequest request) { /** * Executes a search on the background thread with the worker thread's context restored. Restores * both the profiling context and the cancellable task so the request is tagged with the ppl task - * as its parent (see {@link #cancellableTask}). + * as its parent (see {@link #cancellableTask}). The thread's own work is accounted to the query + * too: core measures the search's tasks, but not this plugin pool's thread. */ private OpenSearchResponse backgroundSearch(ProfileContext ctx, OpenSearchRequest request) { boolean restored = false; @@ -137,7 +139,7 @@ private OpenSearchResponse backgroundSearch(ProfileContext ctx, OpenSearchReques OpenSearchQueryManager.setCancellableTask(cancellableTask); restored = true; } - try { + try (ThreadResourceAccounting.Scope ignored = ThreadResourceAccounting.enter(cancellableTask)) { return QueryProfiling.withCurrentContext(ctx, () -> client.search(request)); } finally { if (restored) { diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScannerTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScannerTest.java index f82938e519d..0cb92e6c09d 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScannerTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScannerTest.java @@ -15,11 +15,13 @@ import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Optional; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicReference; @@ -31,6 +33,7 @@ import org.opensearch.sql.executor.Warning; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.executor.OpenSearchQueryManager; +import org.opensearch.sql.opensearch.executor.ThreadResourceAccounting; import org.opensearch.sql.opensearch.request.OpenSearchQueryRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.response.OpenSearchResponse; @@ -226,6 +229,37 @@ void propagatesCancellableTaskToBackgroundSearchThread() { + " can tag the search request with the ppl task as parent"); } + @Test + void accountsBackgroundSearchThreadToTheQuerysTask() { + // Core measures the search's own tasks but not this plugin pool's thread, so the thread's + // work is measured for the query, from before the search starts until it returns. + CancellableTask task = + mock(CancellableTask.class, withSettings().extraInterfaces(ThreadResourceAccounting.class)); + Thread caller = Thread.currentThread(); + List events = new CopyOnWriteArrayList<>(); + when(((ThreadResourceAccounting) task).enterThread()) + .thenAnswer( + invocation -> { + events.add(Thread.currentThread() == caller ? "enter on caller" : "enter"); + return (ThreadResourceAccounting.Scope) () -> events.add("close"); + }); + OpenSearchQueryManager.setCancellableTask(task); + // Scanner captures the task at construction time (on this worker thread). + scanner = new BackgroundSearchScanner(client, 10, 10); + OpenSearchResponse response = mockResponse(false, false, 5); + when(client.search(request)) + .thenAnswer( + invocation -> { + events.add("search"); + return response; + }); + + scanner.startScanning(request); + scanner.fetchNextBatch(request); + + assertEquals(List.of("enter", "search", "close"), events); + } + @Test void clearsRestoredCancellableTaskAfterBackgroundSearch() { // The scanner must not leak the restored task onto the shared background pool thread. diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java index 90adc4ba45e..e5ec9b91b22 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/PPLQueryTask.java @@ -5,21 +5,16 @@ package org.opensearch.sql.plugin.transport; -import com.sun.management.ThreadMXBean; -import java.lang.management.ManagementFactory; import java.util.Map; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.atomic.LongAdder; import org.opensearch.core.tasks.TaskId; import org.opensearch.sql.opensearch.executor.ThreadResourceAccounting; import org.opensearch.tasks.CancellableTask; +import org.opensearch.tasks.TaskResourceTrackingService; public class PPLQueryTask extends CancellableTask implements ThreadResourceAccounting { - /** Null when the JVM can't report per-thread CPU and allocation; accounting is then a no-op. */ - private static final ThreadMXBean THREAD_MX_BEAN = resolveThreadMXBean(); - // The fields below are captured on the request thread and read later by the report listener // (which runs on a different thread), for the Query Insights record. @@ -32,6 +27,9 @@ public class PPLQueryTask extends CancellableTask implements ThreadResourceAccou /** Whether the query failed; recorded so a failed PPL query is flagged in Top N. */ private volatile boolean queryInsightsFailed = false; + /** Milliseconds from the task's start to its response; -1 until the response is sent. */ + private volatile long queryInsightsLatencyMillis = -1L; + /** Whether the statement is an {@code explain}, which plans but runs no search. */ private volatile boolean queryInsightsExplain = false; @@ -48,11 +46,11 @@ public class PPLQueryTask extends CancellableTask implements ThreadResourceAccou */ private volatile String queryInsightsParentMarker; - /** Whether this query's thread usage is measured; off unless Query Insights will report it. */ - private volatile boolean resourceAccountingEnabled = false; - - private final LongAdder cpuNanos = new LongAdder(); - private final LongAdder allocatedBytes = new LongAdder(); + /** + * Core's tracking service, which the engine's thread scopes report to; null (scopes are no-ops) + * on every route but execute: explain, analyze, async submission and the analytics engine. + */ + private volatile TaskResourceTrackingService resourceTracking; /** Threads with an open scope, so a nested scope on the same thread isn't counted twice. */ private final Set accountedThreads = ConcurrentHashMap.newKeySet(); @@ -115,72 +113,76 @@ public boolean isQueryInsightsFailed() { return queryInsightsFailed; } - public void setResourceAccountingEnabled(boolean enabled) { - this.resourceAccountingEnabled = enabled && THREAD_MX_BEAN != null; + public void setQueryInsightsLatencyMillis(long latencyMillis) { + this.queryInsightsLatencyMillis = latencyMillis; + } + + public long getQueryInsightsLatencyMillis() { + return queryInsightsLatencyMillis; } - /** CPU time used on the SQL plugin's threads; child searches are measured by core. */ + /** + * Has the engine's thread scopes report to {@code service}. Set on the execute routes, where the + * engine runs the query; other routes leave the scopes as no-ops. + */ + public void setResourceTracking(TaskResourceTrackingService service) { + this.resourceTracking = service; + } + + /** + * CPU time core tracked on this task: the engine's threads, plus any work core's task-aware pools + * ran for it. Child searches are tracked on their own tasks, so they are not included. + */ public long getCpuNanos() { - return cpuNanos.sum(); + return getTotalResourceStats().getCpuTimeInNanos(); } - /** Bytes allocated on the SQL plugin's threads; child searches are measured by core. */ + /** Bytes allocated by the threads core tracked on this task; see {@link #getCpuNanos()}. */ public long getAllocatedBytes() { - return allocatedBytes.sum(); + return getTotalResourceStats().getMemoryInBytes(); } /** - * Measures the current thread until the scope closes. The scope also holds the task's - * resource-tracking thread count, so the completion listener the Query Insights report hangs off - * fires only after every scope has recorded its usage. Nothing is written to the task's {@code - * resource_stats}, so {@code _tasks} output and {@code task_resource_tracking.enabled} are - * untouched. + * Tracked like a search task, so core's {@code startTracking} registers it when {@code + * task_resource_tracking.enabled} is on and the JVM can measure threads. Otherwise the tracking + * service ignores the engine's scopes, so the cluster setting decides whether anything is + * recorded. + */ + @Override + public boolean supportsResourceTracking() { + return true; + } + + /** + * Tracks the current thread against this task until the scope closes, through core's tracking + * service: the usage lands in the task's {@code resource_stats}, so {@code _tasks} and Live + * Queries show it while the query runs (list tasks refreshes threads still working). A tracked + * thread also holds the task's resource-tracking count, so the completion listener the Query + * Insights report hangs off fires only after every scope has recorded its usage. */ @Override public Scope enterThread() { - if (!resourceAccountingEnabled) { + final TaskResourceTrackingService service = resourceTracking; + if (service == null) { return Scope.NOOP; } final Thread thread = Thread.currentThread(); final long threadId = thread.threadId(); + // Core rejects a second start on a thread it is already tracking for this task. if (!accountedThreads.add(threadId)) { return Scope.NOOP; } - incrementResourceTrackingThreads(); - final long cpuStart = THREAD_MX_BEAN.getCurrentThreadCpuTime(); - final long allocatedStart = THREAD_MX_BEAN.getCurrentThreadAllocatedBytes(); + service.taskExecutionStartedOnThread(getId(), threadId); return () -> { assert Thread.currentThread() == thread : "scope closed on a different thread"; try { - addDelta(cpuNanos, cpuStart, THREAD_MX_BEAN.getCurrentThreadCpuTime()); - addDelta(allocatedBytes, allocatedStart, THREAD_MX_BEAN.getCurrentThreadAllocatedBytes()); + service.taskExecutionFinishedOnThread(getId(), threadId); } finally { accountedThreads.remove(threadId); - decrementResourceTrackingThreads(); } }; } - /** The bean returns -1 when the measurement is unsupported or disabled; record nothing then. */ - private static void addDelta(LongAdder total, long start, long end) { - if (start >= 0 && end >= start) { - total.add(end - start); - } - } - - private static ThreadMXBean resolveThreadMXBean() { - try { - if (ManagementFactory.getThreadMXBean() instanceof ThreadMXBean bean - && bean.isCurrentThreadCpuTimeSupported() - && bean.isThreadAllocatedMemorySupported()) { - return bean; - } - } catch (Exception e) { - // Fall through: accounting stays off. - } - return null; - } - @Override public boolean shouldCancelChildrenOnCancellation() { return true; 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 9f36629fe8a..b0bbb4f08db 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 @@ -108,6 +108,7 @@ public class TransportPPLQueryAction private final TransportService transportServiceRef; private final QueryJobService queryJobService; private final SecurityAdapter securityAdapter; + private final TaskResourceTrackingService taskResourceTrackingService; @Inject public TransportPPLQueryAction( @@ -121,13 +122,15 @@ public TransportPPLQueryAction( Tracer tracer, OpenSearchQueryJobService queryJobService, OpenSearchSecurityAdapter securityAdapter, - RemoteClusterServiceProvider remoteClusters) { + RemoteClusterServiceProvider remoteClusters, + TaskResourceTrackingService taskResourceTrackingService) { super(PPLQueryAction.NAME, transportService, actionFilters, TransportPPLQueryRequest::new); this.clientRef = client; this.clusterServiceRef = clusterService; this.transportServiceRef = transportService; this.queryJobService = queryJobService; this.securityAdapter = securityAdapter; + this.taskResourceTrackingService = taskResourceTrackingService; // The plugin's components are created before any TransportService exists; fill it here. remoteClusters.set(transportService.getRemoteClusterService()); @@ -248,18 +251,14 @@ private boolean isQueryInsightsRecordingEnabled() { } /** - * Whether to measure this query's thread usage. Follows core's {@code - * task_resource_tracking.enabled}: with it off, core doesn't measure the child searches either, - * so measuring only the PPL threads would produce a partial total. Treated as off if unreadable. + * Has the engine's thread scopes track this query's threads through core, so its CPU and memory + * show in {@code _tasks} (and Live Queries) while it runs and in its Top N record after. Core + * records nothing unless it registered the task, which it does only when {@code + * task_resource_tracking.enabled} is on, the same switch that governs the child searches. */ - private boolean isTaskResourceTrackingEnabled() { - try { - return clusterServiceRef - .getClusterSettings() - .get(TaskResourceTrackingService.TASK_RESOURCE_TRACKING_ENABLED); - } catch (Exception e) { - LOG.debug("Failed to read task_resource_tracking.enabled; not measuring PPL threads", e); - return false; + private void trackEngineThreads(PPLQueryTask pplQueryTask) { + if (pplQueryTask != null) { + pplQueryTask.setResourceTracking(taskResourceTrackingService); } } @@ -317,9 +316,13 @@ private void writeQueryInsightsRecord(PPLQueryTask pplQueryTask) { long cpuNanos = pplQueryTask.getCpuNanos(); long memoryBytes = pplQueryTask.getAllocatedBytes(); - // Wall-clock latency (monotonic); the profile total collapses to ~0 across the finish thread. - long latencyMillis = - Math.max(0L, (System.nanoTime() - pplQueryTask.getStartTimeNanos()) / 1_000_000L); + // Taken when the response was sent (see recordCompletion), not now: this runs once the last + // thread scope closes, which a lookahead search still running on the background pool can + // hold open well past the response. + long latencyMillis = pplQueryTask.getQueryInsightsLatencyMillis(); + if (latencyMillis < 0) { + latencyMillis = elapsedMillis(pplQueryTask); + } // Start time, not report time: Query Insights stamps a search record with the request's // start (SearchRequest#getOrCreateAbsoluteStartMillis). Stamping completion here would put a @@ -411,7 +414,7 @@ protected void doExecute( ActionListener clearingListener = wrapWithProfilingClear(tracedListener); // Read back later by writeQueryInsightsRecord, which runs after this listener completes. - clearingListener = markFailedOnFailure(clearingListener, pplQueryTask); + clearingListener = recordCompletion(clearingListener, pplQueryTask); try { // Route to analytics engine for non-Lucene (e.g., Parquet-backed) indices. @@ -487,14 +490,14 @@ protected void doExecute( // plain DSL search) with this query's marker. Child searches still see it: the executor // captures the thread context when the work is scheduled, which happens inside this block. try (ThreadContext.StoredContext ignored = threadContext.newStoredContext(true)) { - // Measure thread usage only on this route, where a record consumes it. Every other route - // leaves accounting off, so the engine's scopes are no-ops there. - pplQueryTask.setResourceAccountingEnabled(isTaskResourceTrackingEnabled()); + trackEngineThreads(pplQueryTask); stampQueryInsightsParentHeader(pplQueryTask); registerQueryInsightsReport(pplQueryTask); executePplQuery(pplService, transformedRequest, callerListener, queryInsightsSink); } } else { + // Tracked without Query Insights recording too: _tasks and Live Queries read the task. + trackEngineThreads(pplQueryTask); executePplQuery(pplService, transformedRequest, clearingListener, queryInsightsSink); } } catch (Exception e) { @@ -694,8 +697,11 @@ private void executePplQuery( queryInsightsSink); } - /** Returns {@code delegate} unchanged when there is no PPL task to mark. */ - private ActionListener markFailedOnFailure( + /** + * Records on the task how the query ended and its latency as the response goes out, for the Query + * Insights record written later. Returns {@code delegate} unchanged when there is no PPL task. + */ + private ActionListener recordCompletion( ActionListener delegate, PPLQueryTask pplQueryTask) { if (pplQueryTask == null) { return delegate; @@ -703,17 +709,27 @@ private ActionListener markFailedOnFailure( return new ActionListener<>() { @Override public void onResponse(TransportPPLQueryResponse response) { + pplQueryTask.setQueryInsightsLatencyMillis(elapsedMillis(pplQueryTask)); delegate.onResponse(response); } @Override public void onFailure(Exception e) { + pplQueryTask.setQueryInsightsLatencyMillis(elapsedMillis(pplQueryTask)); pplQueryTask.setQueryInsightsFailed(true); delegate.onFailure(e); } }; } + /** + * Wall-clock (monotonic) time since the task started. The profile total is not used: it collapses + * to ~0 across the thread that finishes the query. + */ + private static long elapsedMillis(PPLQueryTask pplQueryTask) { + return Math.max(0L, (System.nanoTime() - pplQueryTask.getStartTimeNanos()) / 1_000_000L); + } + private ActionListener wrapWithProfilingClear( ActionListener delegate) { return new ActionListener<>() { diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java index 626eb621bbe..a8a15969394 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/PPLQueryTaskTest.java @@ -6,14 +6,21 @@ package org.opensearch.sql.plugin.transport; import static org.junit.Assert.*; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; import org.junit.Test; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.action.NotifyOnceListener; import org.opensearch.core.tasks.TaskId; import org.opensearch.sql.opensearch.executor.ThreadResourceAccounting; import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskResourceTrackingService; +import org.opensearch.threadpool.ThreadPool; public class PPLQueryTaskTest { @@ -77,41 +84,89 @@ private PPLQueryTask newTask() { Map.of()); } + /** Core's tracking service as a node builds it, with {@code task_resource_tracking.enabled}. */ + private static TaskResourceTrackingService trackingService(boolean enabled) { + Settings settings = + Settings.builder() + .put(TaskResourceTrackingService.TASK_RESOURCE_TRACKING_ENABLED.getKey(), enabled) + .build(); + ThreadPool threadPool = mock(ThreadPool.class); + when(threadPool.getThreadContext()).thenReturn(new ThreadContext(Settings.EMPTY)); + return new TaskResourceTrackingService( + settings, + new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS), + threadPool); + } + + /** A task as the execute route leaves it: registered by core, scopes wired to the service. */ + private PPLQueryTask trackedTask(TaskResourceTrackingService service) { + PPLQueryTask task = newTask(); + // What TransportAction.execute does before the action runs; a no-op when tracking is off. + service.startTracking(task).close(); + task.setResourceTracking(service); + return task; + } + + @Test + public void testTrackedLikeASearchTask() { + assertTrue(newTask().supportsResourceTracking()); + } + @Test - public void testAccountingOffByDefault() { + public void testScopesAreNoOpsOnRoutesThatRunNoQuery() { PPLQueryTask task = newTask(); assertSame(ThreadResourceAccounting.Scope.NOOP, task.enterThread()); - // Never touches core's tracking, so _tasks resource_stats stays empty. - assertFalse(task.supportsResourceTracking()); assertTrue(task.getResourceStats().isEmpty()); } @Test - public void testAccountingRecordsUsageWithoutTouchingResourceStats() { - PPLQueryTask task = newTask(); - task.setResourceAccountingEnabled(true); + public void testScopeRecordsThreadUsageInResourceStats() { + PPLQueryTask task = trackedTask(trackingService(true)); + try (ThreadResourceAccounting.Scope ignored = task.enterThread()) { + burnCpu(); + } + // The task's resource_stats, which _tasks and Live Queries read; the record reads the same. + assertTrue(task.getTotalResourceStats().getCpuTimeInNanos() > 0); + assertTrue(task.getTotalResourceStats().getMemoryInBytes() > 0); + assertEquals(task.getTotalResourceStats().getCpuTimeInNanos(), task.getCpuNanos()); + assertEquals(task.getTotalResourceStats().getMemoryInBytes(), task.getAllocatedBytes()); + } + + @Test + public void testOpenScopeShowsOnRefresh() { + // List tasks refreshes threads still working, so a query's usage shows before its scope ends. + TaskResourceTrackingService service = trackingService(true); + PPLQueryTask task = trackedTask(service); + try (ThreadResourceAccounting.Scope ignored = task.enterThread()) { + burnCpu(); + service.refreshResourceStats(task); + assertTrue(task.getTotalResourceStats().getCpuTimeInNanos() > 0); + } + } + + @Test + public void testNothingRecordedWhenTaskResourceTrackingIsOff() { + PPLQueryTask task = trackedTask(trackingService(false)); try (ThreadResourceAccounting.Scope ignored = task.enterThread()) { burnCpu(); } - assertTrue(task.getCpuNanos() > 0); - assertTrue(task.getAllocatedBytes() > 0); - assertFalse(task.supportsResourceTracking()); assertTrue(task.getResourceStats().isEmpty()); + assertEquals(0L, task.getCpuNanos()); } @Test - public void testNestedScopeOnSameThreadIsNotCountedTwice() { - PPLQueryTask task = newTask(); - task.setResourceAccountingEnabled(true); + public void testNestedScopeOnSameThreadIsNotTrackedTwice() { + PPLQueryTask task = trackedTask(trackingService(true)); try (ThreadResourceAccounting.Scope outer = task.enterThread()) { + // Core rejects a second start on a thread it already tracks for the task. assertSame(ThreadResourceAccounting.Scope.NOOP, task.enterThread()); } + assertEquals(1, task.getResourceStats().get(Thread.currentThread().threadId()).size()); } @Test public void testReportWaitsForOpenScopes() { - PPLQueryTask task = newTask(); - task.setResourceAccountingEnabled(true); + PPLQueryTask task = trackedTask(trackingService(true)); AtomicInteger fired = new AtomicInteger(); assertTrue( task.addResourceTrackingCompletionListener( From 9394b378f40036bf7b1395b4d2c86045344a469a Mon Sep 17 00:00:00 2001 From: Kishore Kumaar Natarajan Date: Tue, 27 Oct 2026 15:04:31 -0700 Subject: [PATCH 26/26] test(ppl): Cover Query Insights through doExecute QueryInsightsParentHeaderScopeTest re-implemented the thread-context steps doExecute takes and tested that copy, so it would keep passing if doExecute changed. Replace it with a test that calls doExecute with only the engine stubbed and core's task resource tracking real. The test decodes the record sent to Query Insights and checks each field: the query text as submitted, the parent marker, node, start timestamp, user, CPU and memory equal to the task's totals, and the failed flag. It also checks that the marker reaches the engine's thread but not the calling thread or the caller's callback, that a nested query keeps the outer marker and sends nothing, that nothing is sent with recording off or Query Insights absent while CPU is still tracked, and that latency stops at the response when a lookahead search keeps running after it. Signed-off-by: Kishore Kumaar Natarajan --- .../QueryInsightsParentHeaderScopeTest.java | 88 ---- ...nsportPPLQueryActionQueryInsightsTest.java | 485 ++++++++++++++++++ 2 files changed, 485 insertions(+), 88 deletions(-) delete mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsParentHeaderScopeTest.java create mode 100644 plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportPPLQueryActionQueryInsightsTest.java diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsParentHeaderScopeTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsParentHeaderScopeTest.java deleted file mode 100644 index 6d514c3d2ca..00000000000 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsParentHeaderScopeTest.java +++ /dev/null @@ -1,88 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.plugin.transport; - -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertNull; - -import java.util.concurrent.atomic.AtomicReference; -import org.junit.Test; -import org.opensearch.action.support.ContextPreservingActionListener; -import org.opensearch.common.settings.Settings; -import org.opensearch.common.util.concurrent.ThreadContext; -import org.opensearch.core.action.ActionListener; - -/** - * Pins the thread-context contract {@code TransportPPLQueryAction} relies on when it stamps the - * Query Insights parent header: the header must reach the searches scheduled during execution, but - * must not outlive the call on the calling thread or reappear in the caller's callback. Without - * that, a caller issuing several PPL queries through {@code NodeClient} would tag later queries (or - * a plain DSL search) with the first query's marker. - */ -public class QueryInsightsParentHeaderScopeTest { - - private static final String MARKER = "PPL:node-1:42"; - - @Test - public void headerReachesScheduledWorkButDoesNotLeak() { - ThreadContext threadContext = new ThreadContext(Settings.EMPTY); - AtomicReference seenByScheduledWork = new AtomicReference<>("unset"); - AtomicReference seenByCallback = new AtomicReference<>("unset"); - - ActionListener caller = - ActionListener.wrap( - response -> - seenByCallback.set(threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER)), - e -> {}); - // Captured before stamping, mirroring doExecute. - ActionListener callerListener = - ContextPreservingActionListener.wrapPreservingContext(caller, threadContext); - - Runnable scheduledWork; - try (ThreadContext.StoredContext ignored = threadContext.newStoredContext(true)) { - threadContext.putHeader(QueryInsightsMarker.PARENT_HEADER, MARKER); - // Executors capture the thread context when work is scheduled; this is what carries the - // header to the sql-worker and background-scan threads. - scheduledWork = - threadContext.preserveContext( - () -> - seenByScheduledWork.set( - threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER))); - } - - assertNull( - "header must not outlive the call on the calling thread", - threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER)); - - scheduledWork.run(); - assertEquals( - "searches scheduled during execution must still see the header", - MARKER, - seenByScheduledWork.get()); - - callerListener.onResponse("ok"); - assertNull("the caller's callback must not inherit the header", seenByCallback.get()); - } - - @Test - public void existingCallerHeaderIsPreservedAndRestored() { - ThreadContext threadContext = new ThreadContext(Settings.EMPTY); - threadContext.putHeader(QueryInsightsMarker.PARENT_HEADER, "caller-supplied"); - - try (ThreadContext.StoredContext ignored = threadContext.newStoredContext(true)) { - // doExecute only stamps when absent, so a caller that already carries a marker keeps it. - if (threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER) == null) { - threadContext.putHeader(QueryInsightsMarker.PARENT_HEADER, MARKER); - } - assertEquals("caller-supplied", threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER)); - } - - assertEquals( - "the caller's own header must survive the scoped stamp", - "caller-supplied", - threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER)); - } -} diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportPPLQueryActionQueryInsightsTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportPPLQueryActionQueryInsightsTest.java new file mode 100644 index 00000000000..6fe4ca4192b --- /dev/null +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportPPLQueryActionQueryInsightsTest.java @@ -0,0 +1,485 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.CALLS_REAL_METHODS; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import java.util.function.Supplier; +import org.json.JSONObject; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.inject.Injector; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.action.NotifyOnceListener; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.sql.common.response.ResponseListener; +import org.opensearch.sql.executor.ExecutionEngine.QueryResponse; +import org.opensearch.sql.executor.ExecutionEngine.Schema; +import org.opensearch.sql.executor.pagination.Cursor; +import org.opensearch.sql.legacy.metrics.BasicCounter; +import org.opensearch.sql.legacy.metrics.MetricName; +import org.opensearch.sql.legacy.metrics.Metrics; +import org.opensearch.sql.legacy.metrics.NumericMetric; +import org.opensearch.sql.opensearch.executor.OpenSearchQueryManager; +import org.opensearch.sql.opensearch.executor.ThreadResourceAccounting; +import org.opensearch.sql.ppl.PPLService; +import org.opensearch.sql.ppl.QueryInsightsMetadata; +import org.opensearch.sql.ppl.domain.PPLQueryRequest; +import org.opensearch.tasks.CancellableTask; +import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskResourceTrackingService; +import org.opensearch.telemetry.tracing.noop.NoopTracer; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.BytesTransportRequest; +import org.opensearch.transport.TransportService; +import org.opensearch.transport.client.node.NodeClient; + +/** + * Drives {@link TransportPPLQueryAction#doExecute} with only the engine stubbed, to pin what + * reaches Query Insights: the record and each of its fields, the parent marker the child searches + * see, and that the marker stays off the caller's thread and callback. Core's task resource + * tracking is real, so the record's CPU, memory and timing come through the same path as on a node. + */ +public class TransportPPLQueryActionQueryInsightsTest { + + private static final String QUERY = "source=accounts | where age > 30 | fields firstname"; + private static final String ANONYMIZED = "source=accounts | where age > *** | fields firstname"; + private static final String USER_INFO = "admin||all_access|__user__"; + private static final long TASK_ID = 42L; + private static final String MARKER = "PPL:node-1:42"; + private static final long LOOKAHEAD_HOLD_MILLIS = 500L; + + /** Registered by the Query Insights plugin, so absent when it is not installed. */ + private static final Setting PPL_RECORDING = + Setting.boolSetting( + TransportPPLQueryAction.QUERY_INSIGHTS_PPL_ENABLED_KEY, + false, + Setting.Property.NodeScope, + Setting.Property.Dynamic); + + private final ThreadContext threadContext = new ThreadContext(Settings.EMPTY); + private final PPLService pplService = mock(PPLService.class); + private final ClusterService clusterService = mock(ClusterService.class); + private final TransportService transportService = mock(TransportService.class); + private final DiscoveryNode localNode = mock(DiscoveryNode.class); + private final ExecutorService worker = Executors.newSingleThreadExecutor(); + private final ExecutorService backgroundPool = Executors.newSingleThreadExecutor(); + private final List registeredMetrics = new ArrayList<>(); + + private final List reports = new CopyOnWriteArrayList<>(); + private final CountDownLatch reported = new CountDownLatch(1); + private volatile long reportedAtNanos; + + private final CompletableFuture outcome = new CompletableFuture<>(); + private volatile long respondedAtNanos; + private final AtomicReference headerSeenByCaller = new AtomicReference<>("unset"); + private final AtomicReference headerSeenByEngine = new AtomicReference<>("unset"); + + private TaskResourceTrackingService tracking; + private TransportPPLQueryAction action; + + @Before + public void setUp() throws Exception { + ThreadPool threadPool = mock(ThreadPool.class); + when(threadPool.getThreadContext()).thenReturn(threadContext); + NodeClient client = mock(NodeClient.class); + when(client.threadPool()).thenReturn(threadPool); + when(transportService.getThreadPool()).thenReturn(threadPool); + when(localNode.getId()).thenReturn("node-1"); + when(clusterService.localNode()).thenReturn(localNode); + doAnswer( + invocation -> { + reportedAtNanos = System.nanoTime(); + reports.add(invocation.getArgument(2)); + reported.countDown(); + return null; + }) + .when(transportService) + .sendRequest(eq(localNode), eq(QueryInsightsReporter.ACTION_NAME), any(), any()); + + Settings trackingOn = + Settings.builder() + .put(TaskResourceTrackingService.TASK_RESOURCE_TRACKING_ENABLED.getKey(), true) + .build(); + tracking = + new TaskResourceTrackingService( + trackingOn, + new ClusterSettings(trackingOn, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS), + threadPool); + + Injector injector = mock(Injector.class); + when(injector.getInstance(PPLService.class)).thenReturn(pplService); + + // Skips the constructor, which builds the plugin's Guice graph; doExecute reads only these. + action = mock(TransportPPLQueryAction.class, CALLS_REAL_METHODS); + setField("pplEnabled", (Supplier) () -> true); + setField("injector", injector); + setField("tracer", NoopTracer.INSTANCE); + setField("clientRef", client); + setField("clusterServiceRef", clusterService); + setField("transportServiceRef", transportService); + setField("taskResourceTrackingService", tracking); + + registerMetricIfAbsent(MetricName.PPL_REQ_TOTAL); + registerMetricIfAbsent(MetricName.PPL_REQ_COUNT_TOTAL); + + // What the security plugin leaves on the request thread for an authenticated user. + threadContext.putTransient("_opendistro_security_user_info", USER_INFO); + } + + @After + public void tearDown() { + worker.shutdownNow(); + backgroundPool.shutdownNow(); + registeredMetrics.forEach(Metrics.getInstance()::unregisterMetric); + OpenSearchQueryManager.clearCancellableTask(); + } + + @Test + public void reportsTheCompletedQueryWithItsCostAndUser() throws Exception { + recording(true); + engine( + (listener, queryTask) -> { + burnCpu(); + listener.onResponse(emptyResponse()); + }); + + PPLQueryTask task = execute(); + + assertTrue("no record sent", reported.await(10, TimeUnit.SECONDS)); + assertTrue(outcome.get(10, TimeUnit.SECONDS) instanceof TransportPPLQueryResponse); + Record record = onlyRecord(); + assertEquals("PPL", record.source()); + assertEquals(MARKER, record.parentMarker()); + assertEquals("node-1", record.nodeId()); + assertEquals("the query as submitted, literals included", QUERY, record.queryText()); + assertEquals( + "the query's start, like a search record", task.getStartTime(), record.timestamp()); + assertTrue("engine CPU", record.cpuNanos() > 0); + assertTrue("engine memory", record.memoryBytes() > 0); + assertEquals("the totals _tasks shows", task.getCpuNanos(), record.cpuNanos()); + assertEquals(task.getAllocatedBytes(), record.memoryBytes()); + assertEquals(USER_INFO, record.userInfo()); + assertFalse(record.failed()); + } + + @Test + public void childSearchesSeeTheMarkerButTheCallerDoesNot() throws Exception { + recording(true); + engine((listener, queryTask) -> listener.onResponse(emptyResponse())); + + execute(); + + assertNull( + "marker left on the calling thread", + threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER)); + outcome.get(10, TimeUnit.SECONDS); + assertEquals( + "the engine's thread, which issues the child searches, carries the marker", + MARKER, + headerSeenByEngine.get()); + assertNull("marker leaked into the caller's callback", headerSeenByCaller.get()); + } + + @Test + public void latencyIsTakenWhenTheResponseIsSent() throws Exception { + recording(true); + // A lookahead search the scanner started before the query answered, still running after. + engine( + (listener, queryTask) -> { + CountDownLatch scoped = new CountDownLatch(1); + CountDownLatch answered = new CountDownLatch(1); + backgroundPool.execute( + () -> { + try (ThreadResourceAccounting.Scope ignored = + ThreadResourceAccounting.enter(queryTask)) { + scoped.countDown(); + answered.await(); + Thread.sleep(LOOKAHEAD_HOLD_MILLIS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + scoped.await(); + listener.onResponse(emptyResponse()); + answered.countDown(); + }); + + PPLQueryTask task = execute(); + + assertTrue("no record sent", reported.await(10, TimeUnit.SECONDS)); + assertTrue( + "the record waits for the lookahead's scope, so it goes out after the response", + reportedAtNanos - respondedAtNanos + >= TimeUnit.MILLISECONDS.toNanos(LOOKAHEAD_HOLD_MILLIS / 2)); + long respondedAtMillis = (respondedAtNanos - task.getStartTimeNanos()) / 1_000_000L; + long latencyMillis = onlyRecord().latencyMillis(); + assertTrue( + "latency " + + latencyMillis + + "ms includes time after the response, sent at " + + respondedAtMillis + + "ms", + latencyMillis <= respondedAtMillis); + } + + @Test + public void failedQueryIsReportedAsFailed() throws Exception { + recording(true); + engine( + (listener, queryTask) -> { + burnCpu(); + listener.onFailure(new IllegalStateException("boom")); + }); + + execute(); + + assertTrue("no record sent", reported.await(10, TimeUnit.SECONDS)); + assertTrue(outcome.get(10, TimeUnit.SECONDS) instanceof IllegalStateException); + Record record = onlyRecord(); + assertTrue("failed flag", record.failed()); + assertEquals(QUERY, record.queryText()); + assertTrue("engine CPU", record.cpuNanos() > 0); + } + + @Test + public void recordingOffSendsNothingButStillTracksTheEngine() throws Exception { + recording(false); + engine( + (listener, queryTask) -> { + burnCpu(); + listener.onResponse(emptyResponse()); + }); + + PPLQueryTask task = execute(); + outcome.get(10, TimeUnit.SECONDS); + awaitEngine(); + + assertTrue("record sent with recording off", reports.isEmpty()); + assertNull("marker stamped with recording off", headerSeenByEngine.get()); + assertTrue("_tasks and Live Queries still read the engine's CPU", task.getCpuNanos() > 0); + } + + @Test + public void queryInsightsNotInstalledSendsNothing() throws Exception { + recording(null); + engine((listener, queryTask) -> listener.onResponse(emptyResponse())); + + execute(); + outcome.get(10, TimeUnit.SECONDS); + awaitEngine(); + + assertTrue(reports.isEmpty()); + assertNull(headerSeenByEngine.get()); + } + + @Test + public void nestedQueryKeepsTheOuterMarkerAndSendsNoRecord() throws Exception { + recording(true); + // An outer query's marker is already in the context, so the child searches roll up into it. + threadContext.putHeader(QueryInsightsMarker.PARENT_HEADER, "PPL:node-0:7"); + engine((listener, queryTask) -> listener.onResponse(emptyResponse())); + + execute(); + outcome.get(10, TimeUnit.SECONDS); + awaitEngine(); + + assertEquals("PPL:node-0:7", headerSeenByEngine.get()); + assertEquals( + "the outer marker survives on the calling thread", + "PPL:node-0:7", + threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER)); + assertTrue("a nested query must not report", reports.isEmpty()); + } + + /** What the engine does with a parsed query, on its worker thread, inside the query's scope. */ + private interface Engine { + void run(ResponseListener listener, CancellableTask queryTask) throws Exception; + } + + private void engine(Engine engine) { + doAnswer( + invocation -> { + ResponseListener listener = invocation.getArgument(1); + Consumer sink = invocation.getArgument(3); + // What PPLService does once the query parses. + sink.accept(new QueryInsightsMetadata(ANONYMIZED, false)); + // What OpenSearchQueryManager.submit does: hand the query to a worker thread, which + // carries the caller's thread context and runs inside a scope for the query's task. + CancellableTask queryTask = OpenSearchQueryManager.getCancellableTask(); + worker.execute( + threadContext.preserveContext( + () -> { + headerSeenByEngine.set( + threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER)); + try (ThreadResourceAccounting.Scope ignored = + ThreadResourceAccounting.enter(queryTask)) { + engine.run(listener, queryTask); + } catch (Exception e) { + listener.onFailure(e); + } + })); + return null; + }) + .when(pplService) + .execute(any(), any(), any(), any()); + } + + /** Runs the query the way TransportAction.execute does: register, execute, unregister. */ + private PPLQueryTask execute() { + TransportPPLQueryRequest request = + new TransportPPLQueryRequest( + new PPLQueryRequest(QUERY, new JSONObject().put("query", QUERY), "/_plugins/_ppl", "")); + PPLQueryTask task = + request.createTask( + TASK_ID, "transport", PPLQueryAction.NAME, TaskId.EMPTY_TASK_ID, Map.of()); + // TaskManager.register: core tracks the task until its last thread is done. The context + // startTracking stamps is closed at once, so the checks see only what doExecute leaves behind. + tracking.startTracking(task).close(); + task.addResourceTrackingCompletionListener( + new NotifyOnceListener<>() { + @Override + protected void innerOnResponse(Task completed) { + tracking.stopTracking(completed); + } + + @Override + protected void innerOnFailure(Exception e) {} + }); + action.doExecute(task, request, caller(task)); + return task; + } + + /** The listener TransportAction.execute puts in front of the caller's: unregister, answer. */ + private ActionListener caller(PPLQueryTask task) { + return new ActionListener<>() { + @Override + public void onResponse(TransportPPLQueryResponse response) { + finish(task, response); + } + + @Override + public void onFailure(Exception e) { + finish(task, e); + } + }; + } + + private void finish(PPLQueryTask task, Object result) { + // TaskManager.unregister releases the task's own thread. + task.decrementResourceTrackingThreads(); + respondedAtNanos = System.nanoTime(); + headerSeenByCaller.set(threadContext.getHeader(QueryInsightsMarker.PARENT_HEADER)); + outcome.complete(result); + } + + private void recording(Boolean enabled) { + ClusterSettings settings = + enabled == null + ? new ClusterSettings(Settings.EMPTY, Set.of()) + : new ClusterSettings( + Settings.builder().put(PPL_RECORDING.getKey(), enabled).build(), + Set.>of(PPL_RECORDING)); + when(clusterService.getClusterSettings()).thenReturn(settings); + } + + /** Waits for the engine's scope to close, after which no record can still be on its way. */ + private void awaitEngine() throws InterruptedException { + worker.shutdown(); + assertTrue(worker.awaitTermination(10, TimeUnit.SECONDS)); + } + + /** A record's fields, in the order {@link QueryInsightsReporter} writes them. */ + private record Record( + String source, + String parentMarker, + String nodeId, + String queryText, + long timestamp, + long latencyMillis, + long cpuNanos, + long memoryBytes, + String userInfo, + boolean failed) {} + + private Record onlyRecord() throws IOException { + assertEquals("records sent", 1, reports.size()); + try (StreamInput in = reports.get(0).bytes().streamInput()) { + assertEquals(QueryInsightsReporter.FORMAT_VERSION, in.readVInt()); + Record record = + new Record( + in.readString(), + in.readString(), + in.readString(), + in.readString(), + in.readVLong(), + in.readVLong(), + in.readVLong(), + in.readVLong(), + in.readString(), + in.readBoolean()); + assertEquals("trailing bytes", 0, in.available()); + return record; + } + } + + private static QueryResponse emptyResponse() { + return new QueryResponse(new Schema(List.of()), List.of(), Cursor.None); + } + + private static void burnCpu() { + long x = 0; + for (int i = 0; i < 1_000_000; i++) { + x += Integer.toString(i).hashCode(); + } + assertNotEquals(42, x); + } + + private void registerMetricIfAbsent(MetricName name) { + if (Metrics.getInstance().getMetric(name.getName()) == null) { + Metrics.getInstance().registerMetric(new NumericMetric<>(name.getName(), new BasicCounter())); + registeredMetrics.add(name.getName()); + } + } + + private void setField(String name, Object value) throws Exception { + Field field = TransportPPLQueryAction.class.getDeclaredField(name); + field.setAccessible(true); + field.set(action, value); + } +}