dom/telegraf
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Add Base64-encoding/decoding for Google Cloud PubSub plugins (#5543)

Browse Source
This commit is contained in:
emily
2019-03-06 17:34:47 -08:00
committed by Daniel Nelson
parent 51cc0fe6d8
commit dd67144660
7 changed files with 143 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
plugins/outputs/cloud_pubsub/README.md
View File

@@ -52,6 +52,9 @@ generate it using `telegraf --usage cloud_pubsub`.
## Optional. Specifies a timeout for requests to the PubSub API.
# publish_timeout = "30s"
## Optional. If true, published PubSub message data will be base64-encoded.
# base64_data = false
## Optional. PubSub attributes to add to metrics.
# [[inputs.pubsub.attributes]]
# my_attr = "tag_value"

17
plugins/outputs/cloud_pubsub/pubsub.go
View File

@@ -6,6 +6,7 @@ import (
"sync"
"cloud.google.com/go/pubsub"
"encoding/base64"
"github.com/influxdata/telegraf"
"github.com/influxdata/telegraf/internal"
"github.com/influxdata/telegraf/plugins/outputs"
@@ -56,6 +57,9 @@ const sampleConfig = `
## Optional. Specifies a timeout for requests to the PubSub API.
# publish_timeout = "30s"
## Optional. If true, published PubSub message data will be base64-encoded.
# base64_data = false
## Optional. PubSub attributes to add to metrics.
# [[inputs.pubsub.attributes]]
# my_attr = "tag_value"
@@ -72,6 +76,7 @@ type PubSub struct {
PublishByteThreshold int `toml:"publish_byte_threshold"`
PublishNumGoroutines int `toml:"publish_num_go_routines"`
PublishTimeout internal.Duration `toml:"publish_timeout"`
Base64Data bool `toml:"base64_data"`
t topic
c *pubsub.Client
@@ -207,6 +212,12 @@ func (ps *PubSub) toMessages(metrics []telegraf.Metric) ([]*pubsub.Message, erro
if err != nil {
return nil, err
}
if ps.Base64Data {
encoded := base64.StdEncoding.EncodeToString(b)
b = []byte(encoded)
}
msg := &pubsub.Message{Data: b}
if ps.Attributes != nil {
msg.Attributes = ps.Attributes
@@ -220,6 +231,12 @@ func (ps *PubSub) toMessages(metrics []telegraf.Metric) ([]*pubsub.Message, erro
if err != nil {
return nil, err
}
if ps.Base64Data {
encoded := base64.StdEncoding.EncodeToString(b)
b = []byte(encoded)
}
msgs[i] = &pubsub.Message{
Data: b,
}

50
plugins/outputs/cloud_pubsub/pubsub_test.go
View File

@@ -4,6 +4,7 @@ import (
"testing"
"cloud.google.com/go/pubsub"
"encoding/base64"
"github.com/influxdata/telegraf"
"github.com/influxdata/telegraf/plugins/parsers"
"github.com/influxdata/telegraf/testutil"
@@ -26,7 +27,7 @@ func TestPubSub_WriteSingle(t *testing.T) {
}
for _, testM := range testMetrics {
verifyMetricPublished(t, testM.m, topic.published)
verifyRawMetricPublished(t, testM.m, topic.published)
}
}
@@ -48,7 +49,7 @@ func TestPubSub_WriteWithAttribute(t *testing.T) {
}
for _, testM := range testMetrics {
msg := verifyMetricPublished(t, testM.m, topic.published)
msg := verifyRawMetricPublished(t, testM.m, topic.published)
assert.Equalf(t, "bar1", msg.Attributes["foo1"], "expected attribute foo1=bar1")
assert.Equalf(t, "bar2", msg.Attributes["foo2"], "expected attribute foo2=bar2")
}
@@ -70,7 +71,7 @@ func TestPubSub_WriteMultiple(t *testing.T) {
}
for _, testM := range testMetrics {
verifyMetricPublished(t, testM.m, topic.published)
verifyRawMetricPublished(t, testM.m, topic.published)
}
assert.Equalf(t, 1, topic.bundleCount, "unexpected bundle count")
}
@@ -94,7 +95,7 @@ func TestPubSub_WriteOverCountThreshold(t *testing.T) {
}
for _, testM := range testMetrics {
verifyMetricPublished(t, testM.m, topic.published)
verifyRawMetricPublished(t, testM.m, topic.published)
}
assert.Equalf(t, 2, topic.bundleCount, "unexpected bundle count")
}
@@ -117,11 +118,33 @@ func TestPubSub_WriteOverByteThreshold(t *testing.T) {
}
for _, testM := range testMetrics {
verifyMetricPublished(t, testM.m, topic.published)
verifyRawMetricPublished(t, testM.m, topic.published)
}
assert.Equalf(t, 2, topic.bundleCount, "unexpected bundle count")
}
func TestPubSub_WriteBase64Single(t *testing.T) {
testMetrics := []testMetric{
{testutil.TestMetric("value_1", "test"), false /*return error */},
{testutil.TestMetric("value_2", "test"), false},
}
settings := pubsub.DefaultPublishSettings
settings.CountThreshold = 1
ps, topic, metrics := getTestResources(t, settings, testMetrics)
ps.Base64Data = true
err := ps.Write(metrics)
if err != nil {
t.Fatalf("got unexpected error: %v", err)
}
for _, testM := range testMetrics {
verifyMetricPublished(t, testM.m, topic.published, true /* base64encoded */)
}
}
func TestPubSub_Error(t *testing.T) {
testMetrics := []testMetric{
// Force this batch to return error
@@ -141,7 +164,11 @@ func TestPubSub_Error(t *testing.T) {
}
}
func verifyMetricPublished(t *testing.T, m telegraf.Metric, published map[string]*pubsub.Message) *pubsub.Message {
func verifyRawMetricPublished(t *testing.T, m telegraf.Metric, published map[string]*pubsub.Message) *pubsub.Message {
return verifyMetricPublished(t, m, published, false)
}
func verifyMetricPublished(t *testing.T, m telegraf.Metric, published map[string]*pubsub.Message, base64Encoded bool) *pubsub.Message {
p, _ := parsers.NewInfluxParser()
v, _ := m.GetField("value")
@@ -150,7 +177,16 @@ func verifyMetricPublished(t *testing.T, m telegraf.Metric, published map[string
t.Fatalf("expected metric to get published (value: %s)", v.(string))
}
parsed, err := p.Parse(psMsg.Data)
data := psMsg.Data
if base64Encoded {
v, err := base64.StdEncoding.DecodeString(string(psMsg.Data))
if err != nil {
t.Fatalf("Unable to decode expected base64-encoded message: %s", err)
}
data = []byte(v)
}
parsed, err := p.Parse(data)
if err != nil {
t.Fatalf("could not parse influxdb metric from published message: %s", string(psMsg.Data))
}

12
plugins/outputs/cloud_pubsub/topic_stubbed.go
View File

@@ -10,6 +10,7 @@ import (
"time"
"cloud.google.com/go/pubsub"
"encoding/base64"
"github.com/influxdata/telegraf"
"github.com/influxdata/telegraf/internal"
"github.com/influxdata/telegraf/plugins/parsers"
@@ -180,8 +181,17 @@ func (t *stubTopic) parseIDs(msg *pubsub.Message) []string {
p, _ := parsers.NewInfluxParser()
metrics, err := p.Parse(msg.Data)
if err != nil {
t.Fatalf("unexpected parsing error: %v", err)
// Just attempt to base64-decode first before returning error.
d, err := base64.StdEncoding.DecodeString(string(msg.Data))
if err != nil {
t.Errorf("unable to base64-decode potential test message: %v", err)
}
metrics, err = p.Parse(d)
if err != nil {
t.Fatalf("unexpected parsing error: %v", err)
}
}
ids := make([]string, len(metrics))
for i, met := range metrics {
id, _ := met.GetField("value")