feat(notification/check): trim yield from flux if provided

notification/check/threshold.go

@@ -152,6 +152,28 @@ func replaceDurationsWithEvery(pkg *ast.Package, every *notification.Duration) {
 	})
 }
 
+// TODO(desa): we'll likely want to remove all other arguments to range that are provided, but for now this should work.
+// When we decide to implement the full feature we'll have to do something more sophisticated.
+func removeStopFromRange(pkg *ast.Package) {
+	ast.Visit(pkg, func(n ast.Node) {
+		if call, ok := n.(*ast.CallExpression); ok {
+			if id, ok := call.Callee.(*ast.Identifier); ok && id.Name == "range" {
+				for _, args := range call.Arguments {
+					if obj, ok := args.(*ast.ObjectExpression); ok {
+						props := obj.Properties[:0]
+						for _, prop := range obj.Properties {
+							if prop.Key.Key() == "start" {
+								props = append(props, prop)
+							}
+						}
+						obj.Properties = props
+					}
+				}
+			}
+		}
+	})
+}
+
 func assignPipelineToData(f *ast.File) error {
 	if len(f.Body) != 1 {
 		return fmt.Errorf("expected there to be a single statement in the flux script body, recieved %d", len(f.Body))
@@ -164,7 +186,18 @@ func assignPipelineToData(f *ast.File) error {
 		return fmt.Errorf("statement is not an *ast.Expression statement, recieved %T", stmt)
 	}
 
-	f.Body[0] = flux.DefineVariable("data", e.Expression)
+	exp := e.Expression
+
+	pipe, ok := exp.(*ast.PipeExpression)
+	if !ok {
+		return fmt.Errorf("expression is not an *ast.PipeExpression statement, recieved %T", exp)
+	}
+
+	if id, ok := pipe.Call.Callee.(*ast.Identifier); ok && id.Name == "yield" {
+		exp = pipe.Argument
+	}
+
+	f.Body[0] = flux.DefineVariable("data", exp)
 	return nil
 }
 
@@ -173,6 +206,7 @@ func assignPipelineToData(f *ast.File) error {
 func (t Threshold) GenerateFluxASTReal() (*ast.Package, error) {
 	p := parser.ParseSource(t.Query.Text)
 	replaceDurationsWithEvery(p, t.Every)
+	removeStopFromRange(p)
 
 	if errs := ast.GetErrors(p); len(errs) != 0 {
 		return nil, multiError(errs)

notification/check/threshold_test.go

@@ -26,7 +26,185 @@ func TestThreshold_GenerateFlux(t *testing.T) {
 		wants wants
 	}{
 		{
-			name: "all levels",
+			name: "all levels with yield and stop",
+			args: args{
+				threshold: check.Threshold{
+					Base: check.Base{
+						ID:   10,
+						Name: "moo",
+						Tags: []notification.Tag{
+							{Key: "aaa", Value: "vaaa"},
+							{Key: "bbb", Value: "vbbb"},
+						},
+						Every:                 mustDuration("1h"),
+						StatusMessageTemplate: "whoa! {check.yeah}",
+						Query: influxdb.DashboardQuery{
+							Text: `from(bucket: "foo") |> range(start: -1d, stop: now()) |> aggregateWindow(every: 1m, fn: mean) |> yield()`,
+						},
+					},
+					Thresholds: []check.ThresholdConfig{
+						check.Greater{
+							ThresholdConfigBase: check.ThresholdConfigBase{
+								Level: notification.Ok,
+							},
+							Value: l,
+						},
+						check.Lesser{
+							ThresholdConfigBase: check.ThresholdConfigBase{
+								Level: notification.Info,
+							},
+							Value: u,
+						},
+						check.Range{
+							ThresholdConfigBase: check.ThresholdConfigBase{
+								Level: notification.Warn,
+							},
+							Min:    l,
+							Max:    u,
+							Within: true,
+						},
+						check.Range{
+							ThresholdConfigBase: check.ThresholdConfigBase{
+								Level: notification.Critical,
+							},
+							Min:    l,
+							Max:    u,
+							Within: true,
+						},
+					},
+				},
+			},
+			wants: wants{
+				script: `package main
+import "influxdata/influxdb/alerts"
+import "influxdata/influxdb/v1"
+
+data = from(bucket: "foo")
+	|> range(start: -1h)
+	|> aggregateWindow(every: 1h, fn: mean)
+
+option task = {name: "moo", every: 1h}
+
+check = {
+	_check_id: "000000000000000a",
+	_check_name: "moo",
+	_check_type: "threshold",
+	tags: {aaa: "vaaa", bbb: "vbbb"},
+}
+ok = (r) =>
+	(r._value > 10.0)
+info = (r) =>
+	(r._value < 40.0)
+warn = (r) =>
+	(r._value < 40.0 and r._value > 10.0)
+crit = (r) =>
+	(r._value < 40.0 and r._value > 10.0)
+messageFn = (r, check) =>
+	("whoa! {check.yeah}")
+
+data
+	|> v1.fieldsAsCols()
+	|> alerts.check(
+		data: check,
+		messageFn: messageFn,
+		ok: ok,
+		info: info,
+		warn: warn,
+		crit: crit,
+	)`,
+			},
+		},
+		{
+			name: "all levels with yield",
+			args: args{
+				threshold: check.Threshold{
+					Base: check.Base{
+						ID:   10,
+						Name: "moo",
+						Tags: []notification.Tag{
+							{Key: "aaa", Value: "vaaa"},
+							{Key: "bbb", Value: "vbbb"},
+						},
+						Every:                 mustDuration("1h"),
+						StatusMessageTemplate: "whoa! {check.yeah}",
+						Query: influxdb.DashboardQuery{
+							Text: `from(bucket: "foo") |> range(start: -1d) |> aggregateWindow(every: 1m, fn: mean) |> yield()`,
+						},
+					},
+					Thresholds: []check.ThresholdConfig{
+						check.Greater{
+							ThresholdConfigBase: check.ThresholdConfigBase{
+								Level: notification.Ok,
+							},
+							Value: l,
+						},
+						check.Lesser{
+							ThresholdConfigBase: check.ThresholdConfigBase{
+								Level: notification.Info,
+							},
+							Value: u,
+						},
+						check.Range{
+							ThresholdConfigBase: check.ThresholdConfigBase{
+								Level: notification.Warn,
+							},
+							Min:    l,
+							Max:    u,
+							Within: true,
+						},
+						check.Range{
+							ThresholdConfigBase: check.ThresholdConfigBase{
+								Level: notification.Critical,
+							},
+							Min:    l,
+							Max:    u,
+							Within: true,
+						},
+					},
+				},
+			},
+			wants: wants{
+				script: `package main
+import "influxdata/influxdb/alerts"
+import "influxdata/influxdb/v1"
+
+data = from(bucket: "foo")
+	|> range(start: -1h)
+	|> aggregateWindow(every: 1h, fn: mean)
+
+option task = {name: "moo", every: 1h}
+
+check = {
+	_check_id: "000000000000000a",
+	_check_name: "moo",
+	_check_type: "threshold",
+	tags: {aaa: "vaaa", bbb: "vbbb"},
+}
+ok = (r) =>
+	(r._value > 10.0)
+info = (r) =>
+	(r._value < 40.0)
+warn = (r) =>
+	(r._value < 40.0 and r._value > 10.0)
+crit = (r) =>
+	(r._value < 40.0 and r._value > 10.0)
+messageFn = (r, check) =>
+	("whoa! {check.yeah}")
+
+data
+	|> v1.fieldsAsCols()
+	|> alerts.check(
+		data: check,
+		messageFn: messageFn,
+		ok: ok,
+		info: info,
+		warn: warn,
+		crit: crit,
+	)`,
+			},
+		},
+		{
+			name: "all levels without yield",
 			args: args{
 				threshold: check.Threshold{
 					Base: check.Base{