fix: starlark pop operation for non-existing keys (#9954)

influxdata · Oct 21, 2021 · 112ef7f · 112ef7f
1 parent e50b415
commit 112ef7f
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 0 deletions.
diff --git a/plugins/processors/starlark/field_dict.go b/plugins/processors/starlark/field_dict.go
@@ -175,6 +175,7 @@ func (d FieldDict) Delete(k starlark.Value) (v starlark.Value, found bool, err e
 			sv, err := asStarlarkValue(value)
 			return sv, ok, err
 		}
+		return starlark.None, false, nil
 	}
 
 	return starlark.None, false, errors.New("key must be of type 'str'")

diff --git a/plugins/processors/starlark/starlark_test.go b/plugins/processors/starlark/starlark_test.go
@@ -705,6 +705,49 @@ def apply(metric):
 				),
 			},
 		},
+		{
+			name: "pop tag (default)",
+			source: `
+def apply(metric):
+	metric.tags['host2'] = metric.tags.pop('url', 'foo.org')
+	return metric
+`,
+			input: []telegraf.Metric{
+				testutil.MustMetric("cpu",
+					map[string]string{
+						"host": "example.org",
+					},
+					map[string]interface{}{"time_idle": 0},
+					time.Unix(0, 0),
+				),
+				testutil.MustMetric("cpu",
+					map[string]string{
+						"host": "example.org",
+						"url":  "bar.org",
+					},
+					map[string]interface{}{"time_idle": 0},
+					time.Unix(0, 0),
+				),
+			},
+			expected: []telegraf.Metric{
+				testutil.MustMetric("cpu",
+					map[string]string{
+						"host":  "example.org",
+						"host2": "foo.org",
+					},
+					map[string]interface{}{"time_idle": 0},
+					time.Unix(0, 0),
+				),
+				testutil.MustMetric("cpu",
+					map[string]string{
+						"host":  "example.org",
+						"host2": "bar.org",
+					},
+					map[string]interface{}{"time_idle": 0},
+					time.Unix(0, 0),
+				),
+			},
+		},
 		{
 			name: "popitem tags",
 			source: `
@@ -1773,6 +1816,53 @@ def apply(metric):
 				),
 			},
 		},
+		{
+			name: "pop field (default)",
+			source: `
+def apply(metric):
+	metric.fields['idle_count'] = metric.fields.pop('count', 10)
+	return metric
+`,
+			input: []telegraf.Metric{
+				testutil.MustMetric("cpu",
+					map[string]string{},
+					map[string]interface{}{
+						"time_idle":  0,
+						"time_guest": 0,
+					},
+					time.Unix(0, 0),
+				),
+				testutil.MustMetric("cpu",
+					map[string]string{},
+					map[string]interface{}{
+						"time_idle":  0,
+						"time_guest": 0,
+						"count":      0,
+					},
+					time.Unix(0, 0),
+				),
+			},
+			expected: []telegraf.Metric{
+				testutil.MustMetric("cpu",
+					map[string]string{},
+					map[string]interface{}{
+						"time_idle":  0,
+						"time_guest": 0,
+						"idle_count": 10,
+					},
+					time.Unix(0, 0),
+				),
+				testutil.MustMetric("cpu",
+					map[string]string{},
+					map[string]interface{}{
+						"time_idle":  0,
+						"time_guest": 0,
+						"idle_count": 0,
+					},
+					time.Unix(0, 0),
+				),
+			},
+		},
 		{
 			name: "popitem field",
 			source: `

diff --git a/plugins/processors/starlark/tag_dict.go b/plugins/processors/starlark/tag_dict.go
@@ -162,6 +162,7 @@ func (d TagDict) Delete(k starlark.Value) (v starlark.Value, found bool, err err
 			v := starlark.String(value)
 			return v, ok, err
 		}
+		return starlark.None, false, nil
 	}
 
 	return starlark.None, false, errors.New("key must be of type 'str'")