Wrap execute and explain with context.measure() for profiling

ahkcs · ahkcs · commit 2545d34fac98 · 2026-03-30T15:43:57.000-07:00
Wrap analyticsEngine.execute() and analyticsEngine.explain() calls
with context.measure(MetricName.EXECUTE, ...) so execution time is
captured in the profiling metrics. Planning is auto-profiled by
UnifiedQueryPlanner.

Signed-off-by: Kai Huang &lt;kaihuang@amazon.com&gt;
Signed-off-by: Kai Huang &lt;ahkcs@amazon.com&gt;
diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/rest/RestUnifiedQueryAction.java
@@ -31,6 +31,7 @@
 import org.opensearch.sql.executor.analytics.AnalyticsExecutionEngine;
 import org.opensearch.sql.executor.analytics.QueryPlanExecutor;
 import org.opensearch.sql.lang.LangSpec;
+import org.opensearch.sql.monitor.profile.MetricName;
 import org.opensearch.sql.plugin.rest.analytics.stub.StubSchemaProvider;
 import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse;
 import org.opensearch.sql.ppl.domain.PPLQueryRequest;
@@ -122,7 +123,14 @@ private void doExecute(
       CalcitePlanContext planContext = context.getPlanContext();
       plan = addQuerySizeLimit(plan, planContext);
 
-      analyticsEngine.execute(plan, planContext, createQueryListener(queryType, listener));
+      RelNode finalPlan = plan;
+      context.measure(
+          MetricName.EXECUTE,
+          () -> {
+            analyticsEngine.execute(
+                finalPlan, planContext, createQueryListener(queryType, listener));
+            return null;
+          });
     } catch (Exception e) {
       listener.onFailure(e);
     }
@@ -140,7 +148,13 @@ private void doExplain(
       CalcitePlanContext planContext = context.getPlanContext();
       plan = addQuerySizeLimit(plan, planContext);
 
-      analyticsEngine.explain(plan, pplRequest.mode(), planContext, listener);
+      RelNode finalPlan = plan;
+      context.measure(
+          MetricName.EXECUTE,
+          () -> {
+            analyticsEngine.explain(finalPlan, pplRequest.mode(), planContext, listener);
+            return null;
+          });
     } catch (Exception e) {
       listener.onFailure(e);
     }
