Merge branch 'main' into feat/report-builder

Author: Swiddis

api/README.md

@@ -8,7 +8,8 @@ This module provides components organized into two main areas aligned with the [
 
 ### Unified Language Specification
 
-- **`UnifiedQueryPlanner`**: Accepts PPL (Piped Processing Language) or SQL queries and returns Calcite `RelNode` logical plans as intermediate representation.
+- **`UnifiedQueryParser`**: Parses PPL (Piped Processing Language) or SQL queries and returns the native parse result (`UnresolvedPlan` for PPL, `SqlNode` for Calcite SQL).
+- **`UnifiedQueryPlanner`**: Accepts PPL or SQL queries and returns Calcite `RelNode` logical plans as intermediate representation.
 - **`UnifiedQueryTranspiler`**: Converts Calcite logical plans (`RelNode`) into SQL strings for various target databases using different SQL dialects.
 
 ### Unified Execution Runtime
@@ -42,6 +43,20 @@ UnifiedQueryContext context = UnifiedQueryContext.builder()
     .build();
 ```
 
+### UnifiedQueryParser
+
+Use `UnifiedQueryParser` to parse queries into their native parse tree. The parser is owned by `UnifiedQueryContext` and returns the native parse result for each language.
+
+```java
+// PPL parsing
+UnresolvedPlan ast = (UnresolvedPlan) context.getParser().parse("source = logs | where status = 200");
+
+// SQL parsing (with QueryType.SQL context)
+SqlNode sqlNode = (SqlNode) sqlContext.getParser().parse("SELECT * FROM logs WHERE status = 200");
+```
+
+Callers can then use each language's native visitor infrastructure (`AbstractNodeVisitor` for PPL, `SqlBasicVisitor` for Calcite SQL) on the typed result for further analysis.
+
 ### UnifiedQueryPlanner
 
 Use `UnifiedQueryPlanner` to parse and analyze PPL or SQL queries into Calcite logical plans. The planner accepts a `UnifiedQueryContext` and can be reused for multiple queries.
@@ -179,6 +194,59 @@ try (UnifiedQueryContext context = UnifiedQueryContext.builder()
 }
 ```
 
+## Profiling
+
+The unified query API supports the same [profiling capability](../docs/user/ppl/interfaces/endpoint.md#profile-experimental) as the PPL REST endpoint. When enabled, each unified query component automatically collects per-phase timing metrics. For code outside unified query components (e.g., `PreparedStatement.executeQuery()` or response formatting), `context.measure()` records custom phases into the same profile.
+
+```java
+try (UnifiedQueryContext context = UnifiedQueryContext.builder()
+    .language(QueryType.PPL)
+    .catalog("catalog", schema)
+    .defaultNamespace("catalog")
+    .profiling(true)
+    .build()) {
+
+  // Auto-profiled: ANALYZE
+  RelNode plan = new UnifiedQueryPlanner(context).plan(query);
+
+  // Auto-profiled: OPTIMIZE
+  PreparedStatement stmt = new UnifiedQueryCompiler(context).compile(plan);
+
+  // User-profiled via measure()
+  ResultSet rs = context.measure(MetricName.EXECUTE, stmt::executeQuery);
+  String json = context.measure(MetricName.FORMAT, () -> formatter.format(result));
+
+  // Retrieve profile snapshot
+  QueryProfile profile = context.getProfile();
+}
+```
+
+The returned `QueryProfile` follows the same JSON structure as the REST API:
+
+```json
+{
+  "summary": {
+    "total_time_ms": 33.34
+  },
+  "phases": {
+    "analyze": { "time_ms": 8.68 },
+    "optimize": { "time_ms": 18.2 },
+    "execute": { "time_ms": 4.87 },
+    "format": { "time_ms": 0.05 }
+  },
+  "plan": {
+    "node": "EnumerableCalc",
+    "time_ms": 4.82,
+    "rows": 2,
+    "children": [
+      { "node": "CalciteEnumerableIndexScan", "time_ms": 4.12, "rows": 2 }
+    ]
+  }
+}
+```
+
+When profiling is disabled (the default), all components execute with zero overhead.
+
 ## Development & Testing
 
 A set of unit tests is provided to validate planner behavior.

api/src/main/java/org/opensearch/sql/api/UnifiedQueryContext.java

@@ -13,7 +13,10 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import lombok.Value;
+import java.util.Optional;
+import java.util.concurrent.Callable;
+import lombok.AllArgsConstructor;
+import lombok.Getter;
 import org.apache.calcite.avatica.util.Casing;
 import org.apache.calcite.jdbc.CalciteSchema;
 import org.apache.calcite.plan.RelTraitDef;
@@ -24,24 +27,64 @@
 import org.apache.calcite.tools.FrameworkConfig;
 import org.apache.calcite.tools.Frameworks;
 import org.apache.calcite.tools.Programs;
+import org.opensearch.sql.api.parser.CalciteSqlQueryParser;
+import org.opensearch.sql.api.parser.PPLQueryParser;
+import org.opensearch.sql.api.parser.UnifiedQueryParser;
 import org.opensearch.sql.calcite.CalcitePlanContext;
 import org.opensearch.sql.calcite.SysLimit;
 import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.executor.QueryType;
+import org.opensearch.sql.monitor.profile.MetricName;
+import org.opensearch.sql.monitor.profile.ProfileMetric;
+import org.opensearch.sql.monitor.profile.QueryProfile;
+import org.opensearch.sql.monitor.profile.QueryProfiling;
 
 /**
  * A reusable abstraction shared across unified query components (planner, compiler, etc.). This
  * centralizes configuration for catalog schemas, query type, execution limits, and other settings,
  * enabling consistent behavior across all unified query operations.
  */
-@Value
+@AllArgsConstructor
+@Getter
 public class UnifiedQueryContext implements AutoCloseable {
 
   /** CalcitePlanContext containing Calcite framework configuration and query type. */
-  CalcitePlanContext planContext;
+  private final CalcitePlanContext planContext;
 
   /** Settings containing execution limits and feature flags used by parsers and planners. */
-  Settings settings;
+  private final Settings settings;
+
+  /** Query parser created eagerly from this context's configuration. */
+  private final UnifiedQueryParser<?> parser;
+
+  /**
+   * Returns the profiling result. Call after query execution to retrieve collected metrics. Returns
+   * empty if profiling was not enabled.
+   */
+  public Optional<QueryProfile> getProfile() {
+    return Optional.ofNullable(QueryProfiling.current().finish());
+  }
+
+  /**
+   * Measures the execution time of the given action and records it as a profiling metric. When
+   * profiling is disabled, the action executes with no overhead. Use this for phases outside
+   * unified query components (e.g., execution, formatting).
+   *
+   * @param <T> the return type of the action
+   * @param metricName the metric to record
+   * @param action the action to measure
+   * @return the result of the action
+   * @throws Exception if the action throws
+   */
+  public <T> T measure(MetricName metricName, Callable<T> action) throws Exception {
+    ProfileMetric metric = QueryProfiling.current().getOrCreateMetric(metricName);
+    long start = System.nanoTime();
+    try {
+      return action.call();
+    } finally {
+      metric.set(System.nanoTime() - start);
+    }
+  }
 
   /**
    * Closes the underlying resource managed by this context.
@@ -50,6 +93,7 @@ public class UnifiedQueryContext implements AutoCloseable {
    */
   @Override
   public void close() throws Exception {
+    QueryProfiling.clear();
     if (planContext != null && planContext.connection != null) {
       planContext.connection.close();
     }
@@ -66,6 +110,7 @@ public static class Builder {
     private final Map<String, Schema> catalogs = new HashMap<>();
     private String defaultNamespace;
     private boolean cacheMetadata = false;
+    private boolean profiling = false;
 
     /**
      * Setting values with defaults from SysLimit.DEFAULT. Only includes planning-required settings
@@ -125,6 +170,18 @@ public Builder cacheMetadata(boolean cache) {
       return this;
     }
 
+    /**
+     * Enables or disables query profiling. When enabled, profiling metrics are collected during
+     * query planning and execution, retrievable via {@link UnifiedQueryContext#getProfile()}.
+     *
+     * @param enabled whether to enable profiling
+     * @return this builder instance
+     */
+    public Builder profiling(boolean enabled) {
+      this.profiling = enabled;
+      return this;
+    }
+
     /**
      * Sets a specific setting value by name.
      *
@@ -152,7 +209,15 @@ public UnifiedQueryContext build() {
       CalcitePlanContext planContext =
           CalcitePlanContext.create(
               buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType);
-      return new UnifiedQueryContext(planContext, settings);
+      QueryProfiling.activate(profiling);
+      return new UnifiedQueryContext(planContext, settings, createParser(planContext, settings));
+    }
+
+    private UnifiedQueryParser<?> createParser(CalcitePlanContext planContext, Settings settings) {
+      return switch (queryType) {
+        case PPL -> new PPLQueryParser(settings);
+        case SQL -> new CalciteSqlQueryParser(planContext);
+      };
     }
 
     private Settings buildSettings() {

api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java

@@ -5,8 +5,9 @@
 
 package org.opensearch.sql.api;
 
+import static org.opensearch.sql.monitor.profile.MetricName.ANALYZE;
+
 import lombok.RequiredArgsConstructor;
-import org.antlr.v4.runtime.tree.ParseTree;
 import org.apache.calcite.rel.RelCollation;
 import org.apache.calcite.rel.RelCollations;
 import org.apache.calcite.rel.RelNode;
@@ -16,15 +17,11 @@
 import org.apache.calcite.sql.SqlNode;
 import org.apache.calcite.tools.Frameworks;
 import org.apache.calcite.tools.Planner;
-import org.opensearch.sql.ast.statement.Query;
-import org.opensearch.sql.ast.statement.Statement;
+import org.opensearch.sql.api.parser.UnifiedQueryParser;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.calcite.CalciteRelNodeVisitor;
 import org.opensearch.sql.common.antlr.SyntaxCheckException;
 import org.opensearch.sql.executor.QueryType;
-import org.opensearch.sql.ppl.antlr.PPLSyntaxParser;
-import org.opensearch.sql.ppl.parser.AstBuilder;
-import org.opensearch.sql.ppl.parser.AstStatementBuilder;
 
 /**
  * {@code UnifiedQueryPlanner} provides a high-level API for parsing and analyzing queries using the
@@ -36,12 +33,16 @@ public class UnifiedQueryPlanner {
   /** Planning strategy selected at construction time based on query type. */
   private final PlanningStrategy strategy;
 
+  /** Unified query context for profiling support. */
+  private final UnifiedQueryContext context;
+
   /**
    * Constructs a UnifiedQueryPlanner with a unified query context.
    *
    * @param context the unified query context containing CalcitePlanContext
    */
   public UnifiedQueryPlanner(UnifiedQueryContext context) {
+    this.context = context;
     this.strategy =
         context.getPlanContext().queryType == QueryType.SQL
             ? new CalciteNativeStrategy(context)
@@ -57,7 +58,7 @@ public UnifiedQueryPlanner(UnifiedQueryContext context) {
    */
   public RelNode plan(String query) {
     try {
-      return strategy.plan(query);
+      return context.measure(ANALYZE, () -> strategy.plan(query));
     } catch (SyntaxCheckException e) {
       // Re-throw syntax error without wrapping
       throw e;
@@ -87,36 +88,26 @@ public RelNode plan(String query) throws Exception {
     }
   }
 
-  /** AST-based planning via ANTLR parser → UnresolvedPlan → CalciteRelNodeVisitor. */
-  @RequiredArgsConstructor
+  /** AST-based planning via context-owned parser → UnresolvedPlan → CalciteRelNodeVisitor. */
   private static class CustomVisitorStrategy implements PlanningStrategy {
     private final UnifiedQueryContext context;
-    private final PPLSyntaxParser parser = new PPLSyntaxParser();
+    private final UnifiedQueryParser<UnresolvedPlan> parser;
     private final CalciteRelNodeVisitor relNodeVisitor =
         new CalciteRelNodeVisitor(new EmptyDataSourceService());
 
+    @SuppressWarnings("unchecked")
+    CustomVisitorStrategy(UnifiedQueryContext context) {
+      this.context = context;
+      this.parser = (UnifiedQueryParser<UnresolvedPlan>) context.getParser();
+    }
+
     @Override
     public RelNode plan(String query) {
-      UnresolvedPlan ast = parse(query);
+      UnresolvedPlan ast = parser.parse(query);
       RelNode logical = relNodeVisitor.analyze(ast, context.getPlanContext());
       return preserveCollation(logical);
     }
 
-    private UnresolvedPlan parse(String query) {
-      ParseTree cst = parser.parse(query);
-      AstStatementBuilder astStmtBuilder =
-          new AstStatementBuilder(
-              new AstBuilder(query, context.getSettings()),
-              AstStatementBuilder.StatementBuilderContext.builder().build());
-      Statement statement = cst.accept(astStmtBuilder);
-
-      if (statement instanceof Query) {
-        return ((Query) statement).getPlan();
-      }
-      throw new UnsupportedOperationException(
-          "Only query statements are supported but got " + statement.getClass().getSimpleName());
-    }
-
     private RelNode preserveCollation(RelNode logical) {
       RelCollation collation = logical.getTraitSet().getCollation();
       if (!(logical instanceof Sort) && collation != RelCollations.EMPTY) {

api/src/main/java/org/opensearch/sql/api/compiler/UnifiedQueryCompiler.java

@@ -5,6 +5,8 @@
 
 package org.opensearch.sql.api.compiler;
 
+import static org.opensearch.sql.monitor.profile.MetricName.OPTIMIZE;
+
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import lombok.NonNull;
@@ -46,26 +48,29 @@ public UnifiedQueryCompiler(UnifiedQueryContext context) {
    */
   public PreparedStatement compile(@NonNull RelNode plan) {
     try {
-      // Apply shuttle to convert LogicalTableScan to BindableTableScan
-      final RelHomogeneousShuttle shuttle =
-          new RelHomogeneousShuttle() {
-            @Override
-            public RelNode visit(TableScan scan) {
-              final RelOptTable table = scan.getTable();
-              if (scan instanceof LogicalTableScan
-                  && Bindables.BindableTableScan.canHandle(table)) {
-                return Bindables.BindableTableScan.create(scan.getCluster(), table);
-              }
-              return super.visit(scan);
-            }
-          };
-      RelNode transformedPlan = plan.accept(shuttle);
-
-      Connection connection = context.getPlanContext().connection;
-      final RelRunner runner = connection.unwrap(RelRunner.class);
-      return runner.prepareStatement(transformedPlan);
+      return context.measure(OPTIMIZE, () -> doCompile(plan));
     } catch (Exception e) {
       throw new IllegalStateException("Failed to compile logical plan", e);
     }
   }
+
+  private PreparedStatement doCompile(RelNode plan) throws Exception {
+    // Apply shuttle to convert LogicalTableScan to BindableTableScan
+    final RelHomogeneousShuttle shuttle =
+        new RelHomogeneousShuttle() {
+          @Override
+          public RelNode visit(TableScan scan) {
+            final RelOptTable table = scan.getTable();
+            if (scan instanceof LogicalTableScan && Bindables.BindableTableScan.canHandle(table)) {
+              return Bindables.BindableTableScan.create(scan.getCluster(), table);
+            }
+            return super.visit(scan);
+          }
+        };
+    RelNode transformedPlan = plan.accept(shuttle);
+
+    Connection connection = context.getPlanContext().connection;
+    final RelRunner runner = connection.unwrap(RelRunner.class);
+    return runner.prepareStatement(transformedPlan);
+  }
 }