diff --git a/docs/user/ppl/admin/settings.md b/docs/user/ppl/admin/settings.md index b1d5eef17fa..56aeffe46d7 100644 --- a/docs/user/ppl/admin/settings.md +++ b/docs/user/ppl/admin/settings.md @@ -472,4 +472,36 @@ Expected output: "transient": {} } ``` - \ No newline at end of file + + +## search.insights.top_queries.ppl.enabled + +### Version + +3.10 + +### 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 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 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 + +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/build.gradle b/integ-test/build.gradle index 2d6b5278fab..10ae36267b4 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" @@ -548,6 +567,23 @@ 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" + } + // 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' @@ -803,6 +839,58 @@ 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' + } +} + +// 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 @@ -832,6 +920,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' @@ -966,6 +1056,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' @@ -1317,6 +1410,10 @@ task integTestRemote(type: RestIntegTestTask) { excludeTestsMatching "org.opensearch.sql.bwc.*IT" } + // 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 // 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..4614ad7649e --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/ppl/QueryInsightsRecordingIT.java @@ -0,0 +1,206 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ppl; + +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.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Predicate; +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 + * 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 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 { + + /** 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 + QueryInsightsRecords.assertSupportsPpl(client()); + QueryInsightsRecords.enableRecording(client()); + } + + @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")); + // 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+")); + 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 query parsed, so it is still recorded. + 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 is not reported. + assertNoPplRecord("qi_it_syntax"); + } + + @Test + public void nothingIsRecordedWhenTheGateIsOff() throws Exception { + indexWithDocs("qi_it_gated"); + QueryInsightsRecords.putPersistentSettings(client(), Map.of(PPL_ENABLED, "false")); + try { + executeQuery("source=qi_it_gated | where v > 1 | fields v"); + } finally { + QueryInsightsRecords.putPersistentSettings(client(), Map.of(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"); + + List children = awaitChildren(client(), 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 record names no indices; its searches carry the join's two sides. + assertEquals(Set.of("qi_it_join_left", "qi_it_join_right"), childIndices); + } + + // --------------------------------------------------------------------------------------------- + + 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. */ + 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(pplRecordFor(index), "PPL record for " + 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 { + return QueryInsightsRecords.awaitRecord( + client(), RequestOptions.DEFAULT, CLASS_START_MILLIS, match, what); + } + + 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/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/QueryInsightsRecordingSecurityIT.java b/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsRecordingSecurityIT.java new file mode 100644 index 00000000000..58f4e74a512 --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsRecordingSecurityIT.java @@ -0,0 +1,154 @@ +/* + * 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 the index its query + * text reads. + */ +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/QueryInsightsReportingSecurityIT.java b/integ-test/src/test/java/org/opensearch/sql/security/QueryInsightsReportingSecurityIT.java new file mode 100644 index 00000000000..5ea3f66479a --- /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; + +/** + * Runs PPL as a low-privilege user on the FGAC (security-enabled) cluster with the Query Insights + * reporting code present. + * + *

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. + * + *

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 { + + 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 { + // 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), + 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/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..118d5b3fed9 --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/util/QueryInsightsRecords.java @@ -0,0 +1,222 @@ +/* + * 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 java.util.regex.Pattern; +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. + * + *

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 { + + 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(); + } + + /** The PPL record of a query that reads {@code index} with {@code source=}. */ + public static Predicate pplRecordFor(String 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) { + 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; + } +} 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..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 @@ -81,6 +81,11 @@ private void schedule( setCancellableTask(cancelTask); + // 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(); timeoutTask.cancel(); @@ -98,6 +103,7 @@ private void schedule( throw e; } finally { + accounting.close(); clearCancellableTask(); } }); 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..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,6 +91,7 @@ private void dispatchInternal( ThreadPool.Names.GENERIC); Cancellable cancelPoller = scheduleCancellationPoller(cancellableTask, executionThread); Hook.Closeable hookHandle = null; + ThreadResourceAccounting.Scope accounting = ThreadResourceAccounting.Scope.NOOP; try { // Restore state from caller thread ThreadContext.putAll(ctx); @@ -106,6 +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 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); @@ -119,6 +122,7 @@ private void dispatchInternal( if (hookHandle != null) { hookHandle.close(); } + 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..5e9706fb6e9 --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/ThreadResourceAccounting.java @@ -0,0 +1,38 @@ +/* + * 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 calling core directly: whether and how a stretch of work is measured 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/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/executor/OpenSearchQueryManagerTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchQueryManagerTest.java index 1463cf48fff..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 @@ -5,18 +5,24 @@ package org.opensearch.sql.opensearch.executor; +import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import java.util.Collections; 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; 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.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; @@ -26,6 +32,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; @@ -76,4 +83,101 @@ public void execute() { assertTrue(isRun.get()); } + + @AfterEach + public void clearTask() { + OpenSearchQueryManager.clearCancellableTask(); + } + + @Test + public void accountingScopeSpansExecution() { + AccountingTask task = new AccountingTask(); + OpenSearchQueryManager.setCancellableTask(task); + + runSubmit(() -> task.openDuringExecute.set(task.openScopes.get() == 1), null); + + 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 accountingScopeClosesWhenExecutionThrows() { + AccountingTask task = new AccountingTask(); + OpenSearchQueryManager.setCancellableTask(task); + RuntimeException boom = new RuntimeException("execute failed"); + + assertThrows(RuntimeException.class, () -> runSubmit(null, boom)); + + assertEquals(1, task.entered.get()); + assertEquals(0, task.openScopes.get()); + } + + @Test + public void propagatesExecutionException() { + RuntimeException boom = new RuntimeException("execute failed"); + RuntimeException thrown = assertThrows(RuntimeException.class, () -> runSubmit(null, boom)); + assertEquals(boom, thrown); + } + + /** 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); + 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() { + if (onExecute != null) { + onExecute.run(); + } + if (toThrow != null) { + throw toThrow; + } + } + }; + + doAnswer( + invocation -> { + if ("sql-worker".equals(invocation.getArgument(2))) { + invocation.getArgument(0).run(); + } + return mockScheduledTask; + }) + .when(threadPool) + .schedule(any(), any(), any()); + + new OpenSearchQueryManager(nodeClient, settings).submit(queryPlan); + } + + /** 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(); + + AccountingTask() { + super(1L, "ppl", "action", "desc", TaskId.EMPTY_TASK_ID, Collections.emptyMap()); + } + + @Override + public Scope enterThread() { + entered.incrementAndGet(); + openScopes.incrementAndGet(); + return openScopes::decrementAndGet; + } + + @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..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 @@ -16,6 +16,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.sql.opensearch.executor.OpenSearchQueryManager.SQL_COMPLEX_WORKER_THREAD_POOL_NAME; import java.util.List; @@ -368,6 +369,74 @@ void cancellableTaskAvailableDuringExecution() { assertEquals(mockTask, taskDuringExecution.get()); } + @Test + 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 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 -> { + 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); + + verify((ThreadResourceAccounting) accountingTask).enterThread(); + verify(scope).close(); + } + + @Test + 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 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 -> { + 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"); + }); + + // 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()); + } + private static RelNode createMockNode(RelNode... children) { RelNode node = mock(RelNode.class); List childList = List.of(children); 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/SQLPlugin.java b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java index 53b4d47c02f..2d3e735f2e6 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java @@ -121,6 +121,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; @@ -365,6 +366,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..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 @@ -6,10 +6,54 @@ package org.opensearch.sql.plugin.transport; import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; 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 { +public class PPLQueryTask extends CancellableTask implements ThreadResourceAccounting { + + // 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; + + /** 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; + + /** 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; + + /** + * 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; + + /** + * 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; + + /** + * 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(); public PPLQueryTask( long id, @@ -21,6 +65,124 @@ public PPLQueryTask( super(id, type, action, description, parentTaskId, headers); } + public void setQueryInsightsUserInfo(String userInfo) { + this.queryInsightsUserInfo = userInfo; + } + + public String getQueryInsightsUserInfo() { + return queryInsightsUserInfo; + } + + public void setQueryInsightsQueryText(String queryText) { + this.queryInsightsQueryText = queryText; + } + + public String getQueryInsightsQueryText() { + return queryInsightsQueryText; + } + + public void setQueryInsightsExplain(boolean explain) { + this.queryInsightsExplain = explain; + } + + public boolean isQueryInsightsExplain() { + return queryInsightsExplain; + } + + public void setQueryInsightsNested(boolean nested) { + this.queryInsightsNested = nested; + } + + 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; + } + + public boolean isQueryInsightsFailed() { + return queryInsightsFailed; + } + + public void setQueryInsightsLatencyMillis(long latencyMillis) { + this.queryInsightsLatencyMillis = latencyMillis; + } + + public long getQueryInsightsLatencyMillis() { + return queryInsightsLatencyMillis; + } + + /** + * 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 getTotalResourceStats().getCpuTimeInNanos(); + } + + /** Bytes allocated by the threads core tracked on this task; see {@link #getCpuNanos()}. */ + public long getAllocatedBytes() { + return getTotalResourceStats().getMemoryInBytes(); + } + + /** + * 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() { + 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; + } + service.taskExecutionStartedOnThread(getId(), threadId); + return () -> { + assert Thread.currentThread() == thread : "scope closed on a different thread"; + try { + service.taskExecutionFinishedOnThread(getId(), threadId); + } finally { + accountedThreads.remove(threadId); + } + }; + } + @Override public boolean shouldCancelChildrenOnCancellation() { return true; 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..df0362cb131 --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/QueryInsightsReporter.java @@ -0,0 +1,83 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +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 = 1; + + 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 + * @param failed whether the query failed, so Query Insights can flag it in Top N + */ + public static void report( + TransportService transportService, + DiscoveryNode localNode, + String querySource, + String parentMarker, + String nodeId, + String queryText, + long timestampMillis, + long latencyMillis, + long cpuNanos, + long memoryBytes, + String userInfo, + boolean failed) { + 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)); + 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 + // 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 b1a9068eae9..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 @@ -31,8 +31,10 @@ 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.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.common.utils.QueryContext; @@ -62,6 +64,7 @@ import org.opensearch.sql.plugin.rest.RestUnifiedQueryAction; import org.opensearch.sql.ppl.PPLQueryRunner; 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; @@ -74,6 +77,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; @@ -101,8 +105,10 @@ public class TransportPPLQueryAction private final NodeClient clientRef; private final ClusterService clusterServiceRef; private final org.opensearch.sql.common.setting.Settings pluginSettingsRef; + private final TransportService transportServiceRef; private final QueryJobService queryJobService; private final SecurityAdapter securityAdapter; + private final TaskResourceTrackingService taskResourceTrackingService; @Inject public TransportPPLQueryAction( @@ -116,12 +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()); @@ -186,6 +195,166 @@ private void buildUnifiedQueryHandlerIfReady() { } } + /** + * 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(); + 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); + pplQueryTask.setQueryInsightsParentMarker(value); + } else { + pplQueryTask.setQueryInsightsNested(true); + } + // 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; + } + } + + /** + * 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 void trackEngineThreads(PPLQueryTask pplQueryTask) { + if (pplQueryTask != null) { + pplQueryTask.setResourceTracking(taskResourceTrackingService); + } + } + + /** + * 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 = + 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); + } + } + + /** + * 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 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 + * children's cost would be counted against both records. + */ + static boolean shouldReportToQueryInsights(PPLQueryTask pplQueryTask) { + return pplQueryTask.getQueryInsightsQueryText() != null + && !pplQueryTask.isQueryInsightsExplain() + && !pplQueryTask.isQueryInsightsNested(); + } + + /** Serialize a completed PPL query and send it to Query Insights; errors are ignored. */ + private void writeQueryInsightsRecord(PPLQueryTask pplQueryTask) { + if (!shouldReportToQueryInsights(pplQueryTask)) { + return; + } + try { + String nodeId = clusterServiceRef.localNode().getId(); + String queryText = pplQueryTask.getQueryInsightsQueryText(); + + long cpuNanos = pplQueryTask.getCpuNanos(); + long memoryBytes = pplQueryTask.getAllocatedBytes(); + + // 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 + // 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(); + + // 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(); + + String userInfo = pplQueryTask.getQueryInsightsUserInfo(); + boolean failed = pplQueryTask.isQueryInsightsFailed(); + + QueryInsightsReporter.report( + transportServiceRef, + clusterServiceRef.localNode(), + "PPL", + parentMarker, + nodeId, + queryText, + timestampMillis, + latencyMillis, + cpuNanos, + memoryBytes, + userInfo, + failed); + } 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. @@ -209,7 +378,8 @@ protected void doExecute( return; } - if (task instanceof PPLQueryTask pplQueryTask) { + final PPLQueryTask pplQueryTask = task instanceof PPLQueryTask ? (PPLQueryTask) task : null; + if (pplQueryTask != null) { OpenSearchQueryManager.setCancellableTask(pplQueryTask); } Metrics.getInstance().getNumericalMetric(MetricName.PPL_REQ_TOTAL).increment(); @@ -243,6 +413,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 = recordCompletion(clearingListener, pplQueryTask); try { // Route to analytics engine for non-Lucene (e.g., Parquet-backed) indices. @@ -276,32 +448,64 @@ protected void doExecute( return; } - Consumer anonymizedQuerySink = - anonymized -> rootSpan.addAttribute("db.query.text", anonymized); + Consumer queryInsightsSink = + metadata -> { + rootSpan.addAttribute("db.query.text", metadata.anonymizedQuery()); + if (pplQueryTask != null) { + // 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()); + } + }; 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 if (transformedRequest.shouldRunAsync()) { - submitAsync(task, pplService, transformedRequest, clearingListener, anonymizedQuerySink); + // Async submission answers RUNNING and the query outlives this task, so the task's + // completion is not the query's end: a record stamped from it would carry near-zero + // latency and reference child searches that have not run yet. Left out of Top N until the + // job path reports for itself. + submitAsync(task, pplService, 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)) { + trackEngineThreads(pplQueryTask); + stampQueryInsightsParentHeader(pplQueryTask); + registerQueryInsightsReport(pplQueryTask); + executePplQuery(pplService, transformedRequest, callerListener, queryInsightsSink); + } } else { - pplService.execute( - transformedRequest, - createListener(transformedRequest, clearingListener), - createExplainResponseListener(transformedRequest, clearingListener), - anonymizedQuerySink); + // Tracked without Query Insights recording too: _tasks and Live Queries read the task. + trackEngineThreads(pplQueryTask); + executePplQuery(pplService, transformedRequest, clearingListener, 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(); } } @@ -310,18 +514,14 @@ private void submitAsync( PPLService pplService, PPLQueryRequest transformedRequest, ActionListener listener, - Consumer anonymizedQuerySink) { + Consumer queryInsightsSink) { if (!(task instanceof PPLQueryTask pplQueryTask)) { throw new IllegalArgumentException("Async PPL queries require a cancellable PPLQueryTask"); } Runnable cancelExecution = () -> pplQueryTask.cancel("async PPL query cancelled"); QueryRunner runner = new PPLQueryRunner( - pplService, - transformedRequest, - anonymizedQuerySink, - Clock.systemUTC(), - cancelExecution); + pplService, transformedRequest, queryInsightsSink, Clock.systemUTC(), cancelExecution); ThreadContext threadContext = clientRef.threadPool().getThreadContext(); ActionListener ctxListener = ContextPreservingActionListener.wrapPreservingContext(listener, threadContext); @@ -485,6 +685,51 @@ 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); + } + + /** + * 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; + } + 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 c9502ac3bbf..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,10 +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 { @@ -62,4 +73,205 @@ 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()); + } + + /** 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 testScopesAreNoOpsOnRoutesThatRunNoQuery() { + PPLQueryTask task = newTask(); + assertSame(ThreadResourceAccounting.Scope.NOOP, task.enterThread()); + assertTrue(task.getResourceStats().isEmpty()); + } + + @Test + 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.getResourceStats().isEmpty()); + assertEquals(0L, task.getCpuNanos()); + } + + @Test + 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 = trackedTask(trackingService(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 + public void testQueryInsightsFailedOffByDefault() { + assertFalse(newTask().isQueryInsightsFailed()); + } + + @Test + public void testQueryInsightsFailedFlagIsSettable() { + PPLQueryTask task = newTask(); + task.setQueryInsightsFailed(true); + 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.setQueryInsightsQueryText("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.setQueryInsightsQueryText("source=table | where identifier > 42"); + task.setQueryInsightsFailed(true); + 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.setQueryInsightsQueryText("source=table | where identifier > 42"); + task.setQueryInsightsNested(true); + 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. + PPLQueryTask task = newTask(); + task.setQueryInsightsQueryText("source=table | where identifier > 42"); + assertFalse("nested must default to false", task.isQueryInsightsNested()); + assertTrue(TransportPPLQueryAction.shouldReportToQueryInsights(task)); + } + + @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/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..be31c485215 --- /dev/null +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/QueryInsightsReporterTest.java @@ -0,0 +1,169 @@ +/* + * 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.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.Base64; +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+8MYgIwMjmxgLwzaunARphZG1pbnx8YWxsX2FjY2Vzc3xfX3VzZXJfXwE="; + + 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 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, + "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 + 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, 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("", in.readString()); + assertFalse(in.readBoolean()); + } + } + + @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(); + } +} diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportPPLQueryActionCancellationTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportPPLQueryActionCancellationTest.java index 2f5f7b8bd49..41c6f1784c9 100644 --- a/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportPPLQueryActionCancellationTest.java +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportPPLQueryActionCancellationTest.java @@ -116,7 +116,7 @@ private void submit(Task task) throws Exception { PPLQueryRequest request = new PPLQueryRequest("source=test", new JSONObject(), "/_plugins/_ppl", "jdbc"); submitAsync.invoke( - action, task, pplService, request, listener, PPLService.NO_ANONYMIZED_QUERY_SINK); + action, task, pplService, request, listener, PPLService.NO_QUERY_INSIGHTS_SINK); } private void setField(String name, Object value) throws Exception { 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); + } +} 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 65d7623d220..0226dcdeb59 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLQueryRunner.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLQueryRunner.java @@ -34,7 +34,7 @@ public final class PPLQueryRunner implements QueryRunner { private final PPLService pplService; private final PPLQueryRequest request; - private final Consumer anonymizedQuerySink; + private final Consumer queryInsightsSink; private final Clock clock; private final Runnable cancelExecution; private final AtomicBoolean started = new AtomicBoolean(); @@ -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 anonymizedQuerySink receives the PII-scrubbed query text; supply {@link - * PPLService#NO_ANONYMIZED_QUERY_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 @@ -53,13 +53,13 @@ public final class PPLQueryRunner implements QueryRunner { public PPLQueryRunner( PPLService pplService, PPLQueryRequest request, - Consumer anonymizedQuerySink, + Consumer queryInsightsSink, Clock clock, Runnable cancelExecution) { this.pplService = Objects.requireNonNull(pplService, "pplService must not be null"); this.request = Objects.requireNonNull(request, "request must not be null"); - this.anonymizedQuerySink = - Objects.requireNonNull(anonymizedQuerySink, "anonymizedQuerySink must not be null"); + this.queryInsightsSink = + Objects.requireNonNull(queryInsightsSink, "queryInsightsSink must not be null"); this.clock = Objects.requireNonNull(clock, "clock must not be null"); this.cancelExecution = Objects.requireNonNull(cancelExecution, "cancelExecution must not be null"); @@ -95,7 +95,7 @@ public void onFailure(Exception e) { future.completeExceptionally(e); } }, - anonymizedQuerySink); + queryInsightsSink); return future; } 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..b6aa37227ad 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 com.google.common.base.Strings; 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; @@ -34,8 +36,8 @@ @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 +74,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 for the parsed query 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 +98,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 for the parsed query 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,49 +120,27 @@ 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 for the parsed query 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); - 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); - anonymizedQuerySink.accept(anonymized); - - UnresolvedPlan unresolvedPlan = ((Query) statement).getPlan(); + 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); @@ -174,7 +153,34 @@ private AbstractPlan plan( PPLQueryRequest request, ResponseListener queryListener, ResponseListener explainListener, - Consumer anonymizedQuerySink) { + 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 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) { + String queryText = request.getRequest(); Statement statement; String anonymized; // Transport-thread work — parse, AST build, anonymize. Trace-only: QueryProfiling isn't @@ -182,25 +188,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); @@ -208,9 +200,17 @@ private AbstractPlan plan( } } log.info("[{}] Incoming request {}", QueryContext.getRequestId(), anonymized); - anonymizedQuerySink.accept(anonymized); + return new Prepared( + statement, new QueryInsightsMetadata(anonymized, 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 (Strings.isNullOrEmpty(request.getFormat())) { + return null; + } + return org.opensearch.sql.protocol.response.format.Format.ofExplain(request.getFormat()) + .orElse(null); } } 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..067cd10c69d --- /dev/null +++ b/ppl/src/main/java/org/opensearch/sql/ppl/QueryInsightsMetadata.java @@ -0,0 +1,16 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.ppl; + +/** + * 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 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, boolean explain) {} 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..2555167650a 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; @@ -20,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; @@ -80,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 @@ -115,6 +131,93 @@ 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: 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.assertFalse(metadata.get().explain()); + } + + @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.assertTrue(metadata.get().explain()); + } + + @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.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(