Span Fix for month and quarter which prometheus doesn't support

Signed-off-by: vamsi-amazon <[email protected]>
opensearch-project · Nov 4, 2022 · 0f8ed02 · 0f8ed02
1 parent 3e30379
commit 0f8ed02
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 11 deletions.
diff --git a/prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java b/prometheus/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java
@@ -106,7 +106,8 @@ private JSONObject readResponse(Response response) throws IOException {
       }
     } else {
       throw new RuntimeException(
-          String.format("Request to Prometheus is Unsuccessful with : %s", response.message()));
+          String.format("Request to Prometheus is Unsuccessful with : %s", Objects.requireNonNull(
+              response.body(), "Response body can't be null").string()));
     }
   }
 

diff --git a/...c/main/java/org/opensearch/sql/prometheus/storage/querybuilder/StepParameterResolver.java b/...c/main/java/org/opensearch/sql/prometheus/storage/querybuilder/StepParameterResolver.java
@@ -12,6 +12,7 @@
 import lombok.NoArgsConstructor;
 import lombok.NonNull;
 import org.apache.commons.lang3.StringUtils;
+import org.opensearch.sql.ast.expression.SpanUnit;
 import org.opensearch.sql.expression.NamedExpression;
 import org.opensearch.sql.expression.span.SpanExpression;
 
@@ -26,29 +27,27 @@ public class StepParameterResolver {
    * Extract step from groupByList or apply heuristic arithmetic
    * on endTime and startTime.
    *
-   *
-   * @param startTime startTime.
-   * @param endTime endTime.
+   * @param startTime   startTime.
+   * @param endTime     endTime.
    * @param groupByList groupByList.
    * @return Step String.
    */
   public static String resolve(@NonNull Long startTime, @NonNull Long endTime,
-                        List<NamedExpression> groupByList) {
+                               List<NamedExpression> groupByList) {
     Optional<SpanExpression> spanExpression = getSpanExpression(groupByList);
     if (spanExpression.isPresent()) {
       if (StringUtils.isEmpty(spanExpression.get().getUnit().getName())) {
         throw new RuntimeException("Missing TimeUnit in the span expression");
       } else {
-        return spanExpression.get().getValue().toString()
-            + spanExpression.get().getUnit().getName();
+        return getTimespan(spanExpression.get());
       }
     } else {
       return Math.max((endTime - startTime) / 250, 1) + "s";
     }
   }
 
   private static Optional<SpanExpression> getSpanExpression(
-        List<NamedExpression> namedExpressionList) {
+      List<NamedExpression> namedExpressionList) {
     if (namedExpressionList == null) {
       return Optional.empty();
     }
@@ -58,6 +57,15 @@ private static Optional<SpanExpression> getSpanExpression(
         .findFirst();
   }
 
+  private static String getTimespan(SpanExpression spanExpression) {
+    if (SpanUnit.M.equals(spanExpression.getUnit())) {
+      return spanExpression.getValue().valueOf(null).integerValue() * 30 + "d";
+    } else if (SpanUnit.Q.equals(spanExpression.getUnit())) {
+      return spanExpression.getValue().valueOf(null).integerValue() * 90 + "d";
+    } else {
+      return spanExpression.getValue().toString() + spanExpression.getUnit().getName();
+    }
+  }
 
 
 }
diff --git a/...t/java/org/opensearch/sql/prometheus/storage/querybuilders/StepParameterResolverTest.java b/...t/java/org/opensearch/sql/prometheus/storage/querybuilders/StepParameterResolverTest.java
@@ -7,20 +7,33 @@
 
 package org.opensearch.sql.prometheus.storage.querybuilders;
 
+import static org.opensearch.sql.expression.DSL.named;
+
 import java.util.Collections;
 import java.util.Date;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
+import org.opensearch.sql.data.type.ExprCoreType;
+import org.opensearch.sql.expression.DSL;
 import org.opensearch.sql.prometheus.storage.querybuilder.StepParameterResolver;
 
 public class StepParameterResolverTest {
 
   @Test
   void testNullChecks() {
-    StepParameterResolver stepParameterResolver = new StepParameterResolver();
     Assertions.assertThrows(NullPointerException.class,
-        () -> stepParameterResolver.resolve(null, new Date().getTime(), Collections.emptyList()));
+        () -> StepParameterResolver.resolve(null, new Date().getTime(), Collections.emptyList()));
     Assertions.assertThrows(NullPointerException.class,
-        () -> stepParameterResolver.resolve(new Date().getTime(), null, Collections.emptyList()));
+        () -> StepParameterResolver.resolve(new Date().getTime(), null, Collections.emptyList()));
+    String step = StepParameterResolver.resolve(new Date().getTime(), new Date().getTime() - 3600,
+        Collections.singletonList(named("span",
+            DSL.span(DSL.ref("@timestamp", ExprCoreType.TIMESTAMP),
+            DSL.literal(2), "M"))));
+    Assertions.assertEquals("60d", step);
+    step = StepParameterResolver.resolve(new Date().getTime(), new Date().getTime() - 3600,
+        Collections.singletonList(named("span",
+            DSL.span(DSL.ref("@timestamp", ExprCoreType.TIMESTAMP),
+                DSL.literal(4), "Q"))));
+    Assertions.assertEquals("360d", step);
   }
 }