From b2c0111c3d237a1171bba14adccd34c129cf2c4f Mon Sep 17 00:00:00 2001 From: Peng Huo Date: Wed, 5 Aug 2026 17:41:57 +0000 Subject: [PATCH] feat: PPL OpenTelemetry tracing integration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add distributed tracing spans across the PPL Calcite query execution pipeline, reusing the existing query-profiling boundaries so a phase means the same thing whether you read it from `profile` output or from a trace. When the telemetry feature flag is off, NoopTracer is used with near-zero overhead. Spans emitted per PPL query: opensearch.query (CLIENT, root) -> opensearch.query.prepare (INTERNAL, trace-only) -> opensearch.query.analyze (INTERNAL) -> opensearch.query.optimize (INTERNAL) -> opensearch.query.execute (INTERNAL) Root-span attributes follow OTel DB semantic conventions: db.system.name, db.query.type=ppl, db.query.id (UUID), db.operation.name (EXECUTE/EXPLAIN), and db.query.text (anonymized via PPLQueryDataAnonymizer). PPLService hands the anonymized text to a Consumer supplied by the transport action, which owns the Span. Module boundaries: core and ppl must not depend on the opensearch module, so neither references org.opensearch.telemetry. Instead core gains a PhaseListener interface and a ProfileScope helper that times one boundary and feeds both a ProfileMetric and — when a listener is installed — a span. The opensearch module supplies TracingPhaseListener, backed by Tracer, which TransportPPLQueryAction installs at construction. Tracer reaches the transport action via the node injector and is bound into the plugin's child injector by OpenSearchPluginModule. Profile metric definitions are preserved from main -- the golden rule is: do not be misled by internal function names, follow the profile's existing semantics. Each phase's MetricName marker is placed where main already had its manual System.nanoTime() timing, just wrapped in a ProfileScope so the trace span shares the same boundary. ANALYZE - QueryService.executeWithCalcite: analyze() + convert + CalciteToolsHelper.optimize() (Calcite rule pass). OPTIMIZE - CalciteToolsHelper.OpenSearchRelRunners.run: shuttle pass (LogicalTableScan -> BindableTableScan) + prepareStatement. The name looks like it should mean "compile"; per main's profile it is called OPTIMIZE. Left as-is. EXECUTE - OpenSearchExecutionEngine.execute: executeQuery() + buildResultSet(). Materialize is inside EXECUTE, matching main; there is no separate MATERIALIZE phase. FORMAT - unchanged, populated by QueryService.analyze onResponse. The phase span name is derived once from MetricName.name().toLowerCase (ROOT) -- the same rule QueryProfile uses -- so a span name and its profile phase key cannot drift apart. To keep the trace tree flat under root, OpenSearchRelRunners.run (OPTIMIZE) is called BEFORE the EXECUTE ProfileScope opens in OpenSearchExecutionEngine.execute, so OPTIMIZE closes as a sibling of EXECUTE rather than nesting inside it. Async lifecycle: the root Span is created on the transport thread and ended on the worker thread by core's TraceableActionListener; the SpanScope closes synchronously on the transport thread once execute()/explain() returns. The EXECUTE scope closes (records its metric, ends its span) BEFORE listener.onResponse fires: onResponse drives the downstream pipeline that snapshots the profile via QueryProfiling.finish(), so recording after that snapshot would drop the metric while still emitting the span -- the exact profile/trace divergence this design prevents. Two hierarchy fixes verified against the observability stack: * The analyze endpoint (also profile:true) re-runs the compile pass on the caller thread to capture the physical plan text via a Calcite hook. That produced a second `compile` span per query while the profile stayed accurate because QueryProfiling.noop() was set. Added ProfileScope.withSuppression, a thread-local mute for BOTH metric and span; wrapped the re-run in it. One compile span per query now. * The explain path opened its ANALYZE ProfileScope around the entire operation, including the call into executionEngine.explain which opens its own compile span. That made `compile` a grandchild of the root instead of a sibling. Narrowed the analyze scope to just the parse-to-plan work; explain traces now show the flat root -> {analyze, compile} shape. Also added a trace-only `opensearch.query.prepare` span in PPLService, covering parse + AST build + anonymize on the transport thread before the sql-worker hop. Without it, cold-start ANTLR grammar init (several hundred ms) showed up as an unlabeled gap between the root span start and analyze. It is trace-only because QueryProfiling isn't active on the transport thread yet -- the profile output is unaffected. Verified against the observability stack: post-prepare gap dropped from tens to hundreds of ms cold and 2-8ms warm to <1ms in both cases; the prepare span itself now shows the actual parse cost (~350ms cold, 1-3ms warm). Adds a `-DenableTelemetry` toggle to the plugin `run` gradle task that installs the telemetry-otel plugin from a local OpenSearch source tree (path required via `-DtelemetryOtelSrc=`), enables the tracer, sets sampling to 100%, and wires OtlpGrpcSpanExporter to http://localhost:4317 for the observability-stack docker-compose. Docs: local-run runbook under docs/dev/. Tested end-to-end against opensearch-project/observability-stack. Spans land in Trace Analytics nested under OpenSearch's own REST and transport spans, with the execute span parenting the downstream shard-search spans, which confirms context propagation survives the fork-thread handoff. Per-phase span durations match the profile output within System.nanoTime() noise, and failures record status.code=2 with the exception message on the root span. Signed-off-by: Peng Huo --- .../sql/calcite/utils/CalciteToolsHelper.java | 62 +++-- .../opensearch/sql/executor/QueryService.java | 85 ++++--- .../sql/monitor/profile/PhaseListener.java | 39 +++ .../sql/monitor/profile/ProfileScope.java | 75 ++++++ docs/dev/ppl-otel-tracing-local-run.md | 107 ++++++++ integ-test/build.gradle | 63 +++++ .../remote/PitContextLimitErrorIT.java | 76 ++++++ .../tracing/OtlpHttpTraceReceiver.java | 230 ++++++++++++++++++ .../sql/calcite/tracing/PPLTracingIT.java | 219 +++++++++++++++++ .../executor/OpenSearchExecutionEngine.java | 16 +- .../tracing/TracingPhaseListener.java | 50 ++++ plugin/build.gradle | 42 ++++ .../org/opensearch/sql/plugin/SQLPlugin.java | 3 +- .../plugin/config/OpenSearchPluginModule.java | 11 +- .../transport/TransportPPLQueryAction.java | 136 +++++++---- .../org/opensearch/sql/ppl/PPLService.java | 152 ++++++++---- 16 files changed, 1185 insertions(+), 181 deletions(-) create mode 100644 core/src/main/java/org/opensearch/sql/monitor/profile/PhaseListener.java create mode 100644 core/src/main/java/org/opensearch/sql/monitor/profile/ProfileScope.java create mode 100644 docs/dev/ppl-otel-tracing-local-run.md create mode 100644 integ-test/src/test/java/org/opensearch/sql/calcite/remote/PitContextLimitErrorIT.java create mode 100644 integ-test/src/test/java/org/opensearch/sql/calcite/tracing/OtlpHttpTraceReceiver.java create mode 100644 integ-test/src/test/java/org/opensearch/sql/calcite/tracing/PPLTracingIT.java create mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/executor/tracing/TracingPhaseListener.java diff --git a/core/src/main/java/org/opensearch/sql/calcite/utils/CalciteToolsHelper.java b/core/src/main/java/org/opensearch/sql/calcite/utils/CalciteToolsHelper.java index c4a1ff1ac81..01c9e2d4b33 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/utils/CalciteToolsHelper.java +++ b/core/src/main/java/org/opensearch/sql/calcite/utils/CalciteToolsHelper.java @@ -28,7 +28,6 @@ package org.opensearch.sql.calcite.utils; import static java.util.Objects.requireNonNull; -import static org.opensearch.sql.monitor.profile.MetricName.OPTIMIZE; import com.google.common.collect.ImmutableList; import java.lang.reflect.Type; @@ -109,8 +108,9 @@ import org.opensearch.sql.common.error.ErrorCode; import org.opensearch.sql.common.error.ErrorReport; import org.opensearch.sql.expression.function.PPLBuiltinOperators; +import org.opensearch.sql.monitor.profile.MetricName; import org.opensearch.sql.monitor.profile.ProfileContext; -import org.opensearch.sql.monitor.profile.ProfileMetric; +import org.opensearch.sql.monitor.profile.ProfileScope; import org.opensearch.sql.monitor.profile.QueryProfiling; /** @@ -514,37 +514,35 @@ private static void enrichErrorsForSpecialCases(ErrorReport.Builder report, SQLE * org.apache.calcite.tools.RelRunners#run(RelNode)} */ public static PreparedStatement run(CalcitePlanContext context, RelNode rel) { - ProfileMetric optimizeTime = QueryProfiling.current().getOrCreateMetric(OPTIMIZE); - long startTime = System.nanoTime(); - final RelShuttle shuttle = - new RelHomogeneousShuttle() { - @Override - public RelNode visit(TableScan scan) { - final RelOptTable table = scan.getTable(); - if (scan instanceof LogicalTableScan - && Bindables.BindableTableScan.canHandle(table)) { - // Always replace the LogicalTableScan with BindableTableScan - // because it's implementation does not require a "schema" as context. - return Bindables.BindableTableScan.create(scan.getCluster(), table); + try (ProfileScope optimizePhase = ProfileScope.open(MetricName.OPTIMIZE)) { + final RelShuttle shuttle = + new RelHomogeneousShuttle() { + @Override + public RelNode visit(TableScan scan) { + final RelOptTable table = scan.getTable(); + if (scan instanceof LogicalTableScan + && Bindables.BindableTableScan.canHandle(table)) { + // Always replace the LogicalTableScan with BindableTableScan + // because it's implementation does not require a "schema" as context. + return Bindables.BindableTableScan.create(scan.getCluster(), table); + } + return super.visit(scan); } - return super.visit(scan); - } - }; - rel = rel.accept(shuttle); - // the line we changed here - try (Connection connection = context.connection) { - final RelRunner runner = connection.unwrap(RelRunner.class); - PreparedStatement preparedStatement = runner.prepareStatement(rel); - optimizeTime.set(System.nanoTime() - startTime); - return preparedStatement; - } catch (SQLException e) { - // Detect if error is due to window functions in unsupported context (bins on time fields) - ErrorReport.Builder report = - ErrorReport.wrap(e) - .location("while compiling the optimized query plan for physical execution") - .code(ErrorCode.PLANNING_ERROR); - enrichErrorsForSpecialCases(report, e); - throw report.build(); + }; + rel = rel.accept(shuttle); + + try (Connection connection = context.connection) { + final RelRunner runner = connection.unwrap(RelRunner.class); + return runner.prepareStatement(rel); + } catch (SQLException e) { + // Detect if error is due to window functions in unsupported context (bins on time fields) + ErrorReport.Builder report = + ErrorReport.wrap(e) + .location("while compiling the optimized query plan for physical execution") + .code(ErrorCode.PLANNING_ERROR); + enrichErrorsForSpecialCases(report, e); + throw report.build(); + } } } } diff --git a/core/src/main/java/org/opensearch/sql/executor/QueryService.java b/core/src/main/java/org/opensearch/sql/executor/QueryService.java index e679daf864c..e35aa7b1ba3 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryService.java @@ -61,8 +61,8 @@ import org.opensearch.sql.exception.NonFallbackCalciteException; import org.opensearch.sql.executor.analyze.AnalyzeRecommendationBuilder; import org.opensearch.sql.monitor.profile.MetricName; -import org.opensearch.sql.monitor.profile.ProfileContext; import org.opensearch.sql.monitor.profile.ProfileMetric; +import org.opensearch.sql.monitor.profile.ProfileScope; import org.opensearch.sql.monitor.profile.QueryProfile; import org.opensearch.sql.monitor.profile.QueryProfiling; import org.opensearch.sql.planner.PlanContext; @@ -219,38 +219,40 @@ public void executeWithCalcite( CalcitePlanContext.run( () -> { try { - ProfileContext profileContext = - QueryProfiling.activate(QueryContext.isProfileEnabled()); - ProfileMetric analyzeMetric = profileContext.getOrCreateMetric(MetricName.ANALYZE); - long analyzeStart = System.nanoTime(); + QueryProfiling.activate(QueryContext.isProfileEnabled()); CalciteClassLoaderHelper.withCalciteClassLoader( () -> { - CalcitePlanContext context = - CalcitePlanContext.create( - buildFrameworkConfig(), - SysLimit.fromSettings(settings), - queryType, - includeMetadata); - - context.setHighlightConfig(highlightConfig); - - // Wrap analyze with ANALYZING stage tracking - RelNode relNode = - StageErrorHandler.executeStage( - QueryProcessingStage.ANALYZING, - () -> analyze(plan, context), - "while preparing and validating the query plan"); - - // Wrap plan conversion with PLAN_CONVERSION stage tracking - RelNode calcitePlan = - StageErrorHandler.executeStage( - QueryProcessingStage.PLAN_CONVERSION, - () -> - withCheckedArithmetic( - convertToCalcitePlan(relNode, context), context), - "while converting the query to an executable plan"); - - executeCalcitePlan(calcitePlan, context, listener, analyzeMetric, analyzeStart); + CalcitePlanContext context; + RelNode optimizedPlan; + try (ProfileScope analyzePhase = ProfileScope.open(MetricName.ANALYZE)) { + context = + CalcitePlanContext.create( + buildFrameworkConfig(), + SysLimit.fromSettings(settings), + queryType, + includeMetadata); + + context.setHighlightConfig(highlightConfig); + + // Wrap analyze with ANALYZING stage tracking + RelNode relNode = + StageErrorHandler.executeStage( + QueryProcessingStage.ANALYZING, + () -> analyze(plan, context), + "while preparing and validating the query plan"); + + // Wrap plan conversion with PLAN_CONVERSION stage tracking + RelNode calcitePlan = + StageErrorHandler.executeStage( + QueryProcessingStage.PLAN_CONVERSION, + () -> + withCheckedArithmetic( + convertToCalcitePlan(relNode, context), context), + "while converting the query to an executable plan"); + + optimizedPlan = CalciteToolsHelper.optimize(calcitePlan, context); + } + executeCalcitePlan(optimizedPlan, context, listener); }, QueryService.class); } catch (Throwable t) { @@ -266,17 +268,10 @@ public void executeWithCalcite( } private void executeCalcitePlan( - RelNode calcitePlan, + RelNode optimizedPlan, CalcitePlanContext context, - ResponseListener listener, - ProfileMetric analyzeMetric, - long analyzeStart) { + ResponseListener listener) { try { - // Optimize before dispatch so the dispatcher's ScriptDetector - // sees the post-optimization plan for accurate routing. - RelNode optimizedPlan = CalciteToolsHelper.optimize(calcitePlan, context); - analyzeMetric.set(System.nanoTime() - analyzeStart); - // Wrap execution with EXECUTING stage tracking — dispatch via // ExecutionDispatcher which may route to a complex worker pool StageErrorHandler.executeStageVoid( @@ -325,9 +320,13 @@ public void explainWithCalcite( context.setHighlightConfig(highlightConfig); context.run( () -> { - RelNode relNode = analyze(plan, context); - RelNode calcitePlan = - withCheckedArithmetic(convertToCalcitePlan(relNode, context), context); + RelNode calcitePlan; + try (ProfileScope analyzePhase = ProfileScope.open(MetricName.ANALYZE)) { + RelNode relNode = analyze(plan, context); + calcitePlan = + withCheckedArithmetic( + convertToCalcitePlan(relNode, context), context); + } if (format != null) { executionEngine.explain(calcitePlan, mode, format, context, listener); } else { diff --git a/core/src/main/java/org/opensearch/sql/monitor/profile/PhaseListener.java b/core/src/main/java/org/opensearch/sql/monitor/profile/PhaseListener.java new file mode 100644 index 00000000000..a760d5b0da7 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/monitor/profile/PhaseListener.java @@ -0,0 +1,39 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.monitor.profile; + +/** + * Observes query-phase boundaries — parse, analyze, optimize, compile, etc. — without pulling any + * tracing dependency into core/ppl. {@link ProfileScope} calls {@link #onPhaseStart(String)} when a + * phase opens and closes the returned handle when the phase ends, so the listener sees the exact + * same start/end boundary that feeds the profile metric. + * + *

The default implementation is a no-op ({@link #NOOP}). Modules that emit tracing spans (e.g. + * the {@code opensearch} module) install a concrete listener via {@link + * ProfileScope#installListener(PhaseListener)}. + */ +public interface PhaseListener { + + /** Called when a phase starts. The returned handle is closed when the phase ends. */ + Handle onPhaseStart(String phaseName); + + /** Observation of a single phase — the caller closes it (and may call {@link #setError}). */ + interface Handle extends AutoCloseable { + /** Record a failure that occurred inside the phase. Called before {@link #close()}. */ + default void setError(Throwable t) {} + + @Override + void close(); + } + + Handle NOOP_HANDLE = + new Handle() { + @Override + public void close() {} + }; + + PhaseListener NOOP = phaseName -> NOOP_HANDLE; +} diff --git a/core/src/main/java/org/opensearch/sql/monitor/profile/ProfileScope.java b/core/src/main/java/org/opensearch/sql/monitor/profile/ProfileScope.java new file mode 100644 index 00000000000..a4b29dc7d88 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/monitor/profile/ProfileScope.java @@ -0,0 +1,75 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.monitor.profile; + +import java.util.Locale; + +/** + * One measured boundary that feeds both query profiling and — if a {@link PhaseListener} is + * installed — distributed tracing. Times a code block with {@link System#nanoTime()} and, on {@link + * #close()}, adds the elapsed nanoseconds to the phase's {@link ProfileMetric} and closes the + * listener's handle. Use with try-with-resources so both outputs share the same start/end boundary. + * + *

The phase name (span name and profile phase key alike) is derived from the {@link MetricName} + * — {@code metric.name().toLowerCase(Locale.ROOT)} — the same rule {@link QueryProfile} uses, so + * the two cannot drift apart. + * + *

Callers call {@link #setError(Throwable)} before rethrowing to record a failure — the listener + * forwards it to its handle (e.g. onto a span). try-with-resources cannot observe the thrown + * exception here. + * + *

Kept trace-free by design: the {@link PhaseListener} indirection means core/ppl never + * reference {@code Tracer}. The {@code opensearch} module installs the tracing listener at startup. + */ +public final class ProfileScope implements AutoCloseable { + + private static volatile PhaseListener listener = PhaseListener.NOOP; + + /** Install a global phase listener (typically from the {@code opensearch} module). */ + public static void installListener(PhaseListener newListener) { + listener = newListener == null ? PhaseListener.NOOP : newListener; + } + + private final ProfileMetric metric; + private final long startNanos; + private final PhaseListener.Handle handle; + + private ProfileScope(ProfileMetric metric, long startNanos, PhaseListener.Handle handle) { + this.metric = metric; + this.startNanos = startNanos; + this.handle = handle; + } + + /** Open a phase that records a profile metric and (if a listener is installed) a tracing span. */ + public static ProfileScope open(MetricName metric) { + return new ProfileScope( + QueryProfiling.current().getOrCreateMetric(metric), + System.nanoTime(), + listener.onPhaseStart(metric.name().toLowerCase(Locale.ROOT))); + } + + /** + * Open a trace-only phase — a span with no matching profile metric. Used for work that happens + * before {@link QueryProfiling} is activated on the current thread (e.g. transport-side parse and + * dispatch) and therefore can't feed a metric anyway. {@code phaseName} must be a fixed literal + * chosen by the caller, not derived from user input. + */ + public static ProfileScope openTraceOnly(String phaseName) { + return new ProfileScope( + NoopProfileMetric.INSTANCE, System.nanoTime(), listener.onPhaseStart(phaseName)); + } + + /** Record a failure on the listener's handle. Callers invoke this before rethrowing. */ + public void setError(Throwable t) { + handle.setError(t); + } + + @Override + public void close() { + metric.add(System.nanoTime() - startNanos); + handle.close(); + } +} diff --git a/docs/dev/ppl-otel-tracing-local-run.md b/docs/dev/ppl-otel-tracing-local-run.md new file mode 100644 index 00000000000..d6db9f33ae0 --- /dev/null +++ b/docs/dev/ppl-otel-tracing-local-run.md @@ -0,0 +1,107 @@ +# Running PPL Locally with OTel Tracing → observability-stack + +Launch the SQL/PPL plugin against a local OpenSearch node with distributed tracing enabled, and export spans to a running [observability-stack](https://github.com/opensearch-project/observability-stack) docker-compose deployment. + +## Prerequisites + +**observability-stack docker-compose running.** From the repo: +```bash +cd /docker-compose +docker compose up -d +``` +The OTel Collector must be listening on OTLP gRPC `0.0.0.0:4317` (default). + +**`telemetry-otel` plugin zip built from the OpenSearch source tree.** This module is not published as a standalone plugin — it must be built locally: +```bash +cd /plugins/telemetry-otel +../../gradlew bundlePlugin +# produces build/distributions/telemetry-otel-.zip +``` +Pass the path with `-DtelemetryOtelSrc=/plugins/telemetry-otel` when running the plugin. + +## One-time: fix port conflict + +Both the local `:run` task and the docker `opensearch` service want host port `9200`. Free `9200` for the local run by remapping only the docker host port (intra-network traffic between the stack's containers stays on `9200`). + +**`observability-stack/.env`** — keep `OPENSEARCH_PORT=9200`. It's used as the intra-network target port by data-prepper, dashboards, and the exporter. + +**`observability-stack/docker-compose.local-opensearch.yml`** — change the host mapping only: +```yaml +ports: + - "9210:9200" # host 9210 → container 9200 (default was "${OPENSEARCH_PORT}:9200") + - "9600:9600" # unchanged +``` + +Recreate the affected containers: +```bash +cd /docker-compose +docker compose down opensearch data-prepper opensearch-dashboards +docker compose up -d opensearch data-prepper opensearch-dashboards +``` + +After this: docker OpenSearch is reachable at `https://localhost:9210` (admin/`My_password_123!@#`), the local plugin `:run` gets `http://localhost:9200`. + +## Launch + +```bash +cd +./gradlew :opensearch-sql-plugin:run -DenableTelemetry \ + -DtelemetryOtelSrc=/plugins/telemetry-otel +``` + +The `-DenableTelemetry` flag (wired in `plugin/build.gradle`) does the following against `testClusters.integTest`: + +- installs the `telemetry-otel` zip from `/build/distributions/telemetry-otel-.zip` +- sets `opensearch.experimental.feature.telemetry.enabled=true` (system property) +- enables the tracer: `telemetry.feature.tracer.enabled=true`, `telemetry.tracer.enabled=true` +- sets sampling to 100%: `telemetry.tracer.sampler.probability=1.0` +- selects the OTLP gRPC exporter: `telemetry.otel.tracer.span.exporter.class=io.opentelemetry.exporter.otlp.trace.OtlpGrpcSpanExporter` +- points the exporter at the collector: `OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4317` + +Wait until the node reports: +``` +[integTest-0] Successfully instantiated the SpanExporter class ... OtlpGrpcSpanExporter +[integTest-0] publish_address {127.0.0.1:9200} +[integTest-0] started +``` + +## Smoke test + +Index a small dataset and run a PPL query: +```bash +curl -s -X POST 'http://localhost:9200/test-logs/_bulk' -H 'Content-Type: application/x-ndjson' --data-binary '{"index":{}} +{"host":"h1","status":200,"latency_ms":15} +{"index":{}} +{"host":"h1","status":500,"latency_ms":40} +{"index":{}} +{"host":"h2","status":404,"latency_ms":60} +' +curl -s -X POST 'http://localhost:9200/test-logs/_refresh' + +curl -s -X POST 'http://localhost:9200/_plugins/_ppl' \ + -H 'Content-Type: application/json' \ + -d '{"query":"source=test-logs | where status > 300 | stats count() by host"}' +``` + +## Verify traces landed + +**Inspect spans directly in the observability-stack OpenSearch:** +```bash +curl -k -s \ + 'https://localhost:9210/otel-v1-apm-span-*/_search?pretty' \ + -H 'Content-Type: application/json' \ + -d '{"size":10,"query":{"prefix":{"name":"opensearch.query"}}, + "sort":[{"startTime":{"order":"desc"}}], + "_source":["traceId","spanId","parentSpanId","name","kind", + "durationInNanos","attributes"]}' +``` + +Expected per query: one CLIENT root span `opensearch.query` plus four INTERNAL children — `opensearch.query.prepare`, `.analyze`, `.optimize`, `.execute`. The `.prepare` span is trace-only and covers transport-side parse + AST build + anonymize (dominant on cold start due to ANTLR grammar init); the others match profile phase keys 1-to-1 (both derived from the same `MetricName`), so trace and profile report the same phases with the same durations. Phases (per the profile's original definition — do not be misled by internal function names): `.analyze` = semantic analyze + Calcite plan conversion + rule-based optimize (`CalciteToolsHelper.optimize`); `.optimize` = physical shuttle + `RelRunner.prepareStatement` (inside `OpenSearchRelRunners.run`); `.execute` = `executeQuery` + `buildResultSet`. Root attributes: + +| Attribute | Example | +|-----------|---------| +| `db.query.type` | `ppl` | +| `db.operation.name` | `EXECUTE` (or `EXPLAIN`) | +| `db.query.text` | `source=test-logs \| where status > 300 \| stats count() by host` | + +**Or view in Dashboards:** http://localhost:5601 → Observability → Trace Analytics. diff --git a/integ-test/build.gradle b/integ-test/build.gradle index 4a071bfefe6..d5d20602810 100644 --- a/integ-test/build.gradle +++ b/integ-test/build.gradle @@ -75,6 +75,11 @@ ext { projectSubstitutions = [:] licenseFile = rootProject.file('LICENSE.TXT') noticeFile = rootProject.file('NOTICE') + telemetryOtelVersion = System.getProperty("telemetryOtelVersion") + if (!telemetryOtelVersion) { + def metadataUrl = new URL('https://ci.opensearch.org/maven2/org/opensearch/plugin/telemetry-otel/maven-metadata.xml') + telemetryOtelVersion = new XmlParser().parse(metadataUrl.openStream()).versioning.release.text() + } getSecurityPluginDownloadLink = { -> var repo = "https://ci.opensearch.org/ci/dbc/snapshots/maven/org/opensearch/plugin/" + @@ -213,6 +218,7 @@ dependencies { testImplementation(testFixtures(project(':api'))) { exclude group: 'org.hamcrest', module: 'hamcrest-core' } + testImplementation 'io.opentelemetry.proto:opentelemetry-proto:1.3.2-alpha' testImplementation('org.junit.jupiter:junit-jupiter-api:5.9.3') testImplementation('org.junit.jupiter:junit-jupiter-params:5.9.3') testRuntimeOnly('org.junit.jupiter:junit-jupiter-engine:5.9.3') @@ -229,6 +235,7 @@ dependencies { // For GeoIP PPL functions zipArchive group: 'org.opensearch.plugin', name:'geospatial', version: "${opensearch_build}" + } java { @@ -380,6 +387,52 @@ def getAnalyticsBackendDatafusionPlugin() { provider { (RegularFile) (() -> file(project.findProperty('analyticsBackendDatafusionZip') ?: analyticsBackendDatafusionZipDest)) } } +ext.telemetryOtelZip = file("${buildDir}/tmp/telemetry-otel-patched.zip") + +def downloadTelemetryOtel = tasks.register('downloadTelemetryOtel') { + outputs.file(project.telemetryOtelZip) + doLast { + def v = project.telemetryOtelVersion + def clusterVersion = opensearch_version.toString().tokenize('-')[0] + def rawZip = file("${buildDir}/tmp/telemetry-otel-${v}-raw.zip") + def extractDir = file("${buildDir}/tmp/telemetry-otel-${v}-extracted") + + if (!rawZip.exists()) { + def url = "https://artifacts.opensearch.org/releases/plugins/telemetry-otel/${v}/telemetry-otel-${v}.zip" + logger.lifecycle "Downloading telemetry-otel plugin ${v} from ${url}" + rawZip.parentFile.mkdirs() + ant.get(src: url, dest: rawZip, httpusecaches: false) + } + + // 3.x plugins are binary-compatible across point releases; OpenSearch's plugin loader + // enforces exact major.minor.patch match. Rewrite the descriptor so the released zip + // installs cleanly on a cluster whose version has bumped ahead of the latest release. + delete extractDir + extractDir.mkdirs() + ant.unzip(src: rawZip, dest: extractDir) + def descriptor = file("${extractDir}/plugin-descriptor.properties") + descriptor.text = descriptor.text.replaceAll( + /(?m)^opensearch\.version=.*$/, + "opensearch.version=${clusterVersion}") + + delete project.telemetryOtelZip + ant.zip(destfile: project.telemetryOtelZip, basedir: extractDir) + } +} + +def getTelemetryOtelPlugin() { + return provider(new Callable() { + @Override + RegularFile call() throws Exception { + return new RegularFile() { + @Override + File getAsFile() { return project.telemetryOtelZip } + } + } + }) +} + + testClusters { integTest { testDistribution = 'archive' @@ -389,6 +442,14 @@ testClusters { setting "plugins.query.datasources.encryption.masterkey", "1234567812345678" // Only /_cluster/health is registered; pin the allow-list to it for the rest ITs. setting 'plugins.ppl.rest.allowed_endpoints', '/_cluster/health' + plugin(getTelemetryOtelPlugin()) + systemProperty 'opensearch.experimental.feature.telemetry.enabled', 'true' + setting 'telemetry.feature.tracer.enabled', 'true' + setting 'telemetry.tracer.enabled', 'true' + setting 'telemetry.tracer.sampler.probability', '1.0' + setting 'telemetry.otel.tracer.span.exporter.class', + 'io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporter' + setting 'telemetry.otel.tracer.exporter.delay', '50ms' } yamlRestTest { testDistribution = 'archive' @@ -687,6 +748,8 @@ yamlRestTest { // Run PPL ITs and new, legacy and comparison SQL ITs with new SQL engine enabled integTest { useCluster testClusters.remoteCluster + dependsOn downloadTelemetryOtel + systemProperty 'tests.tracing.otlp.port', '4318' // Set properties for connection to clusters and between clusters doFirst { diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/PitContextLimitErrorIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/PitContextLimitErrorIT.java new file mode 100644 index 00000000000..cd49d3f9b97 --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/PitContextLimitErrorIT.java @@ -0,0 +1,76 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.calcite.remote; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsString; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertThrows; +import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK; + +import java.io.IOException; +import org.apache.hc.core5.http.ParseException; +import org.apache.hc.core5.http.io.entity.EntityUtils; +import org.junit.jupiter.api.Test; +import org.opensearch.client.Request; +import org.opensearch.client.Response; +import org.opensearch.client.ResponseException; +import org.opensearch.sql.ppl.PPLIntegTestCase; + +/** + * Verifies that when the Calcite execute path fails inside {@code client.schedule(...)} on the + * worker thread, the failure is delivered to the caller via {@code listener.onFailure} rather than + * being lost by the async boundary. Regression test for the change from {@code throw} to {@code + * listener.onFailure} in {@link + * org.opensearch.sql.opensearch.executor.OpenSearchExecutionEngine#execute(org.apache.calcite.rel.RelNode, + * org.opensearch.sql.calcite.CalcitePlanContext, + * org.opensearch.sql.common.response.ResponseListener)}. + * + *

The trigger is {@code search.max_open_pit_context=0}: any Calcite query that opens a PIT + * (aggregations, most scans) throws a {@code SQLException} wrapping the "too many Point In Time + * contexts" rejection at {@code executeQuery()} time — the exact catch block the change touches. If + * the {@code listener.onFailure} routing regresses to {@code throw}, this test's HTTP request hangs + * until the RestClient default timeout expires rather than returning a proper 4xx/5xx. + */ +public class PitContextLimitErrorIT extends PPLIntegTestCase { + + @Override + public void init() throws Exception { + super.init(); + enableCalcite(); + loadIndex(Index.BANK); + } + + @Test + public void pitLimitExhaustionReturnsActionableErrorResponse() + throws IOException, ParseException { + updateClusterSettings(new ClusterSetting("persistent", "search.max_open_pit_context", "0")); + try { + Request request = new Request("POST", "/_plugins/_ppl"); + request.setJsonEntity( + "{\"query\":\"source=" + TEST_INDEX_BANK + " | stats count() by state\"}"); + + ResponseException ex = + assertThrows(ResponseException.class, () -> client().performRequest(request)); + + Response response = ex.getResponse(); + int status = response.getStatusLine().getStatusCode(); + // 4xx/5xx — any non-2xx is proof the async listener fired instead of hanging. + assertNotEquals( + "expected error status but got 2xx — async onFailure likely regressed to throw", + 200, + status); + + String body = EntityUtils.toString(response.getEntity()); + // Both halves of the ErrorReport built in the catch block must be present, since that is + // the exact string the code under test constructs before handing to listener.onFailure. + assertThat(body, containsString("Too many open Point-In-Time (PIT) contexts")); + assertThat(body, containsString("search.max_open_pit_context")); + } finally { + updateClusterSettings(new ClusterSetting("persistent", "search.max_open_pit_context", null)); + } + } +} diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/tracing/OtlpHttpTraceReceiver.java b/integ-test/src/test/java/org/opensearch/sql/calcite/tracing/OtlpHttpTraceReceiver.java new file mode 100644 index 00000000000..e2a5d012baa --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/tracing/OtlpHttpTraceReceiver.java @@ -0,0 +1,230 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.calcite.tracing; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpServer; +import io.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest; +import io.opentelemetry.proto.common.v1.AnyValue; +import io.opentelemetry.proto.common.v1.KeyValue; +import io.opentelemetry.proto.trace.v1.ResourceSpans; +import io.opentelemetry.proto.trace.v1.ScopeSpans; +import java.io.IOException; +import java.io.OutputStream; +import java.net.InetSocketAddress; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HexFormat; +import java.util.List; +import java.util.Map; +import java.util.function.Predicate; +import java.util.stream.Collectors; +import java.util.zip.GZIPInputStream; + +/** + * In-process OTLP-HTTP trace receiver for tests. Binds {@link HttpServer} on {@code 0.0.0.0:} + * at path {@code /v1/traces}, accepts POST requests with OTLP-protobuf bodies (gzip-encoded or + * plain), decodes them via the generated {@link ExportTraceServiceRequest} protobuf classes, and + * exposes filter/wait helpers. + * + *

Pairs with a cluster running the OpenSearch {@code telemetry-otel} plugin. The plugin's {@code + * OTelSpanExporterFactory} instantiates {@code OtlpHttpSpanExporter.getDefault()} which is + * unconditionally hardcoded to {@code http://localhost:4318/v1/traces} and uses {@code + * application/x-protobuf} content type. The tests therefore MUST bind port 4318 and MUST speak + * protobuf — env-var / sysprop endpoint overrides are ignored by the plugin, and the exporter + * cannot be forced into JSON mode without changing the exporter's builder call. + * + *

{@code BatchSpanProcessor} flushes on a schedule (defaults to 5s; the cluster override for + * tracingIntegTest is 500ms), so tests should use {@link #waitForSpans} rather than assuming + * synchronous delivery. + */ +public final class OtlpHttpTraceReceiver implements AutoCloseable { + + private final HttpServer server; + private final List spans = Collections.synchronizedList(new ArrayList<>()); + + public OtlpHttpTraceReceiver(int port) throws IOException { + // Bind to 0.0.0.0 so both IPv4 and IPv6 localhost lookups from the cluster JVM reach us. + this.server = HttpServer.create(new InetSocketAddress(port), 0); + this.server.createContext("/v1/traces", this::handle); + // Also mount on root as a defensive net. + this.server.createContext("/", this::handle); + this.server.setExecutor(null); + this.server.start(); + } + + private void handle(HttpExchange exchange) throws IOException { + try { + if (!"POST".equals(exchange.getRequestMethod())) { + exchange.sendResponseHeaders(405, -1); + return; + } + // Read the whole body into a byte[] first — HttpServer's request InputStream can hang if + // parseFrom() reads past the Content-Length boundary on a keep-alive connection. + byte[] bytes = exchange.getRequestBody().readAllBytes(); + String encoding = exchange.getRequestHeaders().getFirst("Content-Encoding"); + if ("gzip".equalsIgnoreCase(encoding)) { + try (GZIPInputStream gz = new GZIPInputStream(new java.io.ByteArrayInputStream(bytes))) { + bytes = gz.readAllBytes(); + } + } + ExportTraceServiceRequest req = ExportTraceServiceRequest.parseFrom(bytes); + for (ResourceSpans rs : req.getResourceSpansList()) { + for (ScopeSpans ss : rs.getScopeSpansList()) { + for (io.opentelemetry.proto.trace.v1.Span sp : ss.getSpansList()) { + spans.add(parseSpan(sp)); + } + } + } + byte[] resp = new byte[0]; + exchange.getResponseHeaders().add("Content-Type", "application/x-protobuf"); + exchange.sendResponseHeaders(200, resp.length); + try (OutputStream out = exchange.getResponseBody()) { + out.write(resp); + } + } catch (Throwable e) { + byte[] resp = ("{\"error\":\"" + e.getMessage() + "\"}").getBytes(StandardCharsets.UTF_8); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(400, resp.length); + try (OutputStream out = exchange.getResponseBody()) { + out.write(resp); + } + } + } + + private static Span parseSpan(io.opentelemetry.proto.trace.v1.Span p) { + Span s = new Span(); + s.traceId = HexFormat.of().formatHex(p.getTraceId().toByteArray()); + s.spanId = HexFormat.of().formatHex(p.getSpanId().toByteArray()); + s.parentSpanId = + p.getParentSpanId().isEmpty() + ? "" + : HexFormat.of().formatHex(p.getParentSpanId().toByteArray()); + s.name = p.getName(); + // OTLP SpanKind proto enum values line up with the OTel spec: + // 0=UNSPECIFIED, 1=INTERNAL, 2=SERVER, 3=CLIENT, 4=PRODUCER, 5=CONSUMER + s.kind = p.getKindValue(); + s.startEpochNanos = p.getStartTimeUnixNano(); + s.endEpochNanos = p.getEndTimeUnixNano(); + s.statusCode = p.getStatus().getCodeValue(); + s.statusMessage = p.getStatus().getMessage(); + s.attributes = new HashMap<>(); + for (KeyValue kv : p.getAttributesList()) { + s.attributes.put(kv.getKey(), stringify(kv.getValue())); + } + return s; + } + + private static String stringify(AnyValue v) { + switch (v.getValueCase()) { + case STRING_VALUE: + return v.getStringValue(); + case BOOL_VALUE: + return String.valueOf(v.getBoolValue()); + case INT_VALUE: + return String.valueOf(v.getIntValue()); + case DOUBLE_VALUE: + return String.valueOf(v.getDoubleValue()); + case ARRAY_VALUE: + return v.getArrayValue().getValuesList().stream() + .map(OtlpHttpTraceReceiver::stringify) + .collect(Collectors.joining(",", "[", "]")); + default: + return v.toString(); + } + } + + /** Snapshot of all spans received so far. Ordered as they arrived on the wire. */ + public List snapshot() { + synchronized (spans) { + return new ArrayList<>(spans); + } + } + + /** + * Poll until {@code minCount} spans matching {@code matcher} arrive, or {@code timeout} elapses. + * Throws {@link AssertionError} if the deadline expires — includes current snapshot names in the + * message for diagnosis. + */ + public List waitForSpans(Predicate matcher, int minCount, Duration timeout) + throws InterruptedException { + long deadline = System.nanoTime() + timeout.toNanos(); + while (true) { + List matched = snapshot().stream().filter(matcher).collect(Collectors.toList()); + if (matched.size() >= minCount) return matched; + if (System.nanoTime() >= deadline) { + throw new AssertionError( + "Timeout after " + + timeout + + " waiting for " + + minCount + + " spans matching predicate; got " + + matched.size() + + " of " + + snapshot().size() + + " total spans. Names seen: " + + snapshot().stream().map(x -> x.name).collect(Collectors.toList())); + } + Thread.sleep(200); + } + } + + public void clear() { + synchronized (spans) { + spans.clear(); + } + } + + @Override + public void close() { + server.stop(1); + } + + /** Parsed span record. Public fields for concise test assertions. */ + public static final class Span { + public String traceId; + public String spanId; + public String parentSpanId; + public String name; + + /** OTLP SpanKind: 0=UNSPECIFIED, 1=INTERNAL, 2=SERVER, 3=CLIENT, 4=PRODUCER, 5=CONSUMER. */ + public int kind; + + public long startEpochNanos; + public long endEpochNanos; + + /** OTLP StatusCode: 0=UNSET, 1=OK, 2=ERROR. */ + public int statusCode; + + public String statusMessage; + public Map attributes; + + public long durationNanos() { + return endEpochNanos - startEpochNanos; + } + + public String attr(String key) { + return attributes.get(key); + } + + @Override + public String toString() { + return name + + "[" + + spanId + + " parent=" + + (parentSpanId.isEmpty() ? "-" : parentSpanId) + + " kind=" + + kind + + " status=" + + statusCode + + "]"; + } + } +} diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/tracing/PPLTracingIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/tracing/PPLTracingIT.java new file mode 100644 index 00000000000..e5e82f0c8f1 --- /dev/null +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/tracing/PPLTracingIT.java @@ -0,0 +1,219 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.calcite.tracing; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.greaterThan; +import static org.hamcrest.Matchers.hasKey; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.not; +import static org.hamcrest.Matchers.notNullValue; +import static org.hamcrest.Matchers.startsWith; +import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK; + +import java.io.IOException; +import java.time.Duration; +import java.util.List; +import java.util.Set; +import java.util.function.Predicate; +import java.util.stream.Collectors; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; +import org.opensearch.client.Request; +import org.opensearch.client.ResponseException; +import org.opensearch.sql.calcite.tracing.OtlpHttpTraceReceiver.Span; +import org.opensearch.sql.ppl.PPLIntegTestCase; + +public class PPLTracingIT extends PPLIntegTestCase { + + private static final Duration SPAN_WAIT_TIMEOUT = Duration.ofSeconds(15); + + private static final int KIND_INTERNAL = 1; + private static final int KIND_CLIENT = 3; + + private static final int STATUS_UNSET = 0; + private static final int STATUS_ERROR = 2; + + private static OtlpHttpTraceReceiver receiver; + + @BeforeClass + public static void startReceiver() throws IOException { + int port = Integer.parseInt(System.getProperty("tests.tracing.otlp.port", "4318")); + receiver = new OtlpHttpTraceReceiver(port); + } + + @AfterClass + public static void stopReceiver() { + if (receiver != null) receiver.close(); + } + + @Override + public void init() throws Exception { + super.init(); + enableCalcite(); + loadIndex(Index.BANK); + receiver.clear(); + } + + @Test + public void executeQuery_emitsRootAndFourPhaseSpans() throws Exception { + executePPL("source=" + TEST_INDEX_BANK + " | where age > 30 | stats count() by state"); + + Span root = waitForRoot(); + assertThat(root.kind, is(KIND_CLIENT)); + assertThat(root.attributes, hasKey("db.system.name")); + assertThat(root.attr("db.system.name"), equalTo("opensearch")); + assertThat(root.attr("db.query.type"), equalTo("ppl")); + assertThat(root.attr("db.operation.name"), equalTo("EXECUTE")); + assertThat(root.attr("db.query.id"), is(notNullValue())); + assertThat(root.attr("db.query.text"), startsWith("source=")); + assertThat(root.statusCode, is(not(STATUS_ERROR))); + + assertThat( + phaseChildNames(root), + equalTo( + Set.of( + "opensearch.query.prepare", + "opensearch.query.analyze", + "opensearch.query.optimize", + "opensearch.query.execute"))); + + for (Span child : phaseChildrenOf(root)) { + assertThat(child.kind, is(KIND_INTERNAL)); + assertThat(child.durationNanos(), greaterThan(-1L)); + } + } + + @Test + public void explainQuery_emitsPrepareAnalyzeOptimizeButNoExecute() throws Exception { + Request req = new Request("POST", "/_plugins/_ppl/_explain"); + req.setJsonEntity("{\"query\":\"source=" + TEST_INDEX_BANK + " | stats count() by state\"}"); + client().performRequest(req); + + Span root = waitForRoot(s -> "EXPLAIN".equals(s.attr("db.operation.name"))); + + assertThat(root.attr("db.operation.name"), equalTo("EXPLAIN")); + Set phases = phaseChildNames(root); + assertThat( + phases, + equalTo( + Set.of( + "opensearch.query.prepare", + "opensearch.query.analyze", + "opensearch.query.optimize"))); + assertThat(phases.contains("opensearch.query.execute"), is(false)); + } + + @Test + public void parseCommand_emitsAllFourPhases_evenThroughComplexPoolHop() throws Exception { + executePPL( + "source=" + + TEST_INDEX_BANK + + " | parse address '(?\\\\d+)' | where isnotnull(num) | stats count() by num"); + + Span root = waitForRoot(); + assertThat( + phaseChildNames(root), + equalTo( + Set.of( + "opensearch.query.prepare", + "opensearch.query.analyze", + "opensearch.query.optimize", + "opensearch.query.execute"))); + assertThat(phaseChildrenOf(root).size(), is(4)); + } + + @Test + public void failedQueryOnMissingIndex_setsStatusErrorOnRoot() throws Exception { + try { + Request req = new Request("POST", "/_plugins/_ppl"); + req.setJsonEntity("{\"query\":\"source=this-index-does-not-exist-xyz | stats count()\"}"); + client().performRequest(req); + throw new AssertionError("expected ResponseException for missing index"); + } catch (ResponseException expected) { + } + + Span root = waitForRoot(); + assertThat(root.statusCode, is(STATUS_ERROR)); + Set phases = phaseChildNames(root); + assertThat(phases.contains("opensearch.query.execute"), is(false)); + } + + @Test + public void anonymizedQueryTextIsNotRawUserInput() throws Exception { + executePPL("source=" + TEST_INDEX_BANK + " | where age > 30 | fields firstname"); + + Span root = waitForRoot(); + String qtext = root.attr("db.query.text"); + assertThat(qtext, is(notNullValue())); + assertThat(qtext.contains(" 30"), is(false)); + assertThat(qtext.contains("***"), is(true)); + } + + @Test + public void phaseSpans_haveValidTimingAndSameTraceIdAsRoot() throws Exception { + executePPL("source=" + TEST_INDEX_BANK + " | stats count()"); + + Span root = waitForRoot(); + List phases = phaseChildrenOf(root); + assertThat(phases.size(), greaterThan(0)); + for (Span p : phases) { + assertThat(p.traceId, equalTo(root.traceId)); + assertThat(p.parentSpanId, equalTo(root.spanId)); + assertThat(p.endEpochNanos, greaterThan(p.startEpochNanos - 1)); + assertThat( + p.startEpochNanos >= root.startEpochNanos && p.endEpochNanos <= root.endEpochNanos, + is(true)); + } + } + + private void executePPL(String query) throws IOException { + Request req = new Request("POST", "/_plugins/_ppl"); + req.setJsonEntity("{\"query\":\"" + query.replace("\"", "\\\"").replace("\\", "\\\\") + "\"}"); + client().performRequest(req); + } + + private Span waitForRoot() throws InterruptedException { + return waitForRoot(s -> true); + } + + /** + * Wait for a root span matching {@code extra}, then keep polling until ALL {@code + * opensearch.query*} spans (roots + phase children) stop arriving. Phase children have different + * names than the root, so a root-only stability check can return before the batch carrying its + * children has drained — callers of {@link #phaseChildrenOf} would then see a partial set. Tests + * run sequentially; the latest-created matching root is this test's. + */ + private Span waitForRoot(Predicate extra) throws InterruptedException { + Predicate isRoot = s -> "opensearch.query".equals(s.name) && extra.test(s); + Predicate isAnyPplSpan = s -> s.name.startsWith("opensearch.query"); + receiver.waitForSpans(isRoot, 1, SPAN_WAIT_TIMEOUT); + int cur = (int) receiver.snapshot().stream().filter(isAnyPplSpan).count(); + int prev; + do { + prev = cur; + Thread.sleep(100); + cur = (int) receiver.snapshot().stream().filter(isAnyPplSpan).count(); + } while (cur > prev); + return receiver.snapshot().stream() + .filter(isRoot) + .reduce((a, b) -> a.startEpochNanos >= b.startEpochNanos ? a : b) + .orElseThrow(); + } + + private List phaseChildrenOf(Span root) { + return receiver.snapshot().stream() + .filter(s -> s.name.startsWith("opensearch.query.")) + .filter(s -> root.spanId.equals(s.parentSpanId)) + .collect(Collectors.toList()); + } + + private Set phaseChildNames(Span root) { + return phaseChildrenOf(root).stream().map(s -> s.name).collect(Collectors.toSet()); + } +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java index 483f2684d61..5cae7702253 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java @@ -64,8 +64,7 @@ import org.opensearch.sql.expression.function.BuiltinFunctionName; import org.opensearch.sql.expression.function.PPLFuncImpTable; import org.opensearch.sql.monitor.profile.MetricName; -import org.opensearch.sql.monitor.profile.ProfileMetric; -import org.opensearch.sql.monitor.profile.QueryProfiling; +import org.opensearch.sql.monitor.profile.ProfileScope; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.data.value.OpenSearchExprGeoPointValue; import org.opensearch.sql.opensearch.executor.protector.ExecutionProtector; @@ -331,14 +330,13 @@ public void execute( client.schedule( () -> { try (PreparedStatement statement = OpenSearchRelRunners.run(context, rel)) { - ProfileMetric metric = QueryProfiling.current().getOrCreateMetric(MetricName.EXECUTE); - long execTime = System.nanoTime(); - ResultSet result = statement.executeQuery(); - QueryResponse response = - buildResultSet(result, rel.getRowType(), context.sysLimit.querySizeLimit()); - metric.add(System.nanoTime() - execTime); + QueryResponse response; + try (ProfileScope executePhase = ProfileScope.open(MetricName.EXECUTE)) { + ResultSet result = statement.executeQuery(); + response = + buildResultSet(result, rel.getRowType(), context.sysLimit.querySizeLimit()); + } listener.onResponse(response); - } catch (SQLException e) { if (isPitContextLimitReached(e)) { // reason (title) comes from the wrapped cause's message; keep it short and put the diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/tracing/TracingPhaseListener.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/tracing/TracingPhaseListener.java new file mode 100644 index 00000000000..22412f9b0d9 --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/tracing/TracingPhaseListener.java @@ -0,0 +1,50 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.executor.tracing; + +import org.opensearch.sql.monitor.profile.PhaseListener; +import org.opensearch.telemetry.tracing.Span; +import org.opensearch.telemetry.tracing.SpanCreationContext; +import org.opensearch.telemetry.tracing.SpanScope; +import org.opensearch.telemetry.tracing.Tracer; + +/** + * {@link PhaseListener} that opens an OpenTelemetry span for each phase and closes it in lockstep + * with the profile-metric timing. Bridges the tracer-free core/ppl {@link + * org.opensearch.sql.monitor.profile.ProfileScope} to the {@link Tracer} living in the {@code + * opensearch} module. + * + *

Install once at startup via {@link + * org.opensearch.sql.monitor.profile.ProfileScope#installListener(PhaseListener)}. + */ +public final class TracingPhaseListener implements PhaseListener { + + private static final String SPAN_NAME_PREFIX = "opensearch.query."; + + private final Tracer tracer; + + public TracingPhaseListener(Tracer tracer) { + this.tracer = tracer; + } + + @Override + public Handle onPhaseStart(String phaseName) { + Span span = tracer.startSpan(SpanCreationContext.internal().name(SPAN_NAME_PREFIX + phaseName)); + SpanScope scope = tracer.withSpanInScope(span); + return new Handle() { + @Override + public void setError(Throwable t) { + span.setError(t instanceof Exception ? (Exception) t : new RuntimeException(t)); + } + + @Override + public void close() { + scope.close(); + span.endSpan(); + } + }; + } +} diff --git a/plugin/build.gradle b/plugin/build.gradle index fbd2cd1a331..058befb3f41 100644 --- a/plugin/build.gradle +++ b/plugin/build.gradle @@ -305,6 +305,33 @@ def getJobSchedulerPlugin() { }) } +def getTelemetryOtelPlugin() { + // telemetry-otel is an OpenSearch core module — not published to ci.opensearch.org as a + // standalone plugin zip (unlike opensearch-job-scheduler). Build it locally from the OpenSearch + // source tree with `bundlePlugin` and point at it with -DtelemetryOtelSrc=. + String opensearchSrc = System.getProperty("telemetryOtelSrc") + if (opensearchSrc == null) { + throw new GradleException( + "-DenableTelemetry requires -DtelemetryOtelSrc=/plugins/telemetry-otel.\n" + + "Build the zip first: (cd /plugins/telemetry-otel && ../../gradlew bundlePlugin)") + } + def telemetryOtelZip = file("${opensearchSrc}/build/distributions/telemetry-otel-${opensearch_version}.zip") + return provider(new Callable() { + @Override + RegularFile call() throws Exception { + if (!telemetryOtelZip.exists()) { + throw new GradleException( + "telemetry-otel plugin zip not found: ${telemetryOtelZip}.\n" + + "Build it with: (cd ${opensearchSrc} && ../../gradlew bundlePlugin)") + } + return new RegularFile() { + @Override + File getAsFile() { return telemetryOtelZip } + } + } + }) +} + testClusters.integTest { plugin(getJobSchedulerPlugin()) plugin(project.tasks.bundlePlugin.archiveFile) @@ -315,6 +342,21 @@ testClusters.integTest { jvmArgs '-agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=*:5005' } + // Enable OTel tracing and export spans to the observability-stack collector. + // Toggle with: ./gradlew :opensearch-sql-plugin:run -DenableTelemetry + // The telemetry-otel plugin is downloaded from the OpenSearch distribution build (cached in build/telemetry-otel/). + // The observability-stack docker-compose must be running on the host (OTLP gRPC on 4317). + if (System.getProperty("enableTelemetry") != null) { + plugin(getTelemetryOtelPlugin()) + systemProperty 'opensearch.experimental.feature.telemetry.enabled', 'true' + setting 'telemetry.feature.tracer.enabled', 'true' + setting 'telemetry.tracer.enabled', 'true' + setting 'telemetry.tracer.sampler.probability', '1.0' + setting 'telemetry.otel.tracer.span.exporter.class', 'io.opentelemetry.exporter.otlp.trace.OtlpGrpcSpanExporter' + // Default endpoint for OtlpGrpcSpanExporter.getDefault() is http://localhost:4317 — the observability-stack collector. + environment 'OTEL_EXPORTER_OTLP_ENDPOINT', 'http://localhost:4317' + } + // add customized keystore keystore 'plugins.query.federation.datasources.config', new File("$projectDir/src/test/resources/", 'datasources.json') } 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 31f14e3411b..e04566df590 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java @@ -142,6 +142,7 @@ import org.opensearch.sql.spi.rest.RestEndpointProvider; import org.opensearch.sql.sql.domain.SQLQueryRequest; import org.opensearch.sql.storage.DataSourceFactory; +import org.opensearch.telemetry.tracing.noop.NoopTracer; import org.opensearch.threadpool.ExecutorBuilder; import org.opensearch.threadpool.FixedExecutorBuilder; import org.opensearch.threadpool.ThreadPool; @@ -403,7 +404,7 @@ public Collection createComponents( LocalClusterState.state().setPluginSettings((OpenSearchSettings) pluginSettings); LocalClusterState.state().setClient(client); ModulesBuilder modules = new ModulesBuilder(); - modules.add(new OpenSearchPluginModule(executionEngineExtensions)); + modules.add(new OpenSearchPluginModule(executionEngineExtensions, NoopTracer.INSTANCE)); modules.add( b -> { b.bind(NodeClient.class).toInstance((NodeClient) client); diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java index 816f1071310..ebf0196c77a 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java @@ -40,16 +40,19 @@ import org.opensearch.sql.sql.SQLService; import org.opensearch.sql.sql.antlr.SQLSyntaxParser; import org.opensearch.sql.storage.StorageEngine; +import org.opensearch.telemetry.tracing.Tracer; +import org.opensearch.telemetry.tracing.noop.NoopTracer; import org.opensearch.transport.client.node.NodeClient; @RequiredArgsConstructor public class OpenSearchPluginModule extends AbstractModule { private final List executionEngineExtensions; + private final Tracer tracer; /** Default constructor for when no engines are available. */ public OpenSearchPluginModule() { - this(List.of()); + this(List.of(), NoopTracer.INSTANCE); } private final BuiltinFunctionRepository functionRepository = @@ -113,6 +116,12 @@ public SQLService sqlService( return new SQLService(new SQLSyntaxParser(), queryManager, queryPlanFactory, settings); } + @Provides + @Singleton + public Tracer tracer() { + return tracer; + } + /** {@link QueryPlanFactory}. */ @Provides public QueryPlanFactory queryPlanFactory( 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 b0dd7cc1af4..f052a4efc7f 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 @@ -14,6 +14,7 @@ import java.util.Locale; import java.util.Map; import java.util.Optional; +import java.util.function.Consumer; import java.util.function.Supplier; import org.apache.calcite.rel.RelNode; import org.apache.logging.log4j.LogManager; @@ -38,8 +39,10 @@ import org.opensearch.sql.executor.QueryType; import org.opensearch.sql.legacy.metrics.MetricName; import org.opensearch.sql.legacy.metrics.Metrics; +import org.opensearch.sql.monitor.profile.ProfileScope; import org.opensearch.sql.monitor.profile.QueryProfiling; import org.opensearch.sql.opensearch.executor.OpenSearchQueryManager; +import org.opensearch.sql.opensearch.executor.tracing.TracingPhaseListener; import org.opensearch.sql.opensearch.setting.OpenSearchSettings; import org.opensearch.sql.plugin.config.EngineExtensionsHolder; import org.opensearch.sql.plugin.config.OpenSearchPluginModule; @@ -58,6 +61,12 @@ import org.opensearch.sql.protocol.response.format.VisualizationResponseFormatter; import org.opensearch.sql.protocol.response.format.YamlResponseFormatter; import org.opensearch.tasks.Task; +import org.opensearch.telemetry.tracing.Span; +import org.opensearch.telemetry.tracing.SpanCreationContext; +import org.opensearch.telemetry.tracing.SpanScope; +import org.opensearch.telemetry.tracing.Tracer; +import org.opensearch.telemetry.tracing.attributes.Attributes; +import org.opensearch.telemetry.tracing.listener.TraceableActionListener; import org.opensearch.transport.TransportService; import org.opensearch.transport.client.node.NodeClient; @@ -69,6 +78,8 @@ public class TransportPPLQueryAction private final Injector injector; + private final Tracer tracer; + private final Supplier pplEnabled; /** Null when analytics-engine plugin is absent; set via {@link #setQueryPlanExecutor}. */ @@ -86,13 +97,14 @@ public TransportPPLQueryAction( ClusterService clusterService, DataSourceServiceImpl dataSourceService, org.opensearch.common.settings.Settings clusterSettings, - EngineExtensionsHolder extensionsHolder) { + EngineExtensionsHolder extensionsHolder, + Tracer tracer) { super(PPLQueryAction.NAME, transportService, actionFilters, TransportPPLQueryRequest::new); this.clientRef = client; this.clusterServiceRef = clusterService; ModulesBuilder modules = new ModulesBuilder(); - modules.add(new OpenSearchPluginModule(extensionsHolder.engines())); + modules.add(new OpenSearchPluginModule(extensionsHolder.engines(), tracer)); org.opensearch.sql.common.setting.Settings pluginSettings = new OpenSearchSettings(clusterService.getClusterSettings()); this.pluginSettingsRef = pluginSettings; @@ -103,6 +115,8 @@ public TransportPPLQueryAction( b.bind(DataSourceService.class).toInstance(dataSourceService); }); this.injector = Guice.createInjector(modules); + this.tracer = tracer; + ProfileScope.installListener(new TracingPhaseListener(tracer)); this.pplEnabled = () -> MULTI_ALLOW_EXPLICIT_INDEX.get(clusterSettings) @@ -180,52 +194,86 @@ protected void doExecute( // in order to use PPL service, we need to convert TransportPPLQueryRequest to PPLQueryRequest PPLQueryRequest transformedRequest = transportRequest.toPPLQueryRequest(); QueryContext.setProfile(transformedRequest.profile()); - ActionListener clearingListener = wrapWithProfilingClear(listener); - // Route to analytics engine for non-Lucene (e.g., Parquet-backed) indices. - if (unifiedQueryHandler != null - && unifiedQueryHandler.isAnalyticsIndex(transformedRequest.getRequest(), QueryType.PPL)) { - LOG.info("[{}] Routing PPL query to analytics engine", QueryContext.getRequestId()); - // Pass this PPL task so the analytics engine links its query task to it for cancellation. - if (transformedRequest.isExplainRequest()) { - unifiedQueryHandler.explain( - transformedRequest.getRequest(), - QueryType.PPL, - transformedRequest.mode(), - task, - createExplainResponseListener(transformedRequest, clearingListener)); - } else { - // Analytics route only emits JSON; reject unsupported formats (e.g. csv) with a 4xx. - try { - AnalyticsEngineFormatSupport.validateFormat(format(transformedRequest)); - } catch (Exception e) { - clearingListener.onFailure(e); - return; + // Start root span with OTel DB semantic convention attributes + Span rootSpan = + tracer.startSpan( + SpanCreationContext.client() + .name("opensearch.query") + .attributes( + Attributes.create() + .addAttribute("db.system.name", "opensearch") + .addAttribute("db.query.type", "ppl") + .addAttribute("db.query.id", QueryContext.getRequestId()) + .addAttribute( + "db.operation.name", + transformedRequest.isExplainRequest() ? "EXPLAIN" : "EXECUTE"))); + + // Put span in scope so ThreadContext propagation captures it + SpanScope spanScope = tracer.withSpanInScope(rootSpan); + + // Trace wrapper: ends span in async callback, sets error on failure. + ActionListener tracedListener = + TraceableActionListener.create(listener, rootSpan, tracer); + ActionListener clearingListener = + wrapWithProfilingClear(tracedListener); + + try { + // Route to analytics engine for non-Lucene (e.g., Parquet-backed) indices. + if (unifiedQueryHandler != null + && unifiedQueryHandler.isAnalyticsIndex(transformedRequest.getRequest(), QueryType.PPL)) { + LOG.info("[{}] Routing PPL query to analytics engine", QueryContext.getRequestId()); + // Pass this PPL task so the analytics engine links its query task to it for cancellation. + if (transformedRequest.isExplainRequest()) { + unifiedQueryHandler.explain( + transformedRequest.getRequest(), + QueryType.PPL, + transformedRequest.mode(), + task, + createExplainResponseListener(transformedRequest, clearingListener)); + } else { + // Analytics route only emits JSON; reject unsupported formats (e.g. csv) with a 4xx. + try { + AnalyticsEngineFormatSupport.validateFormat(format(transformedRequest)); + } catch (Exception e) { + clearingListener.onFailure(e); + return; + } + unifiedQueryHandler.execute( + transformedRequest.getRequest(), + QueryType.PPL, + transformedRequest.profile(), + transformedRequest.getFetchSize(), + task, + clearingListener); } - unifiedQueryHandler.execute( - transformedRequest.getRequest(), - QueryType.PPL, - transformedRequest.profile(), - transformedRequest.getFetchSize(), - task, - clearingListener); + return; } - return; - } - - PPLService pplService = injector.getInstance(PPLService.class); - if (transformedRequest.isExplainRequest()) { - pplService.explain( - transformedRequest, createExplainResponseListener(transformedRequest, clearingListener)); - } else if (transformedRequest.analyze()) { - pplService.analyze( - transformedRequest, createAnalyzeResponseListener(transformedRequest, clearingListener)); - } else { - pplService.execute( - transformedRequest, - createListener(transformedRequest, clearingListener), - createExplainResponseListener(transformedRequest, clearingListener)); + Consumer anonymizedQuerySink = + anonymized -> rootSpan.addAttribute("db.query.text", anonymized); + PPLService pplService = injector.getInstance(PPLService.class); + if (transformedRequest.isExplainRequest()) { + pplService.explain( + transformedRequest, + createExplainResponseListener(transformedRequest, clearingListener), + anonymizedQuerySink); + } else if (transformedRequest.analyze()) { + pplService.analyze( + transformedRequest, + createAnalyzeResponseListener(transformedRequest, clearingListener), + anonymizedQuerySink); + } else { + pplService.execute( + transformedRequest, + createListener(transformedRequest, clearingListener), + createExplainResponseListener(transformedRequest, clearingListener), + anonymizedQuerySink); + } + } catch (Exception e) { + clearingListener.onFailure(e); + } finally { + spanScope.close(); } } 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 8ff085bfcad..382888274ef 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java @@ -8,6 +8,7 @@ import static org.opensearch.sql.executor.ExecutionEngine.QueryResponse; import static org.opensearch.sql.executor.execution.QueryPlanFactory.NO_CONSUMER_RESPONSE_LISTENER; +import java.util.function.Consumer; import lombok.extern.log4j.Log4j2; import org.antlr.v4.runtime.tree.ParseTree; import org.opensearch.sql.ast.statement.Query; @@ -22,6 +23,7 @@ import org.opensearch.sql.executor.QueryType; import org.opensearch.sql.executor.execution.AbstractPlan; import org.opensearch.sql.executor.execution.QueryPlanFactory; +import org.opensearch.sql.monitor.profile.ProfileScope; import org.opensearch.sql.ppl.antlr.PPLSyntaxParser; import org.opensearch.sql.ppl.domain.PPLQueryRequest; import org.opensearch.sql.ppl.parser.AstBuilder; @@ -31,6 +33,10 @@ /** PPLService. */ @Log4j2 public class PPLService { + + /** Callers that don't care about the anonymized query pass this. */ + public static final Consumer NO_ANONYMIZED_QUERY_SINK = s -> {}; + private final PPLSyntaxParser parser; private final QueryManager queryManager; @@ -66,8 +72,17 @@ public void execute( PPLQueryRequest request, ResponseListener queryListener, ResponseListener explainListener) { + execute(request, queryListener, explainListener, NO_ANONYMIZED_QUERY_SINK); + } + + /** Variant that hands the anonymized query text to {@code anonymizedQuerySink}. */ + public void execute( + PPLQueryRequest request, + ResponseListener queryListener, + ResponseListener explainListener, + Consumer anonymizedQuerySink) { try { - queryManager.submit(plan(request, queryListener, explainListener)); + queryManager.submit(plan(request, queryListener, explainListener, anonymizedQuerySink)); } catch (Exception e) { queryListener.onFailure(e); } @@ -81,8 +96,17 @@ public void execute( * @param listener {@link ResponseListener} for explain response */ public void explain(PPLQueryRequest request, ResponseListener listener) { + explain(request, listener, NO_ANONYMIZED_QUERY_SINK); + } + + /** Variant that hands the anonymized query text to {@code anonymizedQuerySink}. */ + public void explain( + PPLQueryRequest request, + ResponseListener listener, + Consumer anonymizedQuerySink) { try { - queryManager.submit(plan(request, NO_CONSUMER_RESPONSE_LISTENER, listener)); + queryManager.submit( + plan(request, NO_CONSUMER_RESPONSE_LISTENER, listener, anonymizedQuerySink)); } catch (Exception e) { listener.onFailure(e); } @@ -95,29 +119,47 @@ public void explain(PPLQueryRequest request, ResponseListener l * @param listener {@link ResponseListener} for analyze response */ public void analyze(PPLQueryRequest request, ResponseListener listener) { + analyze(request, listener, NO_ANONYMIZED_QUERY_SINK); + } + + /** Variant that hands the anonymized query text to {@code anonymizedQuerySink}. */ + public void analyze( + PPLQueryRequest request, + ResponseListener listener, + Consumer anonymizedQuerySink) { try { String queryText = request.getRequest(); - ParseTree cst = parser.parse(queryText); - Statement statement = - cst.accept( - new AstStatementBuilder( - new AstBuilder(queryText, settings), - 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())); - - log.info( - "[{}] Incoming request {}", - QueryContext.getRequestId(), - anonymizer.anonymizeStatement(statement)); + 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), + 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(); queryManager.submit( @@ -130,34 +172,42 @@ public void analyze(PPLQueryRequest request, ResponseListener l private AbstractPlan plan( PPLQueryRequest request, ResponseListener queryListener, - ResponseListener explainListener) { - // 1.Parse query and convert parse tree (CST) to abstract syntax tree (AST) - ParseTree cst = parser.parse(request.getRequest()); - - boolean includeMetadata = request.getIncludeMetadata(); - - Statement statement = - cst.accept( - new AstStatementBuilder( - new AstBuilder(request.getRequest(), settings), - 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())); - - log.info( - "[{}] Incoming request {}", - QueryContext.getRequestId(), - anonymizer.anonymizeStatement(statement)); + ResponseListener explainListener, + Consumer anonymizedQuerySink) { + 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 { + // 1. Parse query and convert parse tree (CST) to abstract syntax tree (AST) + ParseTree cst = parser.parse(request.getRequest()); + boolean includeMetadata = request.getIncludeMetadata(); + statement = + cst.accept( + new AstStatementBuilder( + new AstBuilder(request.getRequest(), settings), + 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())); + anonymized = anonymizer.anonymizeStatement(statement); + } catch (RuntimeException e) { + preparePhase.setError(e); + throw e; + } + } + log.info("[{}] Incoming request {}", QueryContext.getRequestId(), anonymized); + anonymizedQuerySink.accept(anonymized); return queryExecutionFactory.create(statement, queryListener, explainListener); }