Add support for setting retention policy using tag (#7141)

plugins/outputs/influxdb/README.md

@@ -35,6 +35,13 @@ The InfluxDB output plugin writes metrics to the [InfluxDB v1.x] HTTP or UDP ser
   ## the default retention policy.  Only takes effect when using HTTP.
   # retention_policy = ""
 
+  ## The value of this tag will be used to determine the retention policy.  If this
+  ## tag is not set the 'retention_policy' option is used as the default.
+  # retention_policy_tag = ""
+
+  ## If true, the 'retention_policy_tag' will not be removed from the metric.
+  # exclude_retention_policy_tag = false
+
   ## Write consistency (clusters only), can be: "any", "one", "quorum", "all".
   ## Only takes effect when using HTTP.
   # write_consistency = "any"

plugins/outputs/influxdb/http.go

@@ -83,21 +83,23 @@ func (r WriteResponse) Error() string {
 }
 
 type HTTPConfig struct {
-	URL                  *url.URL
+	URL                       *url.URL
-	UserAgent            string
+	UserAgent                 string
-	Timeout              time.Duration
+	Timeout                   time.Duration
-	Username             string
+	Username                  string
-	Password             string
+	Password                  string
-	TLSConfig            *tls.Config
+	TLSConfig                 *tls.Config
-	Proxy                *url.URL
+	Proxy                     *url.URL
-	Headers              map[string]string
+	Headers                   map[string]string
-	ContentEncoding      string
+	ContentEncoding           string
-	Database             string
+	Database                  string
-	DatabaseTag          string
+	DatabaseTag               string
-	ExcludeDatabaseTag   bool
+	ExcludeDatabaseTag        bool
-	RetentionPolicy      string
+	RetentionPolicy           string
-	Consistency          string
+	RetentionPolicyTag        string
-	SkipDatabaseCreation bool
+	ExcludeRetentionPolicyTag bool
+	Consistency               string
+	SkipDatabaseCreation      bool
 
 	InfluxUintSupport bool `toml:"influx_uint_support"`
 	Serializer        *influx.Serializer
@@ -236,55 +238,66 @@ func (c *httpClient) CreateDatabase(ctx context.Context, database string) error
 	}
 }
 
+type dbrp struct {
+	Database        string
+	RetentionPolicy string
+}
+
 // Write sends the metrics to InfluxDB
 func (c *httpClient) Write(ctx context.Context, metrics []telegraf.Metric) error {
-	batches := make(map[string][]telegraf.Metric)
+	// If these options are not used, we can skip in plugin batching and send
-	if c.config.DatabaseTag == "" {
+	// the full batch in a single request.
-		err := c.writeBatch(ctx, c.config.Database, metrics)
+	if c.config.DatabaseTag == "" && c.config.RetentionPolicyTag == "" {
+		return c.writeBatch(ctx, c.config.Database, c.config.RetentionPolicy, metrics)
+	}
+
+	batches := make(map[dbrp][]telegraf.Metric)
+	for _, metric := range metrics {
+		db, ok := metric.GetTag(c.config.DatabaseTag)
+		if !ok {
+			db = c.config.Database
+		}
+
+		rp, ok := metric.GetTag(c.config.RetentionPolicyTag)
+		if !ok {
+			rp = c.config.RetentionPolicy
+		}
+
+		dbrp := dbrp{
+			Database:        db,
+			RetentionPolicy: rp,
+		}
+
+		if c.config.ExcludeDatabaseTag || c.config.ExcludeRetentionPolicyTag {
+			// Avoid modifying the metric in case we need to retry the request.
+			metric = metric.Copy()
+			metric.Accept()
+			metric.RemoveTag(c.config.DatabaseTag)
+			metric.RemoveTag(c.config.RetentionPolicyTag)
+		}
+
+		batches[dbrp] = append(batches[dbrp], metric)
+	}
+
+	for dbrp, batch := range batches {
+		if !c.config.SkipDatabaseCreation && !c.createdDatabases[dbrp.Database] {
+			err := c.CreateDatabase(ctx, dbrp.Database)
+			if err != nil {
+				c.log.Warnf("When writing to [%s]: database %q creation failed: %v",
+					c.config.URL, dbrp.Database, err)
+			}
+		}
+
+		err := c.writeBatch(ctx, dbrp.Database, dbrp.RetentionPolicy, batch)
 		if err != nil {
 			return err
 		}
-	} else {
-		for _, metric := range metrics {
-			db, ok := metric.GetTag(c.config.DatabaseTag)
-			if !ok {
-				db = c.config.Database
-			}
-
-			if _, ok := batches[db]; !ok {
-				batches[db] = make([]telegraf.Metric, 0)
-			}
-
-			if c.config.ExcludeDatabaseTag {
-				// Avoid modifying the metric in case we need to retry the request.
-				metric = metric.Copy()
-				metric.Accept()
-				metric.RemoveTag(c.config.DatabaseTag)
-			}
-
-			batches[db] = append(batches[db], metric)
-		}
-
-		for db, batch := range batches {
-			if !c.config.SkipDatabaseCreation && !c.createdDatabases[db] {
-				err := c.CreateDatabase(ctx, db)
-				if err != nil {
-					c.log.Warnf("When writing to [%s]: database %q creation failed: %v",
-						c.config.URL, db, err)
-				}
-			}
-
-			err := c.writeBatch(ctx, db, batch)
-			if err != nil {
-				return err
-			}
-		}
 	}
 	return nil
 }
 
-func (c *httpClient) writeBatch(ctx context.Context, db string, metrics []telegraf.Metric) error {
+func (c *httpClient) writeBatch(ctx context.Context, db, rp string, metrics []telegraf.Metric) error {
-	url, err := makeWriteURL(c.config.URL, db, c.config.RetentionPolicy, c.config.Consistency)
+	url, err := makeWriteURL(c.config.URL, db, rp, c.config.Consistency)
 	if err != nil {
 		return err
 	}

plugins/outputs/influxdb/http_test.go

@@ -733,3 +733,200 @@ func TestHTTP_WriteDatabaseTagWorksOnRetry(t *testing.T) {
 	err = client.Write(ctx, metrics)
 	require.NoError(t, err)
 }
+
+func TestDBRPTags(t *testing.T) {
+	ts := httptest.NewServer(http.NotFoundHandler())
+	defer ts.Close()
+
+	u, err := url.Parse(fmt.Sprintf("http://%s", ts.Listener.Addr().String()))
+	require.NoError(t, err)
+
+	tests := []struct {
+		name        string
+		config      influxdb.HTTPConfig
+		metrics     []telegraf.Metric
+		handlerFunc func(t *testing.T, w http.ResponseWriter, r *http.Request)
+		url         string
+	}{
+		{
+			name: "defaults",
+			config: influxdb.HTTPConfig{
+				URL:      u,
+				Database: "telegraf",
+			},
+			metrics: []telegraf.Metric{
+				testutil.MustMetric(
+					"cpu",
+					map[string]string{
+						"database": "foo",
+					},
+					map[string]interface{}{
+						"value": 42.0,
+					},
+					time.Unix(0, 0),
+				),
+			},
+			handlerFunc: func(t *testing.T, w http.ResponseWriter, r *http.Request) {
+				require.Equal(t, r.FormValue("db"), "telegraf")
+				require.Equal(t, r.FormValue("rp"), "")
+				w.WriteHeader(http.StatusNoContent)
+			},
+		},
+		{
+			name: "static retention policy",
+			config: influxdb.HTTPConfig{
+				URL:             u,
+				Database:        "telegraf",
+				RetentionPolicy: "foo",
+			},
+			metrics: []telegraf.Metric{
+				testutil.MustMetric(
+					"cpu",
+					map[string]string{},
+					map[string]interface{}{
+						"value": 42.0,
+					},
+					time.Unix(0, 0),
+				),
+			},
+			handlerFunc: func(t *testing.T, w http.ResponseWriter, r *http.Request) {
+				require.Equal(t, r.FormValue("db"), "telegraf")
+				require.Equal(t, r.FormValue("rp"), "foo")
+				w.WriteHeader(http.StatusNoContent)
+			},
+		},
+		{
+			name: "retention policy tag",
+			config: influxdb.HTTPConfig{
+				URL:                  u,
+				SkipDatabaseCreation: true,
+				Database:             "telegraf",
+				RetentionPolicyTag:   "rp",
+				Log:                  testutil.Logger{},
+			},
+			metrics: []telegraf.Metric{
+				testutil.MustMetric(
+					"cpu",
+					map[string]string{
+						"rp": "foo",
+					},
+					map[string]interface{}{
+						"value": 42.0,
+					},
+					time.Unix(0, 0),
+				),
+			},
+			handlerFunc: func(t *testing.T, w http.ResponseWriter, r *http.Request) {
+				require.Equal(t, r.FormValue("db"), "telegraf")
+				require.Equal(t, r.FormValue("rp"), "foo")
+				body, err := ioutil.ReadAll(r.Body)
+				require.NoError(t, err)
+				require.Contains(t, string(body), "cpu,rp=foo value=42")
+				w.WriteHeader(http.StatusNoContent)
+			},
+		},
+		{
+			name: "retention policy tag fallback to static rp",
+			config: influxdb.HTTPConfig{
+				URL:                  u,
+				SkipDatabaseCreation: true,
+				Database:             "telegraf",
+				RetentionPolicy:      "foo",
+				RetentionPolicyTag:   "rp",
+				Log:                  testutil.Logger{},
+			},
+			metrics: []telegraf.Metric{
+				testutil.MustMetric(
+					"cpu",
+					map[string]string{},
+					map[string]interface{}{
+						"value": 42.0,
+					},
+					time.Unix(0, 0),
+				),
+			},
+			handlerFunc: func(t *testing.T, w http.ResponseWriter, r *http.Request) {
+				require.Equal(t, r.FormValue("db"), "telegraf")
+				require.Equal(t, r.FormValue("rp"), "foo")
+				w.WriteHeader(http.StatusNoContent)
+			},
+		},
+		{
+			name: "retention policy tag fallback to unset rp",
+			config: influxdb.HTTPConfig{
+				URL:                  u,
+				SkipDatabaseCreation: true,
+				Database:             "telegraf",
+				RetentionPolicyTag:   "rp",
+				Log:                  testutil.Logger{},
+			},
+			metrics: []telegraf.Metric{
+				testutil.MustMetric(
+					"cpu",
+					map[string]string{},
+					map[string]interface{}{
+						"value": 42.0,
+					},
+					time.Unix(0, 0),
+				),
+			},
+			handlerFunc: func(t *testing.T, w http.ResponseWriter, r *http.Request) {
+				require.Equal(t, r.FormValue("db"), "telegraf")
+				require.Equal(t, r.FormValue("rp"), "")
+				w.WriteHeader(http.StatusNoContent)
+			},
+		},
+		{
+			name: "exclude retention policy tag",
+			config: influxdb.HTTPConfig{
+				URL:                       u,
+				SkipDatabaseCreation:      true,
+				Database:                  "telegraf",
+				RetentionPolicyTag:        "rp",
+				ExcludeRetentionPolicyTag: true,
+				Log:                       testutil.Logger{},
+			},
+			metrics: []telegraf.Metric{
+				testutil.MustMetric(
+					"cpu",
+					map[string]string{
+						"rp": "foo",
+					},
+					map[string]interface{}{
+						"value": 42.0,
+					},
+					time.Unix(0, 0),
+				),
+			},
+			handlerFunc: func(t *testing.T, w http.ResponseWriter, r *http.Request) {
+				require.Equal(t, r.FormValue("db"), "telegraf")
+				require.Equal(t, r.FormValue("rp"), "foo")
+				body, err := ioutil.ReadAll(r.Body)
+				require.NoError(t, err)
+				require.Contains(t, string(body), "cpu value=42")
+				w.WriteHeader(http.StatusNoContent)
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ts.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				switch r.URL.Path {
+				case "/write":
+					tt.handlerFunc(t, w, r)
+					return
+				default:
+					w.WriteHeader(http.StatusNotFound)
+					return
+				}
+			})
+
+			client, err := influxdb.NewHTTPClient(tt.config)
+			require.NoError(t, err)
+
+			ctx := context.Background()
+			err = client.Write(ctx, tt.metrics)
+			require.NoError(t, err)
+		})
+	}
+}

plugins/outputs/influxdb/influxdb.go

@@ -31,23 +31,25 @@ type Client interface {
 
 // InfluxDB struct is the primary data structure for the plugin
 type InfluxDB struct {
-	URL                  string   // url deprecated in 0.1.9; use urls
+	URL                       string            // url deprecated in 0.1.9; use urls
-	URLs                 []string `toml:"urls"`
+	URLs                      []string          `toml:"urls"`
-	Username             string
+	Username                  string            `toml:"username"`
-	Password             string
+	Password                  string            `toml:"password"`
-	Database             string
+	Database                  string            `toml:"database"`
-	DatabaseTag          string `toml:"database_tag"`
+	DatabaseTag               string            `toml:"database_tag"`
-	ExcludeDatabaseTag   bool   `toml:"exclude_database_tag"`
+	ExcludeDatabaseTag        bool              `toml:"exclude_database_tag"`
-	UserAgent            string
+	RetentionPolicy           string            `toml:"retention_policy"`
-	RetentionPolicy      string
+	RetentionPolicyTag        string            `toml:"retention_policy_tag"`
-	WriteConsistency     string
+	ExcludeRetentionPolicyTag bool              `toml:"exclude_retention_policy_tag"`
-	Timeout              internal.Duration
+	UserAgent                 string            `toml:"user_agent"`
-	UDPPayload           internal.Size     `toml:"udp_payload"`
+	WriteConsistency          string            `toml:"write_consistency"`
-	HTTPProxy            string            `toml:"http_proxy"`
+	Timeout                   internal.Duration `toml:"timeout"`
-	HTTPHeaders          map[string]string `toml:"http_headers"`
+	UDPPayload                internal.Size     `toml:"udp_payload"`
-	ContentEncoding      string            `toml:"content_encoding"`
+	HTTPProxy                 string            `toml:"http_proxy"`
-	SkipDatabaseCreation bool              `toml:"skip_database_creation"`
+	HTTPHeaders               map[string]string `toml:"http_headers"`
-	InfluxUintSupport    bool              `toml:"influx_uint_support"`
+	ContentEncoding           string            `toml:"content_encoding"`
+	SkipDatabaseCreation      bool              `toml:"skip_database_creation"`
+	InfluxUintSupport         bool              `toml:"influx_uint_support"`
 	tls.ClientConfig
 
 	Precision string // precision deprecated in 1.0; value is ignored
@@ -89,6 +91,13 @@ var sampleConfig = `
   ## the default retention policy.  Only takes effect when using HTTP.
   # retention_policy = ""
 
+  ## The value of this tag will be used to determine the retention policy.  If this
+  ## tag is not set the 'retention_policy' option is used as the default.
+  # retention_policy_tag = ""
+
+  ## If true, the 'retention_policy_tag' will not be removed from the metric.
+  # exclude_retention_policy_tag = false
+
   ## Write consistency (clusters only), can be: "any", "one", "quorum", "all".
   ## Only takes effect when using HTTP.
   # write_consistency = "any"
@@ -250,23 +259,25 @@ func (i *InfluxDB) httpClient(ctx context.Context, url *url.URL, proxy *url.URL)
 	}
 
 	config := &HTTPConfig{
-		URL:                  url,
+		URL:                       url,
-		Timeout:              i.Timeout.Duration,
+		Timeout:                   i.Timeout.Duration,
-		TLSConfig:            tlsConfig,
+		TLSConfig:                 tlsConfig,
-		UserAgent:            i.UserAgent,
+		UserAgent:                 i.UserAgent,
-		Username:             i.Username,
+		Username:                  i.Username,
-		Password:             i.Password,
+		Password:                  i.Password,
-		Proxy:                proxy,
+		Proxy:                     proxy,
-		ContentEncoding:      i.ContentEncoding,
+		ContentEncoding:           i.ContentEncoding,
-		Headers:              i.HTTPHeaders,
+		Headers:                   i.HTTPHeaders,
-		Database:             i.Database,
+		Database:                  i.Database,
-		DatabaseTag:          i.DatabaseTag,
+		DatabaseTag:               i.DatabaseTag,
-		ExcludeDatabaseTag:   i.ExcludeDatabaseTag,
+		ExcludeDatabaseTag:        i.ExcludeDatabaseTag,
-		SkipDatabaseCreation: i.SkipDatabaseCreation,
+		SkipDatabaseCreation:      i.SkipDatabaseCreation,
-		RetentionPolicy:      i.RetentionPolicy,
+		RetentionPolicy:           i.RetentionPolicy,
-		Consistency:          i.WriteConsistency,
+		RetentionPolicyTag:        i.RetentionPolicyTag,
-		Serializer:           i.newSerializer(),
+		ExcludeRetentionPolicyTag: i.ExcludeRetentionPolicyTag,
-		Log:                  i.Log,
+		Consistency:               i.WriteConsistency,
+		Serializer:                i.newSerializer(),
+		Log:                       i.Log,
 	}
 
 	c, err := i.CreateHTTPClientF(config)