Add tag based routing in influxdb/influxdb_v2 outputs (#5490)

plugins/outputs/influxdb_v2/README.md

@@ -22,6 +22,10 @@ The InfluxDB output plugin writes metrics to the [InfluxDB v2.x] HTTP service.
   ## Destination bucket to write into.
   bucket = ""
 
+  ## The value of this tag will be used to determine the bucket.  If this
+  ## tag is not set the 'bucket' option is used as the default.
+  # bucket_tag = ""
+
   ## Timeout for HTTP messages.
   # timeout = "5s"
 

plugins/outputs/influxdb_v2/http.go

@@ -20,13 +20,10 @@ import (
 	"github.com/influxdata/telegraf/plugins/serializers/influx"
 )
 
-type APIErrorType int
-
 type APIError struct {
 	StatusCode  int
 	Title       string
 	Description string
-	Type        APIErrorType
 }
 
 func (e APIError) Error() string {
@@ -47,6 +44,7 @@ type HTTPConfig struct {
 	Token           string
 	Organization    string
 	Bucket          string
+	BucketTag       string
 	Timeout         time.Duration
 	Headers         map[string]string
 	Proxy           *url.URL
@@ -58,10 +56,12 @@ type HTTPConfig struct {
 }
 
 type httpClient struct {
-	WriteURL        string
 	ContentEncoding string
 	Timeout         time.Duration
 	Headers         map[string]string
+	Organization    string
+	Bucket          string
+	BucketTag       string
 
 	client     *http.Client
 	serializer *influx.Serializer
@@ -103,14 +103,6 @@ func NewHTTPClient(config *HTTPConfig) (*httpClient, error) {
 		serializer = influx.NewSerializer()
 	}
 
-	writeURL, err := makeWriteURL(
-		*config.URL,
-		config.Organization,
-		config.Bucket)
-	if err != nil {
-		return nil, err
-	}
-
 	var transport *http.Transport
 	switch config.URL.Scheme {
 	case "http", "https":
@@ -139,10 +131,12 @@ func NewHTTPClient(config *HTTPConfig) (*httpClient, error) {
 			Transport: transport,
 		},
 		url:             config.URL,
-		WriteURL:        writeURL,
 		ContentEncoding: config.ContentEncoding,
 		Timeout:         timeout,
 		Headers:         headers,
+		Organization:    config.Organization,
+		Bucket:          config.Bucket,
+		BucketTag:       config.BucketTag,
 	}
 	return client, nil
 }
@@ -173,8 +167,45 @@ func (c *httpClient) Write(ctx context.Context, metrics []telegraf.Metric) error
 	if c.retryTime.After(time.Now()) {
 		return errors.New("Retry time has not elapsed")
 	}
+
+	batches := make(map[string][]telegraf.Metric)
+	if c.BucketTag == "" {
+		err := c.writeBatch(ctx, c.Bucket, metrics)
+		if err != nil {
+			return err
+		}
+	} else {
+		for _, metric := range metrics {
+			bucket, ok := metric.GetTag(c.BucketTag)
+			if !ok {
+				bucket = c.Bucket
+			}
+
+			if _, ok := batches[bucket]; !ok {
+				batches[bucket] = make([]telegraf.Metric, 0)
+			}
+
+			batches[bucket] = append(batches[bucket], metric)
+		}
+
+		for bucket, batch := range batches {
+			err := c.writeBatch(ctx, bucket, batch)
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func (c *httpClient) writeBatch(ctx context.Context, bucket string, metrics []telegraf.Metric) error {
+	url, err := makeWriteURL(*c.url, c.Organization, bucket)
+	if err != nil {
+		return err
+	}
+
 	reader := influx.NewReader(metrics, c.serializer)
-	req, err := c.makeWriteRequest(reader)
+	req, err := c.makeWriteRequest(url, reader)
 	if err != nil {
 		return err
 	}
@@ -227,7 +258,7 @@ func (c *httpClient) Write(ctx context.Context, metrics []telegraf.Metric) error
 	}
 }
 
-func (c *httpClient) makeWriteRequest(body io.Reader) (*http.Request, error) {
+func (c *httpClient) makeWriteRequest(url string, body io.Reader) (*http.Request, error) {
 	var err error
 	if c.ContentEncoding == "gzip" {
 		body, err = internal.CompressWithGzip(body)
@@ -236,7 +267,7 @@ func (c *httpClient) makeWriteRequest(body io.Reader) (*http.Request, error) {
 		}
 	}
 
-	req, err := http.NewRequest("POST", c.WriteURL, body)
+	req, err := http.NewRequest("POST", url, body)
 	if err != nil {
 		return nil, err
 	}

plugins/outputs/influxdb_v2/http_internal_test.go

@@ -1,7 +1,6 @@
 package influxdb_v2
 
 import (
-	"io"
 	"net/url"
 	"testing"
 
@@ -46,14 +45,3 @@ func TestMakeWriteURL(t *testing.T) {
 		}
 	}
 }
-
-func TestMakeWriteRequest(t *testing.T) {
-	reader, _ := io.Pipe()
-	cli := httpClient{
-		WriteURL:        "http://localhost:9999/v2/write?bucket=telegraf&org=influx",
-		ContentEncoding: "gzip",
-		Headers:         map[string]string{"x": "y"},
-	}
-	_, err := cli.makeWriteRequest(reader)
-	require.NoError(t, err)
-}

plugins/outputs/influxdb_v2/influxdb.go

@@ -38,6 +38,10 @@ var sampleConfig = `
   ## Destination bucket to write into.
   bucket = ""
 
+  ## The value of this tag will be used to determine the bucket.  If this
+  ## tag is not set the 'bucket' option is used as the default.
+  # bucket_tag = ""
+
   ## Timeout for HTTP messages.
   # timeout = "5s"
 
@@ -77,6 +81,7 @@ type InfluxDB struct {
 	Token           string            `toml:"token"`
 	Organization    string            `toml:"organization"`
 	Bucket          string            `toml:"bucket"`
+	BucketTag       string            `toml:"bucket_tag"`
 	Timeout         internal.Duration `toml:"timeout"`
 	HTTPHeaders     map[string]string `toml:"http_headers"`
 	HTTPProxy       string            `toml:"http_proxy"`
@@ -174,6 +179,7 @@ func (i *InfluxDB) getHTTPClient(ctx context.Context, url *url.URL, proxy *url.U
 		Token:           i.Token,
 		Organization:    i.Organization,
 		Bucket:          i.Bucket,
+		BucketTag:       i.BucketTag,
 		Timeout:         i.Timeout.Duration,
 		Headers:         i.HTTPHeaders,
 		Proxy:           proxy,