webapi: improve normalize behavior for sample-count (#1671)

Before it was getting expanded exposing the underlying details
and showing the user a more complex query that is hard to for
them to understand. Use named rewrite to preserve the operator
when doing the normalization.

atlas-core/src/main/scala/com/netflix/atlas/core/model/CustomVocabulary.scala

@@ -157,7 +157,7 @@ object CustomVocabulary {
         val avg = MathExpr.Divide(numerator, denominator)
         val ctxt = Context(context.interpreter, Nil, Map.empty)
         val rewrite = Some(this)
-        MathExpr.NamedRewrite(name, q, avg, ctxt, rewrite) :: s
+        MathExpr.NamedRewrite(name, q, Nil, avg, ctxt, rewrite) :: s
     }
 
     /**

atlas-core/src/main/scala/com/netflix/atlas/core/model/MathExpr.scala

@@ -990,6 +990,7 @@ object MathExpr {
   case class NamedRewrite(
     name: String,
     displayExpr: Expr,
+    displayParams: List[Any],
     evalExpr: TimeSeriesExpr,
     context: Context,
     groupByRewrite: Option[(Expr, List[String]) => Expr] = None
@@ -1000,14 +1001,19 @@ object MathExpr {
     override def toString: String = toString(displayExpr)
 
     private def toString(expr: Expr): String = {
+      val op =
+        if (displayParams.isEmpty)
+          s":$name"
+        else
+          displayParams.mkString("", ",", s",:$name")
       expr match {
         case q: Query =>
           // If the displayExpr is a query type, then the rewrite is simulating an
           // aggregate function. Modifications to the aggregate need to be represented
           // after the operation as part of the expression string. There are two
           // categories: offsets applied to the data function and group by.
           val buffer = new java.lang.StringBuilder
-          buffer.append(s"$q,:$name")
+          buffer.append(s"$q,$op")
           getOffset(evalExpr).foreach(d => buffer.append(s",$d,:offset"))
 
           val grouping = evalExpr.finalGrouping
@@ -1022,14 +1028,14 @@ object MathExpr {
           // would alter the toString behavior. So the grouping match check is based on the
           // original display expression.
           val evalOffset = getOffset(evalExpr)
-          evalOffset.fold(s"$t,:$name") { d =>
+          evalOffset.fold(s"$t,$op") { d =>
             val displayOffset = getOffset(t).getOrElse(Duration.ZERO)
-            if (d != displayOffset) s"${t.withOffset(d)},:$name" else s"$t,:$name"
+            if (d != displayOffset) s"${t.withOffset(d)},$op" else s"$t,$op"
           }
         case _ =>
           val grouping = evalExpr.finalGrouping
           val by = if (grouping.nonEmpty) grouping.mkString(",(,", ",", ",),:by") else ""
-          s"$expr,:$name$by"
+          s"$expr,$op$by"
       }
     }
 

atlas-core/src/main/scala/com/netflix/atlas/core/model/MathVocabulary.scala

@@ -505,9 +505,9 @@ object MathVocabulary extends Vocabulary {
 
     private def addCommonKeys(expr: TimeSeriesExpr, keys: List[String]): TimeSeriesExpr = {
       val newExpr = expr.rewrite {
-        case nr @ MathExpr.NamedRewrite(_, _: Query, e, _, _) if e.isGrouped =>
+        case nr @ MathExpr.NamedRewrite(_, _: Query, _, e, _, _) if e.isGrouped =>
           nr.copy(evalExpr = addCommonKeys(e, keys))
-        case nr @ MathExpr.NamedRewrite(_, _: Query, _, _, _) =>
+        case nr @ MathExpr.NamedRewrite(_, _: Query, _, _, _, _) =>
           nr.groupBy(keys)
         case af: AggregateFunction =>
           DataExpr.GroupBy(af, keys)
@@ -612,11 +612,11 @@ object MathVocabulary extends Vocabulary {
       case (n: String) :: TimeSeriesType(rw) :: (orig: Expr) :: stack =>
         // If the original is already an expr type, e.g. a Query, then we should
         // preserve it without modification. So we first match for Expr.
-        MathExpr.NamedRewrite(n, orig, rw, context) :: stack
+        MathExpr.NamedRewrite(n, orig, Nil, rw, context) :: stack
       case (n: String) :: TimeSeriesType(rw) :: TimeSeriesType(orig) :: stack =>
         // This is a more general match that will coerce the original into a
         // TimeSeriesExpr if it is not one already, e.g., a constant.
-        MathExpr.NamedRewrite(n, orig, rw, context) :: stack
+        MathExpr.NamedRewrite(n, orig, Nil, rw, context) :: stack
     }
 
     override def summary: String =
@@ -1284,7 +1284,8 @@ object MathVocabulary extends Vocabulary {
         case DoubleType(max) :: DoubleType(min) :: (q: Query) :: stack =>
           val rangeQuery = q.and(bucketQuery(min, max))
           val expr = DataExpr.Sum(rangeQuery)
-          context.copy(stack = expr :: stack)
+          val nr = MathExpr.NamedRewrite(name, q, List(min, max), expr, context)
+          context.copy(stack = nr :: stack)
         case _ =>
           invalidStack
       }

atlas-core/src/test/scala/com/netflix/atlas/core/model/CustomVocabularySuite.scala

@@ -63,31 +63,31 @@ class CustomVocabularySuite extends FunSuite {
 
   test("simple average") {
     val expr = eval(s"$cpuUser,:node-avg").rewrite {
-      case MathExpr.NamedRewrite("node-avg", _, e, _, _) => e
+      case MathExpr.NamedRewrite("node-avg", _, _, e, _, _) => e
     }
     val expected = eval(s"$cpuUser,:sum,$numInstances,:sum,:div")
     assertEquals(expr, expected)
   }
 
   test("expr with cluster") {
     val expr = eval(s"$cpuUser,cluster,foo,:eq,:and,:node-avg").rewrite {
-      case MathExpr.NamedRewrite("node-avg", _, e, _, _) => e
+      case MathExpr.NamedRewrite("node-avg", _, _, e, _, _) => e
     }
     val expected = eval(s"$cpuUser,:sum,$numInstances,:sum,:div,cluster,foo,:eq,:cq")
     assertEquals(expr, expected)
   }
 
   test("expr with cq using non-infrastructure tags") {
     val expr = eval(s"$cpuUser,:node-avg,core,1,:eq,:cq").rewrite {
-      case MathExpr.NamedRewrite("node-avg", _, e, _, _) => e
+      case MathExpr.NamedRewrite("node-avg", _, _, e, _, _) => e
     }
     val expected = eval(s"$cpuUser,core,1,:eq,:and,:sum,$numInstances,:sum,:div")
     assertEquals(expr, expected)
   }
 
   test("expr grouped by infrastructure tags") {
     val expr = eval(s"$cpuUser,cluster,foo,:eq,:and,:node-avg,(,zone,),:by").rewrite {
-      case MathExpr.NamedRewrite("node-avg", _, e, _, _) => e
+      case MathExpr.NamedRewrite("node-avg", _, _, e, _, _) => e
     }
     val expected =
       eval(s"$cpuUser,:sum,(,zone,),:by,$numInstances,:sum,(,zone,),:by,:div,cluster,foo,:eq,:cq")
@@ -96,7 +96,7 @@ class CustomVocabularySuite extends FunSuite {
 
   test("expr grouped by non-infrastructure tags") {
     val expr = eval(s"$cpuUser,cluster,foo,:eq,:and,:node-avg,(,name,),:by").rewrite {
-      case MathExpr.NamedRewrite("node-avg", _, e, _, _) => e
+      case MathExpr.NamedRewrite("node-avg", _, _, e, _, _) => e
     }
     val expected = eval(s"$cpuUser,:sum,(,name,),:by,$numInstances,:sum,:div,cluster,foo,:eq,:cq")
     assertEquals(expr, expected)
@@ -105,7 +105,7 @@ class CustomVocabularySuite extends FunSuite {
   test("expr grouped by non-infrastructure tags with offset") {
     val displayExpr = eval(s"$cpuUser,cluster,foo,:eq,:and,:node-avg,(,name,),:by,1h,:offset")
     val evalExpr = displayExpr.rewrite {
-      case MathExpr.NamedRewrite("node-avg", _, e, _, _) => e
+      case MathExpr.NamedRewrite("node-avg", _, _, e, _, _) => e
     }
     val expected = eval(
       s"$cpuUser,:sum,(,name,),:by,PT1H,:offset,$numInstances,:sum,PT1H,:offset,:div,cluster,foo,:eq,:cq"
@@ -131,7 +131,7 @@ class CustomVocabularySuite extends FunSuite {
 
   test("expr with not") {
     val expr = eval(s"$cpuUser,foo,bar,:eq,:not,:and,cluster,foo,:eq,:and,:node-avg").rewrite {
-      case MathExpr.NamedRewrite("node-avg", _, e, _, _) => e
+      case MathExpr.NamedRewrite("node-avg", _, _, e, _, _) => e
     }
     val expected =
       eval(s"$cpuUser,foo,bar,:eq,:not,:and,:sum,$numInstances,:sum,:div,cluster,foo,:eq,:cq")
@@ -140,7 +140,7 @@ class CustomVocabularySuite extends FunSuite {
 
   test("group by mixed keys") {
     val expr = eval("name,(,a,b,c,),:in,app,beacon,:eq,:and,:node-avg,(,name,asg,),:by").rewrite {
-      case MathExpr.NamedRewrite("node-avg", _, e, _, _) => e
+      case MathExpr.NamedRewrite("node-avg", _, _, e, _, _) => e
     }
     val expected = eval(
       s"name,(,a,b,c,),:in,:sum,(,name,asg,),:by,$numInstances,:sum,(,asg,),:by,:div,app,beacon,:eq,:cq"

atlas-eval/src/main/scala/com/netflix/atlas/eval/graph/SimpleLegends.scala

@@ -118,7 +118,7 @@ object SimpleLegends extends StrictLogging {
     // a simple aggregate like sum based on the display expression.
     expr
       .rewrite {
-        case MathExpr.NamedRewrite(n, q: Query, evalExpr, _, _) if n.endsWith("-avg") =>
+        case MathExpr.NamedRewrite(n, q: Query, _, evalExpr, _, _) if n.endsWith("-avg") =>
           val aggr = DataExpr.Sum(q)
           if (evalExpr.isGrouped) DataExpr.GroupBy(aggr, evalExpr.finalGrouping) else aggr
       }

atlas-webapi/src/test/scala/com/netflix/atlas/webapi/ExprApiSuite.scala

@@ -137,6 +137,18 @@ class ExprApiSuite extends MUnitRouteSuite {
     assertEquals(data, List("name,(,sps1,sps2,),:in,:sum"))
   }
 
+  testGet("/api/v1/expr/normalize?q=name,latency,:eq,0,5,:sample-count") {
+    assertEquals(response.status, StatusCodes.OK)
+    val data = Json.decode[List[String]](responseAs[String])
+    assertEquals(data, List("name,latency,:eq,0.0,5.0,:sample-count"))
+  }
+
+  testGet("/api/v1/expr/normalize?q=name,latency,:eq,0,5,:sample-count,(,app,),:by") {
+    assertEquals(response.status, StatusCodes.OK)
+    val data = Json.decode[List[String]](responseAs[String])
+    assertEquals(data, List("name,latency,:eq,0.0,5.0,:sample-count,(,app,),:by"))
+  }
+
   testGet(
     "/api/v1/expr/normalize?q=(,name,:swap,:eq,nf.cluster,foo,:eq,:and,:sum,),foo,:sset,cpu,foo,:fcall,disk,foo,:fcall"
   ) {