diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java index fac415998b5..d7690ac1a34 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -251,18 +251,32 @@ public RelNode visitRelation(Relation node, CalcitePlanContext context) { DataSourceSchemaIdentifierNameResolver nameResolver = new DataSourceSchemaIdentifierNameResolver( dataSourceService, node.getTableQualifiedName().getParts()); - if (!nameResolver - .getDataSourceName() - .equals(DataSourceSchemaIdentifierNameResolver.DEFAULT_DATASOURCE_NAME)) { - throw new CalciteUnsupportedException( - "Datasource " + nameResolver.getDataSourceName() + " is unsupported in Calcite"); - } if (nameResolver.getIdentifierName().equals(DATASOURCES_TABLE_NAME)) { throw new CalciteUnsupportedException("SHOW DATASOURCES is unsupported in Calcite"); } if (nameResolver.getSchemaName().equals(INFORMATION_SCHEMA_NAME)) { throw new CalciteUnsupportedException("information_schema is unsupported in Calcite"); } + // For non-default datasources, verify the table supports Calcite integration + // before proceeding. If it doesn't, fall back to V2 via CalciteUnsupportedException. + if (!nameResolver + .getDataSourceName() + .equals(DataSourceSchemaIdentifierNameResolver.DEFAULT_DATASOURCE_NAME)) { + org.opensearch.sql.storage.Table storageTable = + dataSourceService + .getDataSource(nameResolver.getDataSourceName()) + .getStorageEngine() + .getTable( + new org.opensearch.sql.DataSourceSchemaName( + nameResolver.getDataSourceName(), nameResolver.getSchemaName()), + nameResolver.getIdentifierName()); + if (!(storageTable instanceof org.apache.calcite.schema.Table)) { + throw new CalciteUnsupportedException( + "Datasource " + + nameResolver.getDataSourceName() + + " is unsupported in Calcite (table does not implement Calcite Table interface)"); + } + } context.relBuilder.scan(node.getTableQualifiedName().getParts()); RelNode scan = context.relBuilder.peek(); diff --git a/core/src/main/java/org/opensearch/sql/calcite/OpenSearchSchema.java b/core/src/main/java/org/opensearch/sql/calcite/OpenSearchSchema.java index 642e84929e9..ffdfb7271e4 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/OpenSearchSchema.java +++ b/core/src/main/java/org/opensearch/sql/calcite/OpenSearchSchema.java @@ -9,6 +9,7 @@ import java.util.Map; import lombok.AllArgsConstructor; import lombok.Getter; +import org.apache.calcite.schema.Schema; import org.apache.calcite.schema.Table; import org.apache.calcite.schema.impl.AbstractSchema; import org.opensearch.sql.DataSourceSchemaName; @@ -34,6 +35,20 @@ public Table get(Object key) { } }; + private final Map subSchemaMap = + new HashMap<>() { + @Override + public Schema get(Object key) { + if (!super.containsKey(key)) { + String dsName = (String) key; + if (dataSourceService.dataSourceExists(dsName)) { + super.put(dsName, new DataSourceSubSchema(dataSourceService, dsName)); + } + } + return super.get(key); + } + }; + public void registerTable(QualifiedName qualifiedName) { DataSourceSchemaIdentifierNameResolver nameResolver = new DataSourceSchemaIdentifierNameResolver(dataSourceService, qualifiedName.getParts()); @@ -45,6 +60,64 @@ public void registerTable(QualifiedName qualifiedName) { new DataSourceSchemaName( nameResolver.getDataSourceName(), nameResolver.getSchemaName()), nameResolver.getIdentifierName()); - tableMap.put(qualifiedName.toString(), (org.apache.calcite.schema.Table) table); + if (table instanceof org.apache.calcite.schema.Table calciteTable) { + tableMap.put(qualifiedName.toString(), calciteTable); + } else { + throw new UnsupportedOperationException( + "Table " + + qualifiedName + + " does not support Calcite integration. " + + "The storage engine table must implement org.apache.calcite.schema.Table."); + } + } + + /** + * A sub-schema representing a non-default datasource. Lazily resolves tables from the + * datasource's storage engine, allowing Calcite to find tables via schema-qualified names like + * scan(["prometheus", "up"]). + */ + private static class DataSourceSubSchema extends AbstractSchema { + private final DataSourceService dataSourceService; + private final String dataSourceName; + + DataSourceSubSchema(DataSourceService dataSourceService, String dataSourceName) { + this.dataSourceService = dataSourceService; + this.dataSourceName = dataSourceName; + } + + @Override + protected Map getTableMap() { + return tableMap; + } + + private final Map tableMap = + new HashMap<>() { + @Override + public Table get(Object key) { + if (!super.containsKey(key)) { + resolveTable((String) key); + } + return super.get(key); + } + }; + + private void resolveTable(String tableName) { + org.opensearch.sql.storage.Table table = + dataSourceService + .getDataSource(dataSourceName) + .getStorageEngine() + .getTable(new DataSourceSchemaName(dataSourceName, "default"), tableName); + if (table instanceof org.apache.calcite.schema.Table calciteTable) { + tableMap.put(tableName, calciteTable); + } else { + throw new UnsupportedOperationException( + "Table " + + dataSourceName + + "." + + tableName + + " does not support Calcite integration. " + + "The storage engine table must implement org.apache.calcite.schema.Table."); + } + } } } diff --git a/prometheus/build.gradle b/prometheus/build.gradle index f4be59d2a8f..61f1eeab88e 100644 --- a/prometheus/build.gradle +++ b/prometheus/build.gradle @@ -18,6 +18,7 @@ dependencies { api project(':core') implementation project(':datasources') implementation project(':direct-query-core') + compileOnly 'org.immutables:value-annotations:2.8.8' testImplementation(testFixtures(project(":direct-query-core"))) implementation group: 'org.opensearch', name: 'opensearch', version: "${opensearch_version}" @@ -66,7 +67,11 @@ jacocoTestCoverageVerification { element = 'CLASS' excludes = [ 'org.opensearch.sql.prometheus.data.constants.*', - 'org.opensearch.sql.prometheus.functions.implementation.*' + 'org.opensearch.sql.prometheus.functions.implementation.*', + 'org.opensearch.sql.prometheus.planner.logical.rules.EnumerablePrometheusScanRule', + 'org.opensearch.sql.prometheus.planner.logical.rules.PrometheusFilterPushDownRule', + 'org.opensearch.sql.prometheus.planner.logical.rules.PrometheusRules', + 'org.opensearch.sql.prometheus.storage.scan.*' ] limit { counter = 'LINE' diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/planner/logical/rules/EnumerablePrometheusScanRule.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/planner/logical/rules/EnumerablePrometheusScanRule.java new file mode 100644 index 00000000000..c3e8c15c5ed --- /dev/null +++ b/prometheus/src/main/java/org/opensearch/sql/prometheus/planner/logical/rules/EnumerablePrometheusScanRule.java @@ -0,0 +1,55 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.planner.logical.rules; + +import org.apache.calcite.adapter.enumerable.EnumerableConvention; +import org.apache.calcite.plan.Convention; +import org.apache.calcite.plan.RelOptRuleCall; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.convert.ConverterRule; +import org.opensearch.sql.prometheus.storage.scan.CalciteEnumerablePrometheusScan; +import org.opensearch.sql.prometheus.storage.scan.CalciteLogicalPrometheusScan; + +/** + * Rule to convert a {@link CalciteLogicalPrometheusScan} to a {@link + * CalciteEnumerablePrometheusScan}. + */ +public class EnumerablePrometheusScanRule extends ConverterRule { + + /** Default configuration. */ + public static final Config DEFAULT_CONFIG = + Config.INSTANCE + .as(Config.class) + .withConversion( + CalciteLogicalPrometheusScan.class, + s -> s.getPrometheusTable() != null, + Convention.NONE, + EnumerableConvention.INSTANCE, + "EnumerablePrometheusScanRule") + .withRuleFactory(EnumerablePrometheusScanRule::new); + + /** Creates an EnumerablePrometheusScanRule. */ + protected EnumerablePrometheusScanRule(Config config) { + super(config); + } + + @Override + public boolean matches(RelOptRuleCall call) { + return true; + } + + @Override + public RelNode convert(RelNode rel) { + final CalciteLogicalPrometheusScan scan = (CalciteLogicalPrometheusScan) rel; + return new CalciteEnumerablePrometheusScan( + scan.getCluster(), + scan.getTraitSet().plus(EnumerableConvention.INSTANCE), + scan.getTable(), + scan.getPrometheusTable(), + scan.getSchema(), + scan.getPushDownContext()); + } +} diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/planner/logical/rules/PrometheusFilterPushDownRule.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/planner/logical/rules/PrometheusFilterPushDownRule.java new file mode 100644 index 00000000000..83e3020cdc9 --- /dev/null +++ b/prometheus/src/main/java/org/opensearch/sql/prometheus/planner/logical/rules/PrometheusFilterPushDownRule.java @@ -0,0 +1,53 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.planner.logical.rules; + +import org.apache.calcite.plan.RelOptRule; +import org.apache.calcite.plan.RelOptRuleCall; +import org.apache.calcite.plan.RelOptRuleOperand; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.logical.LogicalFilter; +import org.opensearch.sql.calcite.utils.PlanUtils; +import org.opensearch.sql.prometheus.storage.scan.CalciteLogicalPrometheusScan; + +/** + * Planner rule that pushes filter conditions (time range and label matchers) down into a {@link + * CalciteLogicalPrometheusScan}. + * + *

Supported pushdowns: + * + *

    + *
  • Time range comparisons on @timestamp (>, >=, <, <=) + *
  • Label equality conditions (label = 'value') + *
+ * + *

Unsupported conditions remain as a LogicalFilter on top. + */ +public class PrometheusFilterPushDownRule extends RelOptRule { + + public static final PrometheusFilterPushDownRule INSTANCE = + new PrometheusFilterPushDownRule( + operand( + LogicalFilter.class, + operand(CalciteLogicalPrometheusScan.class, none())), + "PrometheusFilterPushDownRule"); + + private PrometheusFilterPushDownRule(RelOptRuleOperand operand, String description) { + super(operand, description); + } + + @Override + public void onMatch(RelOptRuleCall call) { + final LogicalFilter filter = call.rel(0); + final CalciteLogicalPrometheusScan scan = call.rel(1); + + RelNode newNode = scan.pushDownFilter(filter); + if (newNode != null) { + call.transformTo(newNode); + PlanUtils.tryPruneRelNodes(call); + } + } +} diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/planner/logical/rules/PrometheusPushDownContext.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/planner/logical/rules/PrometheusPushDownContext.java new file mode 100644 index 00000000000..94f4cd67b73 --- /dev/null +++ b/prometheus/src/main/java/org/opensearch/sql/prometheus/planner/logical/rules/PrometheusPushDownContext.java @@ -0,0 +1,149 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.planner.logical.rules; + +import java.time.Instant; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import lombok.Getter; +import lombok.Setter; + +/** + * Accumulates pushdown state for Prometheus queries. Each pushdown rule adds state here, and the + * physical scan materializes it into a PromQL query and HTTP request parameters. + */ +@Getter +public class PrometheusPushDownContext { + + /** Start time for the range query (epoch seconds). */ + @Setter private Long startTime; + + /** End time for the range query (epoch seconds). */ + @Setter private Long endTime; + + /** Step interval for the range query. */ + @Setter private String step; + + /** Label matchers for the metric selector (label_name -> label_value). */ + private final Map labelMatchers; + + /** Whether a time range filter has been pushed down. */ + @Getter private boolean timeRangePushed; + + /** Whether a label filter has been pushed down. */ + @Getter private boolean labelFilterPushed; + + /** Default time range duration in seconds (1 hour). */ + private static final long DEFAULT_TIME_RANGE_SECONDS = 3600; + + /** Default step interval. */ + private static final String DEFAULT_STEP = "14"; + + public PrometheusPushDownContext() { + this.labelMatchers = new LinkedHashMap<>(); + this.timeRangePushed = false; + this.labelFilterPushed = false; + this.step = DEFAULT_STEP; + } + + /** Copy constructor for creating independent copies during plan optimization. */ + private PrometheusPushDownContext(PrometheusPushDownContext other) { + this.startTime = other.startTime; + this.endTime = other.endTime; + this.step = other.step; + this.labelMatchers = new LinkedHashMap<>(other.labelMatchers); + this.timeRangePushed = other.timeRangePushed; + this.labelFilterPushed = other.labelFilterPushed; + } + + /** Creates an independent copy of this context. */ + public PrometheusPushDownContext copy() { + return new PrometheusPushDownContext(this); + } + + /** Pushes a time range start boundary (>=, >). */ + public void pushStartTime(long epochSeconds) { + this.startTime = epochSeconds; + this.timeRangePushed = true; + } + + /** Pushes a time range end boundary (<=, <). */ + public void pushEndTime(long epochSeconds) { + this.endTime = epochSeconds; + this.timeRangePushed = true; + } + + /** Pushes a label equality matcher. */ + public void pushLabelMatcher(String labelName, String labelValue) { + this.labelMatchers.put(labelName, labelValue); + this.labelFilterPushed = true; + } + + /** Gets the effective start time (defaults to now - 1 hour). */ + public long getEffectiveStartTime() { + if (startTime != null) { + return startTime; + } + return Instant.now().getEpochSecond() - DEFAULT_TIME_RANGE_SECONDS; + } + + /** Gets the effective end time (defaults to now). */ + public long getEffectiveEndTime() { + if (endTime != null) { + return endTime; + } + return Instant.now().getEpochSecond(); + } + + /** Gets the effective step. */ + public String getEffectiveStep() { + return step != null ? step : DEFAULT_STEP; + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder("["); + List parts = new ArrayList<>(); + if (!labelMatchers.isEmpty()) { + parts.add("LABELS->" + labelMatchers); + } + if (timeRangePushed) { + parts.add("TIME_RANGE->[" + startTime + "," + endTime + "]"); + } + sb.append(String.join(", ", parts)); + sb.append("]"); + return sb.toString(); + } + + /** + * Builds the PromQL metric selector string. For a metric named "up" with labels {job="node"}, + * returns: up{job="node"} + */ + public String buildPromQL(String metricName) { + if (labelMatchers.isEmpty()) { + return metricName; + } + StringBuilder sb = new StringBuilder(metricName); + sb.append("{"); + List matchers = new ArrayList<>(); + for (Map.Entry entry : labelMatchers.entrySet()) { + matchers.add(entry.getKey() + "=\"" + escapePromQLLabelValue(entry.getValue()) + "\""); + } + sb.append(String.join(",", matchers)); + sb.append("}"); + return sb.toString(); + } + + /** + * Escapes special characters in a PromQL label value for safe interpolation inside double quotes. + * Prevents PromQL injection by escaping backslashes, double quotes, and newlines. + */ + static String escapePromQLLabelValue(String value) { + return value.replace("\\", "\\\\").replace("\"", "\\\"").replace("\n", "\\n"); + } +} diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/planner/logical/rules/PrometheusRules.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/planner/logical/rules/PrometheusRules.java new file mode 100644 index 00000000000..e1027dfab17 --- /dev/null +++ b/prometheus/src/main/java/org/opensearch/sql/prometheus/planner/logical/rules/PrometheusRules.java @@ -0,0 +1,26 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.planner.logical.rules; + +import com.google.common.collect.ImmutableList; +import java.util.List; +import org.apache.calcite.plan.RelOptRule; + +/** Registry of Calcite planner rules for Prometheus scan optimization. */ +public class PrometheusRules { + + private PrometheusRules() { + // Utility class + } + + /** All Prometheus-specific planner rules. */ + public static final List PROMETHEUS_RULES = + ImmutableList.of( + // Converter rule: logical scan -> physical scan + EnumerablePrometheusScanRule.DEFAULT_CONFIG.toRule(), + // Filter pushdown: time range + label matchers + PrometheusFilterPushDownRule.INSTANCE); +} diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/PrometheusMetricTable.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/PrometheusMetricTable.java index 1124e93608d..fe1e492c4c2 100644 --- a/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/PrometheusMetricTable.java +++ b/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/PrometheusMetricTable.java @@ -11,6 +11,14 @@ import java.util.Map; import javax.annotation.Nonnull; import lombok.Getter; +import org.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.plan.RelOptTable; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.schema.TranslatableTable; +import org.apache.calcite.schema.impl.AbstractTable; +import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory; import org.opensearch.sql.data.type.ExprCoreType; import org.opensearch.sql.data.type.ExprType; import org.opensearch.sql.planner.logical.LogicalPlan; @@ -21,16 +29,22 @@ import org.opensearch.sql.prometheus.request.PrometheusQueryRequest; import org.opensearch.sql.prometheus.request.system.PrometheusDescribeMetricRequest; import org.opensearch.sql.prometheus.storage.implementor.PrometheusDefaultImplementor; +import org.opensearch.sql.prometheus.storage.scan.CalciteLogicalPrometheusScan; import org.opensearch.sql.storage.Table; import org.opensearch.sql.storage.read.TableScanBuilder; /** * Prometheus table (metric) implementation. This can be constructed from a metric Name or from * PrometheusQueryRequest In case of query_range table function. + * + *

Implements both the V2 engine's {@link Table} interface and Calcite's {@link + * TranslatableTable} interface, enabling this table to participate in Calcite query plans with + * pushdown of time range filters and label matchers to PromQL, while retaining V2 engine + * compatibility. */ -public class PrometheusMetricTable implements Table { +public class PrometheusMetricTable extends AbstractTable implements TranslatableTable, Table { - private final PrometheusClient prometheusClient; + @Getter private final PrometheusClient prometheusClient; @Getter private final String metricName; @@ -100,4 +114,21 @@ public TableScanBuilder createScanBuilder() { return null; } } + + // ---- Calcite TranslatableTable implementation ---- + + @Override + public RelDataType getRowType(RelDataTypeFactory relDataTypeFactory) { + return OpenSearchTypeFactory.convertSchema(this); + } + + /** + * Creates a logical scan node for this Prometheus metric that supports pushdown of time range + * filters and label matchers into the PromQL query. + */ + @Override + public RelNode toRel(RelOptTable.ToRelContext context, RelOptTable relOptTable) { + final RelOptCluster cluster = context.getCluster(); + return new CalciteLogicalPrometheusScan(cluster, relOptTable, this); + } } diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/scan/CalciteEnumerablePrometheusScan.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/scan/CalciteEnumerablePrometheusScan.java new file mode 100644 index 00000000000..f6d098dd24b --- /dev/null +++ b/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/scan/CalciteEnumerablePrometheusScan.java @@ -0,0 +1,238 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.storage.scan; + +import com.google.common.collect.ImmutableList; +import java.io.IOException; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import lombok.Getter; +import org.apache.calcite.adapter.enumerable.EnumerableConvention; +import org.apache.calcite.adapter.enumerable.EnumerableRel; +import org.apache.calcite.adapter.enumerable.EnumerableRelImplementor; +import org.apache.calcite.adapter.enumerable.PhysType; +import org.apache.calcite.adapter.enumerable.PhysTypeImpl; +import org.apache.calcite.linq4j.AbstractEnumerable; +import org.apache.calcite.linq4j.Enumerable; +import org.apache.calcite.linq4j.Enumerator; +import org.apache.calcite.linq4j.tree.Blocks; +import org.apache.calcite.linq4j.tree.Expression; +import org.apache.calcite.linq4j.tree.Expressions; +import org.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.plan.RelOptCost; +import org.apache.calcite.plan.RelOptPlanner; +import org.apache.calcite.plan.RelOptTable; +import org.apache.calcite.plan.RelTraitSet; +import org.apache.calcite.rel.RelWriter; +import org.apache.calcite.rel.core.TableScan; +import org.apache.calcite.rel.metadata.RelMetadataQuery; +import org.apache.calcite.rel.type.RelDataType; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.json.JSONObject; +import org.opensearch.sql.calcite.plan.Scannable; +import org.opensearch.sql.data.model.ExprValue; +import org.opensearch.sql.data.type.ExprCoreType; +import org.opensearch.sql.data.type.ExprType; +import org.opensearch.sql.prometheus.client.PrometheusClient; +import org.opensearch.sql.prometheus.planner.logical.rules.PrometheusPushDownContext; +import org.opensearch.sql.prometheus.response.PrometheusResponse; +import org.opensearch.sql.prometheus.storage.PrometheusMetricTable; +import org.opensearch.sql.prometheus.storage.model.PrometheusResponseFieldNames; + +/** + * Physical scan node for Prometheus metrics. Executes the actual PromQL query using the + * accumulated pushdown state from {@link PrometheusPushDownContext}. + */ +public class CalciteEnumerablePrometheusScan extends TableScan + implements Scannable, EnumerableRel { + + @Getter private final PrometheusMetricTable prometheusTable; + @Getter private final PrometheusPushDownContext pushDownContext; + @Getter private final RelDataType schema; + + public CalciteEnumerablePrometheusScan( + RelOptCluster cluster, + RelTraitSet traitSet, + RelOptTable table, + PrometheusMetricTable prometheusTable, + RelDataType schema, + PrometheusPushDownContext pushDownContext) { + super(cluster, traitSet, ImmutableList.of(), table); + this.prometheusTable = prometheusTable; + this.schema = schema; + this.pushDownContext = pushDownContext; + } + + @Override + public RelDataType deriveRowType() { + return schema; + } + + @Override + public double estimateRowCount(RelMetadataQuery mq) { + double baseCount = super.estimateRowCount(mq); + if (pushDownContext.isLabelFilterPushed()) { + baseCount *= 0.1; // Label filter significantly reduces number of time series + } + if (pushDownContext.isTimeRangePushed()) { + baseCount *= 0.5; + } + return Math.max(baseCount, 1.0); + } + + @Override + public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) { + RelOptCost baseCost = super.computeSelfCost(planner, mq); + if (baseCost == null) { + return null; + } + double factor = 1.0; + if (pushDownContext.isTimeRangePushed()) { + factor *= 0.5; + } + if (pushDownContext.isLabelFilterPushed()) { + factor *= 0.7; + } + return baseCost.multiplyBy(factor); + } + + @Override + public RelWriter explainTerms(RelWriter pw) { + super.explainTerms(pw); + if (!pushDownContext.getLabelMatchers().isEmpty() || pushDownContext.isTimeRangePushed()) { + pw.item("PushDownContext", pushDownContext.toString()); + } + return pw; + } + + @Override + public Result implement(EnumerableRelImplementor implementor, Prefer pref) { + PhysType physType = + PhysTypeImpl.of(implementor.getTypeFactory(), getRowType(), pref.preferArray()); + + Expression scanOperator = implementor.stash(this, CalciteEnumerablePrometheusScan.class); + return implementor.result(physType, Blocks.toBlock(Expressions.call(scanOperator, "scan"))); + } + + @Override + public Enumerable<@Nullable Object> scan() { + return new AbstractEnumerable<>() { + @Override + public Enumerator<@Nullable Object> enumerator() { + return new PrometheusEnumerator(); + } + }; + } + + /** Enumerator that executes the PromQL query and iterates over results. */ + private class PrometheusEnumerator implements Enumerator<@Nullable Object> { + private Iterator responseIterator; + private List fieldOrder; + private Map fieldTypes; + private Object current; + + PrometheusEnumerator() { + try { + PrometheusClient client = prometheusTable.getPrometheusClient(); + String metricName = prometheusTable.getMetricName(); + + String promQL; + long startTime; + long endTime; + String step; + + if (prometheusTable.getPrometheusQueryRequest() != null) { + // Use pre-configured query request (from query_range table function) + var request = prometheusTable.getPrometheusQueryRequest(); + promQL = request.getPromQl(); + startTime = request.getStartTime(); + endTime = request.getEndTime(); + step = request.getStep(); + } else { + // Build PromQL from pushdown context + promQL = pushDownContext.buildPromQL(metricName); + startTime = pushDownContext.getEffectiveStartTime(); + endTime = pushDownContext.getEffectiveEndTime(); + step = pushDownContext.getEffectiveStep(); + } + + JSONObject responseObject = client.queryRange(promQL, startTime, endTime, step); + PrometheusResponseFieldNames fieldNames = new PrometheusResponseFieldNames(); + PrometheusResponse response = new PrometheusResponse(responseObject, fieldNames); + + this.fieldTypes = prometheusTable.getFieldTypes(); + this.fieldOrder = new ArrayList<>(fieldTypes.keySet()); + this.responseIterator = response.iterator(); + } catch (IOException e) { + throw new RuntimeException( + "Error fetching data from Prometheus server: " + e.getMessage(), e); + } + } + + @Override + public Object current() { + return current; + } + + @Override + public boolean moveNext() { + if (responseIterator.hasNext()) { + ExprValue exprValue = responseIterator.next(); + Map tupleValue = exprValue.tupleValue(); + + if (fieldOrder.size() == 1) { + // Single column — Calcite expects a scalar value + String fieldName = fieldOrder.get(0); + ExprValue value = tupleValue.get(fieldName); + current = convertExprValueToCalcite(value, fieldTypes.get(fieldName)); + } else { + // Multiple columns — Calcite expects Object[] + Object[] row = new Object[fieldOrder.size()]; + for (int i = 0; i < fieldOrder.size(); i++) { + String fieldName = fieldOrder.get(i); + ExprValue value = tupleValue.get(fieldName); + row[i] = convertExprValueToCalcite(value, fieldTypes.get(fieldName)); + } + current = row; + } + return true; + } + return false; + } + + @Override + public void reset() { + throw new UnsupportedOperationException("Reset not supported for Prometheus scan"); + } + + @Override + public void close() { + // No resources to close + } + + private Object convertExprValueToCalcite(ExprValue value, ExprType type) { + if (value == null) { + return null; + } + if (type == ExprCoreType.TIMESTAMP) { + return value.timestampValue().toEpochMilli(); + } else if (type == ExprCoreType.DOUBLE) { + return value.doubleValue(); + } else if (type == ExprCoreType.INTEGER) { + return value.integerValue(); + } else if (type == ExprCoreType.LONG) { + return value.longValue(); + } else if (type == ExprCoreType.STRING) { + return value.stringValue(); + } else { + return value.value().toString(); + } + } + } +} diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/scan/CalciteLogicalPrometheusScan.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/scan/CalciteLogicalPrometheusScan.java new file mode 100644 index 00000000000..36e75c21159 --- /dev/null +++ b/prometheus/src/main/java/org/opensearch/sql/prometheus/storage/scan/CalciteLogicalPrometheusScan.java @@ -0,0 +1,346 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.storage.scan; + +import com.google.common.collect.ImmutableList; +import java.util.List; +import lombok.Getter; +import org.apache.calcite.plan.Convention; +import org.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.plan.RelOptCost; +import org.apache.calcite.plan.RelOptPlanner; +import org.apache.calcite.plan.RelOptRule; +import org.apache.calcite.plan.RelOptTable; +import org.apache.calcite.plan.RelTraitSet; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.RelWriter; +import org.apache.calcite.rel.core.Filter; +import org.apache.calcite.rel.core.TableScan; +import org.apache.calcite.rel.metadata.RelMetadataQuery; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexInputRef; +import org.apache.calcite.rex.RexLiteral; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.type.SqlTypeName; +import org.opensearch.sql.prometheus.planner.logical.rules.PrometheusPushDownContext; +import org.opensearch.sql.prometheus.planner.logical.rules.PrometheusRules; +import org.opensearch.sql.prometheus.storage.PrometheusMetricTable; + +/** + * Logical scan node for Prometheus metrics in the Calcite planner. Supports pushdown of time range + * filters and label matchers via {@link PrometheusPushDownContext}. + */ +public class CalciteLogicalPrometheusScan extends TableScan { + + @Getter private final PrometheusMetricTable prometheusTable; + @Getter private final PrometheusPushDownContext pushDownContext; + @Getter private final RelDataType schema; + + public CalciteLogicalPrometheusScan( + RelOptCluster cluster, RelOptTable table, PrometheusMetricTable prometheusTable) { + this( + cluster, + cluster.traitSetOf(Convention.NONE), + table, + prometheusTable, + table.getRowType(), + new PrometheusPushDownContext()); + } + + public CalciteLogicalPrometheusScan( + RelOptCluster cluster, + RelTraitSet traitSet, + RelOptTable table, + PrometheusMetricTable prometheusTable, + RelDataType schema, + PrometheusPushDownContext pushDownContext) { + super(cluster, traitSet, ImmutableList.of(), table); + this.prometheusTable = prometheusTable; + this.schema = schema; + this.pushDownContext = pushDownContext; + } + + @Override + public RelDataType deriveRowType() { + return schema; + } + + @Override + public RelWriter explainTerms(RelWriter pw) { + super.explainTerms(pw); + if (!pushDownContext.getLabelMatchers().isEmpty() || pushDownContext.isTimeRangePushed()) { + pw.item("PushDownContext", pushDownContext.toString()); + } + return pw; + } + + @Override + public void register(RelOptPlanner planner) { + super.register(planner); + for (RelOptRule rule : PrometheusRules.PROMETHEUS_RULES) { + planner.addRule(rule); + } + } + + @Override + public double estimateRowCount(RelMetadataQuery mq) { + double baseCount = super.estimateRowCount(mq); + if (pushDownContext.isLabelFilterPushed()) { + baseCount *= 0.1; // Label filter significantly reduces number of time series + } + if (pushDownContext.isTimeRangePushed()) { + baseCount *= 0.5; + } + return Math.max(baseCount, 1.0); + } + + @Override + public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) { + RelOptCost baseCost = super.computeSelfCost(planner, mq); + if (baseCost == null) { + return null; + } + // Reduce cost when filters are pushed down to prefer pushed-down plans + double factor = 1.0; + if (pushDownContext.isTimeRangePushed()) { + factor *= 0.5; + } + if (pushDownContext.isLabelFilterPushed()) { + factor *= 0.7; + } + return baseCost.multiplyBy(factor); + } + + /** + * Attempts to push a filter condition into the Prometheus scan. Returns a new plan with the + * pushed-down condition removed, or null if nothing could be pushed. + * + *

Handles: - Time range comparisons on @timestamp (>, >=, <, <=) - Label equality conditions + * (label = 'value') + * + * @param filter the Filter RelNode above this scan + * @return a new plan (scan-only if fully pushed, or filter.copy with remaining condition), or + * null if nothing could be pushed + */ + public RelNode pushDownFilter(Filter filter) { + RexNode condition = filter.getCondition(); + List fieldNames = getRowType().getFieldNames(); + PrometheusPushDownContext newContext = pushDownContext.copy(); + RexNode remaining = pushDownCondition(condition, fieldNames, newContext); + + // If nothing was pushed, return null to indicate no transformation + if (!newContext.isTimeRangePushed() && !newContext.isLabelFilterPushed()) { + // Check if anything new was pushed compared to the original context + if (newContext.getLabelMatchers().equals(pushDownContext.getLabelMatchers()) + && java.util.Objects.equals(newContext.getStartTime(), pushDownContext.getStartTime()) + && java.util.Objects.equals(newContext.getEndTime(), pushDownContext.getEndTime())) { + return null; + } + } + + CalciteLogicalPrometheusScan newScan = + new CalciteLogicalPrometheusScan( + getCluster(), getTraitSet(), table, prometheusTable, schema, newContext); + + if (remaining != null) { + // Some conditions couldn't be pushed — use filter.copy() to preserve VolcanoPlanner + // equivalence semantics (same pattern as OpenSearch's FilterIndexScanRule) + return filter.copy(filter.getTraitSet(), newScan, remaining); + } + return newScan; + } + + /** + * Recursively analyzes a condition, pushing what can be pushed into the context and returning the + * remaining condition (or null if fully pushed). + */ + private RexNode pushDownCondition( + RexNode condition, List fieldNames, PrometheusPushDownContext context) { + if (!(condition instanceof RexCall)) { + return condition; // Can't push non-call expressions + } + + RexCall call = (RexCall) condition; + + if (call.getKind() == SqlKind.AND) { + // Process each AND operand independently + List remaining = new java.util.ArrayList<>(); + for (RexNode operand : call.getOperands()) { + RexNode r = pushDownCondition(operand, fieldNames, context); + if (r != null) { + remaining.add(r); + } + } + if (remaining.isEmpty()) { + return null; + } else if (remaining.size() == 1) { + return remaining.get(0); + } else { + return getCluster().getRexBuilder().makeCall(call.getOperator(), remaining); + } + } + + // Try to push this single condition + if (tryPushTimeRange(call, fieldNames, context)) { + return null; // Successfully pushed + } + if (tryPushLabelMatcher(call, fieldNames, context)) { + return null; // Successfully pushed + } + + return condition; // Can't push — return as remaining + } + + /** + * Tries to push a time range comparison (e.g., @timestamp >= '2024-01-01'). Returns true if + * pushed. + */ + private boolean tryPushTimeRange( + RexCall call, List fieldNames, PrometheusPushDownContext context) { + SqlKind kind = call.getKind(); + if (kind != SqlKind.GREATER_THAN + && kind != SqlKind.GREATER_THAN_OR_EQUAL + && kind != SqlKind.LESS_THAN + && kind != SqlKind.LESS_THAN_OR_EQUAL) { + return false; + } + + RexNode left = call.getOperands().get(0); + RexNode right = call.getOperands().get(1); + + // Determine which side is the field reference and which is the literal + String fieldName = null; + RexNode literal = null; + boolean fieldOnLeft = false; + + if (left instanceof RexInputRef && isTimestampLiteral(right)) { + fieldName = fieldNames.get(((RexInputRef) left).getIndex()); + literal = right; + fieldOnLeft = true; + } else if (right instanceof RexInputRef && isTimestampLiteral(left)) { + fieldName = fieldNames.get(((RexInputRef) right).getIndex()); + literal = left; + fieldOnLeft = false; + } + + if (fieldName == null || !fieldName.equals("@timestamp")) { + return false; + } + + long epochSeconds = extractEpochSeconds(literal); + // Determine effective comparison direction + SqlKind effectiveKind = fieldOnLeft ? kind : reverseComparison(kind); + + switch (effectiveKind) { + case GREATER_THAN: + case GREATER_THAN_OR_EQUAL: + context.pushStartTime(epochSeconds); + return true; + case LESS_THAN: + case LESS_THAN_OR_EQUAL: + context.pushEndTime(epochSeconds); + return true; + default: + return false; + } + } + + /** Tries to push a label equality condition (e.g., job = 'prometheus'). Returns true if pushed. */ + private boolean tryPushLabelMatcher( + RexCall call, List fieldNames, PrometheusPushDownContext context) { + if (call.getKind() != SqlKind.EQUALS) { + return false; + } + + RexNode left = call.getOperands().get(0); + RexNode right = call.getOperands().get(1); + + String fieldName = null; + String value = null; + + if (left instanceof RexInputRef && right instanceof RexLiteral) { + fieldName = fieldNames.get(((RexInputRef) left).getIndex()); + value = extractStringLiteral((RexLiteral) right); + } else if (right instanceof RexInputRef && left instanceof RexLiteral) { + fieldName = fieldNames.get(((RexInputRef) right).getIndex()); + value = extractStringLiteral((RexLiteral) left); + } + + if (fieldName == null || value == null) { + return false; + } + + // Don't push @timestamp or @value equality — those aren't label selectors + if (fieldName.equals("@timestamp") || fieldName.equals("@value")) { + return false; + } + + context.pushLabelMatcher(fieldName, value); + return true; + } + + private boolean isTimestampLiteral(RexNode node) { + if (node instanceof RexLiteral) { + return true; // Will try to extract epoch seconds + } + // Handle CAST expressions (e.g., CAST('2024-01-01' AS TIMESTAMP)) + if (node instanceof RexCall) { + RexCall castCall = (RexCall) node; + if (castCall.getKind() == SqlKind.CAST || castCall.getKind() == SqlKind.REINTERPRET) { + return castCall.getOperands().get(0) instanceof RexLiteral; + } + } + return false; + } + + private long extractEpochSeconds(RexNode node) { + if (node instanceof RexLiteral) { + RexLiteral lit = (RexLiteral) node; + SqlTypeName typeName = lit.getType().getSqlTypeName(); + if (typeName == SqlTypeName.TIMESTAMP || typeName == SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE) { + // Calcite stores timestamps as millis since epoch + Long millis = lit.getValueAs(Long.class); + return millis != null ? millis / 1000 : 0; + } else if (typeName == SqlTypeName.BIGINT || typeName == SqlTypeName.INTEGER) { + Long val = lit.getValueAs(Long.class); + return val != null ? val : 0; + } + // Try generic number extraction + Number num = lit.getValueAs(Number.class); + return num != null ? num.longValue() : 0; + } + if (node instanceof RexCall) { + RexCall castCall = (RexCall) node; + return extractEpochSeconds(castCall.getOperands().get(0)); + } + return 0; + } + + private String extractStringLiteral(RexLiteral literal) { + if (literal.getType().getSqlTypeName() == SqlTypeName.CHAR + || literal.getType().getSqlTypeName() == SqlTypeName.VARCHAR) { + return literal.getValueAs(String.class); + } + return null; + } + + private SqlKind reverseComparison(SqlKind kind) { + switch (kind) { + case GREATER_THAN: + return SqlKind.LESS_THAN; + case GREATER_THAN_OR_EQUAL: + return SqlKind.LESS_THAN_OR_EQUAL; + case LESS_THAN: + return SqlKind.GREATER_THAN; + case LESS_THAN_OR_EQUAL: + return SqlKind.GREATER_THAN_OR_EQUAL; + default: + return kind; + } + } +} diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/planner/logical/rules/PrometheusPushDownContextTest.java b/prometheus/src/test/java/org/opensearch/sql/prometheus/planner/logical/rules/PrometheusPushDownContextTest.java new file mode 100644 index 00000000000..1b669826b6d --- /dev/null +++ b/prometheus/src/test/java/org/opensearch/sql/prometheus/planner/logical/rules/PrometheusPushDownContextTest.java @@ -0,0 +1,268 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.prometheus.planner.logical.rules; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.junit.jupiter.api.Test; + +class PrometheusPushDownContextTest { + + @Test + void testDefaultState() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + + assertFalse(ctx.isTimeRangePushed()); + assertFalse(ctx.isLabelFilterPushed()); + assertTrue(ctx.getLabelMatchers().isEmpty()); + assertNull(ctx.getStartTime()); + assertNull(ctx.getEndTime()); + assertEquals("14", ctx.getEffectiveStep()); + } + + @Test + void testPushLabelMatcher() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + + ctx.pushLabelMatcher("service", "frontend"); + + assertTrue(ctx.isLabelFilterPushed()); + assertEquals(1, ctx.getLabelMatchers().size()); + assertEquals("frontend", ctx.getLabelMatchers().get("service")); + } + + @Test + void testPushMultipleLabels() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + + ctx.pushLabelMatcher("service", "frontend"); + ctx.pushLabelMatcher("job", "prometheus"); + + assertEquals(2, ctx.getLabelMatchers().size()); + assertEquals("frontend", ctx.getLabelMatchers().get("service")); + assertEquals("prometheus", ctx.getLabelMatchers().get("job")); + } + + @Test + void testPushStartTime() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + + ctx.pushStartTime(1700000000L); + + assertTrue(ctx.isTimeRangePushed()); + assertEquals(1700000000L, ctx.getStartTime()); + assertEquals(1700000000L, ctx.getEffectiveStartTime()); + } + + @Test + void testPushEndTime() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + + ctx.pushEndTime(1700003600L); + + assertTrue(ctx.isTimeRangePushed()); + assertEquals(1700003600L, ctx.getEndTime()); + assertEquals(1700003600L, ctx.getEffectiveEndTime()); + } + + @Test + void testEffectiveStartTimeDefault() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + + long effectiveStart = ctx.getEffectiveStartTime(); + long now = java.time.Instant.now().getEpochSecond(); + + // Default is now - 3600 seconds (1 hour), allow 5s tolerance + assertTrue(Math.abs(effectiveStart - (now - 3600)) < 5, + "Default start time should be approximately now - 1 hour"); + } + + @Test + void testEffectiveEndTimeDefault() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + + long effectiveEnd = ctx.getEffectiveEndTime(); + long now = java.time.Instant.now().getEpochSecond(); + + // Default is now, allow 5s tolerance + assertTrue(Math.abs(effectiveEnd - now) < 5, + "Default end time should be approximately now"); + } + + @Test + void testSetStep() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + + ctx.setStep("30"); + + assertEquals("30", ctx.getEffectiveStep()); + } + + @Test + void testGetEffectiveStepWhenNull() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + + ctx.setStep(null); + + assertEquals("14", ctx.getEffectiveStep()); + } + + @Test + void testBuildPromQLNoLabels() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + + assertEquals("up", ctx.buildPromQL("up")); + assertEquals("http_requests_total", ctx.buildPromQL("http_requests_total")); + } + + @Test + void testBuildPromQLSingleLabel() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + ctx.pushLabelMatcher("job", "prometheus"); + + assertEquals("up{job=\"prometheus\"}", ctx.buildPromQL("up")); + } + + @Test + void testBuildPromQLMultipleLabels() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + ctx.pushLabelMatcher("job", "prometheus"); + ctx.pushLabelMatcher("instance", "localhost:9090"); + + assertEquals( + "up{job=\"prometheus\",instance=\"localhost:9090\"}", + ctx.buildPromQL("up")); + } + + @Test + void testBuildPromQLEscapesDoubleQuotes() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + ctx.pushLabelMatcher("path", "/api/v1/\"test\""); + + assertEquals( + "metric{path=\"/api/v1/\\\"test\\\"\"}", + ctx.buildPromQL("metric")); + } + + @Test + void testBuildPromQLEscapesBackslashes() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + ctx.pushLabelMatcher("path", "C:\\Users\\admin"); + + assertEquals( + "metric{path=\"C:\\\\Users\\\\admin\"}", + ctx.buildPromQL("metric")); + } + + @Test + void testBuildPromQLEscapesNewlines() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + ctx.pushLabelMatcher("msg", "line1\nline2"); + + assertEquals( + "metric{msg=\"line1\\nline2\"}", + ctx.buildPromQL("metric")); + } + + @Test + void testBuildPromQLEscapesCombined() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + // Value with all special chars: backslash, quote, newline + ctx.pushLabelMatcher("val", "a\\b\"c\nd"); + + assertEquals( + "metric{val=\"a\\\\b\\\"c\\nd\"}", + ctx.buildPromQL("metric")); + } + + @Test + void testBuildPromQLInjectionAttempt() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + // Attempt to break out of label value and inject another matcher + ctx.pushLabelMatcher("service", "frontend\"} OR {__name__=\"secret"); + + String result = ctx.buildPromQL("metric"); + // The injection should be safely escaped inside the quotes + assertEquals( + "metric{service=\"frontend\\\"} OR {__name__=\\\"secret\"}", + result); + } + + @Test + void testEscapePromQLLabelValue() { + assertEquals("simple", PrometheusPushDownContext.escapePromQLLabelValue("simple")); + assertEquals("has\\\\backslash", PrometheusPushDownContext.escapePromQLLabelValue("has\\backslash")); + assertEquals("has\\\"quote", PrometheusPushDownContext.escapePromQLLabelValue("has\"quote")); + assertEquals("has\\nnewline", PrometheusPushDownContext.escapePromQLLabelValue("has\nnewline")); + assertEquals("", PrometheusPushDownContext.escapePromQLLabelValue("")); + } + + @Test + void testCopyIsIndependent() { + PrometheusPushDownContext original = new PrometheusPushDownContext(); + original.pushLabelMatcher("service", "frontend"); + original.pushStartTime(1700000000L); + original.setStep("30"); + + PrometheusPushDownContext copy = original.copy(); + + // Copy has same state + assertEquals(original.getLabelMatchers(), copy.getLabelMatchers()); + assertEquals(original.getStartTime(), copy.getStartTime()); + assertEquals(original.getEffectiveStep(), copy.getEffectiveStep()); + assertTrue(copy.isLabelFilterPushed()); + assertTrue(copy.isTimeRangePushed()); + + // Mutating copy doesn't affect original + copy.pushLabelMatcher("job", "node"); + copy.pushEndTime(1700003600L); + copy.setStep("60"); + + assertEquals(1, original.getLabelMatchers().size()); + assertNull(original.getEndTime()); + assertEquals("30", original.getEffectiveStep()); + assertNotSame(original.getLabelMatchers(), copy.getLabelMatchers()); + } + + @Test + void testToStringEmpty() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + + assertEquals("[]", ctx.toString()); + } + + @Test + void testToStringWithLabels() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + ctx.pushLabelMatcher("service", "frontend"); + + assertEquals("[LABELS->{service=frontend}]", ctx.toString()); + } + + @Test + void testToStringWithTimeRange() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + ctx.pushStartTime(1700000000L); + ctx.pushEndTime(1700003600L); + + assertEquals("[TIME_RANGE->[1700000000,1700003600]]", ctx.toString()); + } + + @Test + void testToStringWithLabelsAndTimeRange() { + PrometheusPushDownContext ctx = new PrometheusPushDownContext(); + ctx.pushLabelMatcher("service", "frontend"); + ctx.pushStartTime(1700000000L); + ctx.pushEndTime(1700003600L); + + assertEquals("[LABELS->{service=frontend}, TIME_RANGE->[1700000000,1700003600]]", ctx.toString()); + } +} diff --git a/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusMetricTableTest.java b/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusMetricTableTest.java index c6b9b63ec5e..beeedb749b9 100644 --- a/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusMetricTableTest.java +++ b/prometheus/src/test/java/org/opensearch/sql/prometheus/storage/PrometheusMetricTableTest.java @@ -10,6 +10,10 @@ import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoMoreInteractions; @@ -31,6 +35,7 @@ import static org.opensearch.sql.prometheus.utils.LogicalPlanUtils.testLogicalPlanNode; import com.google.common.collect.ImmutableList; +import java.io.IOException; import java.text.DateFormat; import java.text.SimpleDateFormat; import java.util.ArrayList; @@ -41,11 +46,17 @@ import java.util.Map; import java.util.stream.Collectors; import lombok.SneakyThrows; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeField; +import org.apache.calcite.schema.TranslatableTable; +import org.apache.calcite.sql.type.SqlTypeName; +import org.json.JSONObject; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory; import org.opensearch.sql.data.type.ExprCoreType; import org.opensearch.sql.data.type.ExprType; import org.opensearch.sql.expression.DSL; @@ -1053,4 +1064,125 @@ void testCreateScanBuilderWithPPLQuery() { TableScanBuilder tableScanBuilder = prometheusMetricTable.createScanBuilder(); Assertions.assertNull(tableScanBuilder); } + + // ---- Calcite ScannableTable tests ---- + + @Test + void testImplementsTranslatableTable() { + PrometheusMetricTable prometheusMetricTable = + new PrometheusMetricTable(client, TestConstants.METRIC_NAME); + assertTrue(prometheusMetricTable instanceof TranslatableTable); + assertTrue(prometheusMetricTable instanceof org.apache.calcite.schema.Table); + } + + @Test + @SneakyThrows + void testGetRowTypeFromMetric() { + when(client.getLabels(TestConstants.METRIC_NAME)).thenReturn(List.of("job", "instance")); + PrometheusMetricTable prometheusMetricTable = + new PrometheusMetricTable(client, TestConstants.METRIC_NAME); + + RelDataType rowType = + prometheusMetricTable.getRowType(OpenSearchTypeFactory.TYPE_FACTORY); + + assertNotNull(rowType); + List fields = rowType.getFieldList(); + assertTrue(fields.size() >= 2, "Should have at least @timestamp and @value fields"); + + // Verify @timestamp and @value fields exist + boolean hasTimestamp = fields.stream().anyMatch(f -> f.getName().equals("@timestamp")); + boolean hasValue = fields.stream().anyMatch(f -> f.getName().equals("@value")); + assertTrue(hasTimestamp, "Should have @timestamp field"); + assertTrue(hasValue, "Should have @value field"); + } + + @Test + @SneakyThrows + void testGetRowTypeFromQueryRequest() { + PrometheusMetricTable prometheusMetricTable = + new PrometheusMetricTable(client, new PrometheusQueryRequest()); + + RelDataType rowType = + prometheusMetricTable.getRowType(OpenSearchTypeFactory.TYPE_FACTORY); + + assertNotNull(rowType); + List fields = rowType.getFieldList(); + // query_range returns @timestamp, @value, and @labels + assertTrue(fields.size() >= 3, "Should have @timestamp, @value, and @labels fields"); + } + + @Test + @SneakyThrows + void testMetricNameAccessible() { + PrometheusMetricTable prometheusMetricTable = + new PrometheusMetricTable(client, "test_metric"); + + assertEquals("test_metric", prometheusMetricTable.getMetricName()); + assertEquals(client, prometheusMetricTable.getPrometheusClient()); + assertNull(prometheusMetricTable.getPrometheusQueryRequest()); + } + + @Test + @SneakyThrows + void testQueryRequestAccessible() { + PrometheusQueryRequest request = new PrometheusQueryRequest(); + request.setPromQl("up"); + request.setStartTime(1435781400L); + request.setEndTime(1435785000L); + request.setStep("14"); + + PrometheusMetricTable prometheusMetricTable = new PrometheusMetricTable(client, request); + + assertEquals(request, prometheusMetricTable.getPrometheusQueryRequest()); + assertEquals(client, prometheusMetricTable.getPrometheusClient()); + assertNull(prometheusMetricTable.getMetricName()); + } + + @Test + @SneakyThrows + void testGetRowTypeHasExpectedFields() { + when(client.getLabels("empty_metric")).thenReturn(List.of("job")); + + PrometheusMetricTable prometheusMetricTable = + new PrometheusMetricTable(client, "empty_metric"); + RelDataType rowType = + prometheusMetricTable.getRowType(OpenSearchTypeFactory.TYPE_FACTORY); + + assertNotNull(rowType); + List fieldNames = rowType.getFieldNames(); + assertTrue(fieldNames.contains("@timestamp"), "Should have @timestamp field"); + assertTrue(fieldNames.contains("@value"), "Should have @value field"); + assertTrue(fieldNames.contains("job"), "Should have job label field"); + } + + @Test + @SneakyThrows + void testToRelReturnsLogicalPrometheusScan() { + PrometheusMetricTable prometheusMetricTable = + new PrometheusMetricTable(client, "test_metric"); + + // Verify the table is a TranslatableTable + assertTrue(prometheusMetricTable instanceof TranslatableTable); + + // Create a minimal Calcite context to invoke toRel() + org.apache.calcite.plan.RelOptCluster cluster = + org.apache.calcite.plan.RelOptCluster.create( + new org.apache.calcite.plan.volcano.VolcanoPlanner(), + new org.apache.calcite.rex.RexBuilder( + org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.TYPE_FACTORY)); + + org.apache.calcite.plan.RelOptTable relOptTable = + org.mockito.Mockito.mock(org.apache.calcite.plan.RelOptTable.class); + org.apache.calcite.plan.RelOptTable.ToRelContext toRelContext = + org.mockito.Mockito.mock(org.apache.calcite.plan.RelOptTable.ToRelContext.class); + when(toRelContext.getCluster()).thenReturn(cluster); + + org.apache.calcite.rel.RelNode result = prometheusMetricTable.toRel(toRelContext, relOptTable); + + assertNotNull(result); + assertTrue( + result + instanceof + org.opensearch.sql.prometheus.storage.scan.CalciteLogicalPrometheusScan); + } }