code: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   *
   */
  
  package org.apache.skywalking.oap.query.graphql.mqe.rt;
  
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.List;
  import java.util.Optional;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.skywalking.oap.query.graphql.mqe.rt.exception.IllegalExpressionException;
  import org.apache.skywalking.oap.query.graphql.mqe.rt.operation.AggregateLabelsOp;
  import org.apache.skywalking.oap.query.graphql.mqe.rt.operation.AggregationOp;
  import org.apache.skywalking.oap.query.graphql.resolver.MetricsQuery;
  import org.apache.skywalking.oap.query.graphql.resolver.RecordsQuery;
  import org.apache.skywalking.oap.query.graphql.type.mql.ExpressionResult;
  import org.apache.skywalking.oap.query.graphql.type.mql.ExpressionResultType;
  import org.apache.skywalking.oap.query.graphql.type.mql.MQEValue;
  import org.apache.skywalking.oap.query.graphql.type.mql.MQEValues;
  import org.apache.skywalking.oap.query.graphql.type.mql.Metadata;
  import org.apache.skywalking.oap.server.core.Const;
  import org.apache.skywalking.oap.server.core.query.DurationUtils;
  import org.apache.skywalking.oap.server.core.query.PointOfTime;
  import org.apache.skywalking.oap.server.core.query.enumeration.Order;
  import org.apache.skywalking.oap.server.core.query.input.Duration;
  import org.apache.skywalking.oap.server.core.query.input.Entity;
  import org.apache.skywalking.oap.server.core.query.input.MetricsCondition;
  import org.apache.skywalking.oap.server.core.query.input.RecordCondition;
  import org.apache.skywalking.oap.server.core.query.input.TopNCondition;
  import org.apache.skywalking.oap.server.core.query.type.KVInt;
  import org.apache.skywalking.oap.server.core.query.type.KeyValue;
  import org.apache.skywalking.oap.server.core.query.type.MetricsValues;
  import org.apache.skywalking.oap.server.core.query.type.Record;
  import org.apache.skywalking.oap.server.core.query.type.SelectedRecord;
  import org.apache.skywalking.oap.server.core.storage.annotation.Column;
  import org.apache.skywalking.oap.server.core.storage.annotation.ValueColumnMetadata;
  import org.apache.skywalking.oap.server.library.util.StringUtil;
  import org.apache.skywalking.mqe.rt.grammar.MQEParser;
  import org.apache.skywalking.mqe.rt.grammar.MQEParserBaseVisitor;
  
  import static org.apache.skywalking.oap.query.graphql.mqe.rt.operation.BinaryOp.doBinaryOp;
  import static org.apache.skywalking.oap.query.graphql.mqe.rt.operation.FunctionOp.doFunction0Op;
  import static org.apache.skywalking.oap.query.graphql.mqe.rt.operation.FunctionOp.doFunction1Op;
  
  @Slf4j
  public class MQEVisitor extends MQEParserBaseVisitor<ExpressionResult> {
      private final MetricsQuery metricsQuery;
      private final RecordsQuery recordsQuery;
      private final Entity entity;
      private final Duration duration;
  
      private final static String GENERAL_LABEL_NAME = "_";
  
      public MQEVisitor(final MetricsQuery metricsQuery,
                        final RecordsQuery recordsQuery,
                        final Entity entity,
                        final Duration duration) {
          this.metricsQuery = metricsQuery;
          this.recordsQuery = recordsQuery;
          this.entity = entity;
          this.duration = duration;
      }
  
      @Override
      public ExpressionResult visitAddSubOp(MQEParser.AddSubOpContext ctx) {
          ExpressionResult left = visit(ctx.expression(0));
          if (StringUtil.isNotBlank(left.getError())) {
              return left;
          }
          ExpressionResult right = visit(ctx.expression(1));
          if (StringUtil.isNotBlank(right.getError())) {
              return right;
          }
          int opType = ctx.addSub().getStart().getType();
          try {
              return doBinaryOp(left, right, opType);
          } catch (IllegalExpressionException e) {
              ExpressionResult result = new ExpressionResult();
              result.setType(ExpressionResultType.UNKNOWN);
              result.setError(e.getMessage());
              return result;
          }
      }
  
      @Override
      public ExpressionResult visitMulDivModOp(MQEParser.MulDivModOpContext ctx) {
          ExpressionResult left = visit(ctx.expression(0));
          if (StringUtil.isNotBlank(left.getError())) {
              return left;
          }
          ExpressionResult right = visit(ctx.expression(1));
          if (StringUtil.isNotBlank(right.getError())) {
              return right;
          }
          int opType = ctx.mulDivMod().getStart().getType();
          try {
              return doBinaryOp(left, right, opType);
          } catch (IllegalExpressionException e) {
              ExpressionResult result = new ExpressionResult();
              result.setType(ExpressionResultType.UNKNOWN);
              result.setError(e.getMessage());
              return result;
          }
      }
  
      @Override
      public ExpressionResult visitScalar(MQEParser.ScalarContext ctx) {
          ExpressionResult result = new ExpressionResult();
          double value = Double.parseDouble(ctx.getText());
          MQEValue mqeValue = new MQEValue();
          mqeValue.setDoubleValue(value);
          mqeValue.setEmptyValue(false);
          MQEValues mqeValues = new MQEValues();
          mqeValues.getValues().add(mqeValue);
          result.getResults().add(mqeValues);
          result.setType(ExpressionResultType.SINGLE_VALUE);
          return result;
      }
  
      @Override
      public ExpressionResult visitAggregationOp(MQEParser.AggregationOpContext ctx) {
          int opType = ctx.aggregation().getStart().getType();
          ExpressionResult expResult = visit(ctx.expression());
          if (StringUtil.isNotEmpty(expResult.getError())) {
              return expResult;
          }
          try {
              return AggregationOp.doAggregationOp(expResult, opType);
          } catch (IllegalExpressionException e) {
              ExpressionResult result = new ExpressionResult();
              result.setType(ExpressionResultType.UNKNOWN);
              result.setError(e.getMessage());
              return result;
          }
      }
  
      @Override
      public ExpressionResult visitAggregateLabelsOp(final MQEParser.AggregateLabelsOpContext ctx) {
          int funcType = ctx.aggregateLabelsFunc().getStart().getType();
          ExpressionResult expResult = visit(ctx.expression());
          if (StringUtil.isNotEmpty(expResult.getError())) {
              return expResult;
          }
  
          if (!expResult.isLabeledResult()) {
              expResult.setError("The result of expression [" + ctx.expression().getText() + "] is not a labeled result.");
              return expResult;
          }
  
          try {
              return AggregateLabelsOp.doAggregateLabelsOp(expResult, funcType);
          } catch (IllegalExpressionException e) {
              ExpressionResult result = new ExpressionResult();
              result.setType(ExpressionResultType.UNKNOWN);
              result.setError(e.getMessage());
              return result;
          }
      }
  
      @Override
      public ExpressionResult visitFunction0OP(MQEParser.Function0OPContext ctx) {
          int opType = ctx.function0().getStart().getType();
          ExpressionResult expResult = visit(ctx.expression());
          if (StringUtil.isNotEmpty(expResult.getError())) {
              return expResult;
          }
          try {
              return doFunction0Op(expResult, opType);
          } catch (IllegalExpressionException e) {
              ExpressionResult result = new ExpressionResult();
              result.setType(ExpressionResultType.UNKNOWN);
              result.setError(e.getMessage());
              return result;
          }
      }
  
      @Override
      public ExpressionResult visitFunction1OP(MQEParser.Function1OPContext ctx) {
          int opType = ctx.function1().getStart().getType();
          ExpressionResult expResult = visit(ctx.expression());
          if (StringUtil.isNotEmpty(expResult.getError())) {
              return expResult;
          }
          try {
              return doFunction1Op(expResult, opType, Integer.parseInt(ctx.parameter().INTEGER().getText()));
          } catch (IllegalExpressionException e) {
              ExpressionResult result = new ExpressionResult();
              result.setType(ExpressionResultType.UNKNOWN);
              result.setError(e.getMessage());
              return result;
          }
      }
  
      @Override
      public ExpressionResult visitTopNOP(MQEParser.TopNOPContext ctx) {
          return visit(ctx.metric());
      }
  
      @Override
      public ExpressionResult visitRelablesOP(MQEParser.RelablesOPContext ctx) {
          ExpressionResult result = visit(ctx.expression());
          if (!result.isLabeledResult()) {
              // Reserve the original result type
              result.setError("The result of expression [" + ctx.expression().getText() + "] is not a labeled result.");
              return result;
          }
  
          List<String> relabelList = Collections.emptyList();
          String newLabelValue = ctx.label().labelValue().getText();
          String labelValueTrim = newLabelValue.substring(1, newLabelValue.length() - 1);
          if (StringUtil.isNotBlank(labelValueTrim)) {
              relabelList = Arrays.asList(labelValueTrim.split(Const.COMMA));
          }
          List<MQEValues> mqeValuesList = result.getResults();
  
          if (mqeValuesList.size() != relabelList.size()) {
              // Resever the original result type
              result.setError("The number of relabels is not equal to the number of labels.");
              return result;
          }
          // For now, we only have a single label named `label`
          for (int i = 0; i < mqeValuesList.size(); i++) {
              mqeValuesList.get(i).getMetric().getLabels().get(0).setValue(relabelList.get(i));
          }
  
          return result;
      }
  
      @Override
      public ExpressionResult visitMetric(MQEParser.MetricContext ctx) {
          ExpressionResult result = new ExpressionResult();
          String metricName = ctx.metricName().getText();
          Optional<ValueColumnMetadata.ValueColumn> valueColumn = ValueColumnMetadata.INSTANCE.readValueColumnDefinition(
              metricName);
          if (valueColumn.isEmpty()) {
              result.setType(ExpressionResultType.UNKNOWN);
              result.setError("Metric: [" + metricName + "] dose not exist.");
              return result;
          }
  
          Column.ValueDataType dataType = valueColumn.get().getDataType();
          try {
              if (Column.ValueDataType.COMMON_VALUE == dataType) {
                  if (ctx.parent instanceof MQEParser.TopNOPContext) {
                      MQEParser.TopNOPContext parent = (MQEParser.TopNOPContext) ctx.parent;
                      querySortMetrics(metricName, Integer.parseInt(parent.parameter().getText()),
                                       Order.valueOf(parent.order().getText().toUpperCase()), result
                      );
                  } else {
                      queryMetrics(metricName, result);
                  }
              } else if (Column.ValueDataType.LABELED_VALUE == dataType) {
                  List<String> queryLabelList = Collections.emptyList();
                  if (ctx.label() != null) {
                      String labelValue = ctx.label().labelValue().getText();
                      String labelValueTrim = labelValue.substring(1, labelValue.length() - 1);
                      if (StringUtil.isNotBlank(labelValueTrim)) {
                          queryLabelList = Arrays.asList(labelValueTrim.split(Const.COMMA));
                      }
                  }
                  queryLabeledMetrics(metricName, queryLabelList, result);
              } else if (Column.ValueDataType.SAMPLED_RECORD == dataType) {
                  if (ctx.parent instanceof MQEParser.TopNOPContext) {
                      MQEParser.TopNOPContext parent = (MQEParser.TopNOPContext) ctx.parent;
                      queryRecords(metricName, Integer.parseInt(parent.parameter().getText()),
                                   Order.valueOf(parent.order().getText().toUpperCase()), result
                      );
                  } else {
                      throw new IllegalExpressionException(
                          "Metric: [" + metricName + "] is topN record, need top_n function for query.");
                  }
              }
          } catch (IllegalExpressionException e) {
              ExpressionResult errorResult = new ExpressionResult();
              errorResult.setType(ExpressionResultType.UNKNOWN);
              errorResult.setError(e.getMessage());
              return errorResult;
          } catch (IOException e) {
              ExpressionResult errorResult = new ExpressionResult();
              errorResult.setType(ExpressionResultType.UNKNOWN);
              errorResult.setError("Internal IO exception, query metrics error.");
              log.error("Query metrics from backend error.", e);
              return errorResult;
          }
          return result;
      }
  
      private void querySortMetrics(String metricName,
                                    int topN,
                                    Order order,
                                    ExpressionResult result) throws IOException {
          TopNCondition topNCondition = new TopNCondition();
          topNCondition.setName(metricName);
          topNCondition.setTopN(topN);
          topNCondition.setParentService(entity.getServiceName());
          topNCondition.setOrder(order);
          topNCondition.setNormal(entity.getNormal());
          List<SelectedRecord> selectedRecords = metricsQuery.sortMetrics(topNCondition, duration);
  
          List<MQEValue> mqeValueList = new ArrayList<>(selectedRecords.size());
          selectedRecords.forEach(selectedRecord -> {
              MQEValue mqeValue = new MQEValue();
              mqeValue.setId(selectedRecord.getName());
              mqeValue.setEmptyValue(false);
              mqeValue.setDoubleValue(Double.parseDouble(selectedRecord.getValue()));
              mqeValueList.add(mqeValue);
          });
          Metadata metadata = new Metadata();
          MQEValues mqeValues = new MQEValues();
          mqeValues.setValues(mqeValueList);
          mqeValues.setMetric(metadata);
          result.getResults().add(mqeValues);
          result.setType(ExpressionResultType.SORTED_LIST);
      }
  
      private void queryRecords(String metricName, int topN, Order order, ExpressionResult result) throws IOException {
          RecordCondition recordCondition = new RecordCondition();
          recordCondition.setName(metricName);
          recordCondition.setTopN(topN);
          recordCondition.setParentEntity(entity);
          recordCondition.setOrder(order);
          List<Record> records = recordsQuery.readRecords(recordCondition, duration);
  
          List<MQEValue> mqeValueList = new ArrayList<>(records.size());
          records.forEach(record -> {
              MQEValue mqeValue = new MQEValue();
              mqeValue.setId(record.getName());
              mqeValue.setEmptyValue(false);
              mqeValue.setDoubleValue(Double.parseDouble(record.getValue()));
              mqeValue.setTraceID(record.getRefId());
              mqeValueList.add(mqeValue);
          });
          Metadata metadata = new Metadata();
          MQEValues mqeValues = new MQEValues();
          mqeValues.setValues(mqeValueList);
          mqeValues.setMetric(metadata);
          result.getResults().add(mqeValues);
          result.setType(ExpressionResultType.RECORD_LIST);
      }
  
      private void queryMetrics(String metricName, ExpressionResult result) throws IOException {
          MetricsCondition metricsCondition = new MetricsCondition();
          metricsCondition.setName(metricName);
          metricsCondition.setEntity(entity);
          MetricsValues metricsValues = metricsQuery.readMetricsValues(metricsCondition, duration);
          List<PointOfTime> times = duration.assembleDurationPoints();
          List<MQEValue> mqeValueList = new ArrayList<>(times.size());
          for (int i = 0; i < times.size(); i++) {
              long retTimestamp = DurationUtils.INSTANCE.parseToDateTime(duration.getStep(), times.get(i).getPoint())
                                                        .getMillis();
              KVInt kvInt = metricsValues.getValues().getValues().get(i);
              MQEValue mqeValue = new MQEValue();
              mqeValue.setId(Long.toString(retTimestamp));
              mqeValue.setEmptyValue(kvInt.isEmptyValue());
              mqeValue.setDoubleValue(kvInt.getValue());
              mqeValueList.add(mqeValue);
          }
          Metadata metadata = new Metadata();
          MQEValues mqeValues = new MQEValues();
          mqeValues.setValues(mqeValueList);
          mqeValues.setMetric(metadata);
          result.getResults().add(mqeValues);
          result.setType(ExpressionResultType.TIME_SERIES_VALUES);
      }
  
      private void queryLabeledMetrics(String metricName,
                                       List<String> queryLabelList,
                                       ExpressionResult result) throws IOException {
          MetricsCondition metricsCondition = new MetricsCondition();
          metricsCondition.setName(metricName);
          metricsCondition.setEntity(entity);
          List<MetricsValues> metricsValuesList = metricsQuery.readLabeledMetricsValues(
              metricsCondition, queryLabelList, duration);
          List<PointOfTime> times = duration.assembleDurationPoints();
          metricsValuesList.forEach(metricsValues -> {
              List<MQEValue> mqeValueList = new ArrayList<>(times.size());
              for (int i = 0; i < times.size(); i++) {
                  long retTimestamp = DurationUtils.INSTANCE.parseToDateTime(duration.getStep(), times.get(i).getPoint())
                                                            .getMillis();
                  KVInt kvInt = metricsValues.getValues().getValues().get(i);
                  MQEValue mqeValue = new MQEValue();
                  mqeValue.setEmptyValue(kvInt.isEmptyValue());
                  mqeValue.setId(Long.toString(retTimestamp));
                  mqeValueList.add(mqeValue);
                  if (!kvInt.isEmptyValue()) {
                      mqeValue.setDoubleValue(kvInt.getValue());
                  }
              }
  
              Metadata metadata = new Metadata();
              KeyValue labelValue = new KeyValue(GENERAL_LABEL_NAME, metricsValues.getLabel());
              metadata.getLabels().add(labelValue);
              MQEValues mqeValues = new MQEValues();
              mqeValues.setValues(mqeValueList);
              mqeValues.setMetric(metadata);
              result.getResults().add(mqeValues);
          });
          result.setType(ExpressionResultType.TIME_SERIES_VALUES);
          result.setLabeledResult(true);
      }
  }
  