Fix influxdb output database quoting (influxdata#2851)

nevins-b · May 25, 2017 · f47924f · f47924f
1 parent a96f85c
commit f47924f
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 3 deletions.
diff --git a/plugins/outputs/influxdb/client/http.go b/plugins/outputs/influxdb/client/http.go
@@ -16,7 +16,6 @@ var (
 	defaultRequestTimeout = time.Second * 5
 )
 
-//
 func NewHTTP(config HTTPConfig, defaultWP WriteParams) (Client, error) {
 	// validate required parameters:
 	if len(config.URL) == 0 {

diff --git a/plugins/outputs/influxdb/influxdb.go b/plugins/outputs/influxdb/influxdb.go
@@ -16,6 +16,11 @@ import (
 	"github.com/influxdata/telegraf/plugins/outputs/influxdb/client"
 )
 
+var (
+	// Quote Ident replacer.
+	qiReplacer = strings.NewReplacer("\n", `\n`, `\`, `\\`, `"`, `\"`)
+)
+
 // InfluxDB struct is the primary data structure for the plugin
 type InfluxDB struct {
 	// URL is only for backwards compatability
@@ -133,7 +138,7 @@ func (i *InfluxDB) Connect() error {
 			}
 			i.clients = append(i.clients, c)
 
-			err = c.Query("CREATE DATABASE " + i.Database)
+			err = c.Query(fmt.Sprintf(`CREATE DATABASE "%s"`, qiReplacer.Replace(i.Database)))
 			if err != nil {
 				if !strings.Contains(err.Error(), "Status Code [403]") {
 					log.Println("I! Database creation failed: " + err.Error())
@@ -191,7 +196,8 @@ func (i *InfluxDB) Write(metrics []telegraf.Metric) error {
 		if _, e := i.clients[n].WriteStream(r, bufsize); e != nil {
 			// If the database was not found, try to recreate it:
 			if strings.Contains(e.Error(), "database not found") {
-				if errc := i.clients[n].Query("CREATE DATABASE  " + i.Database); errc != nil {
+				errc := i.clients[n].Query(fmt.Sprintf(`CREATE DATABASE "%s"`, qiReplacer.Replace(i.Database)))
+				if errc != nil {
 					log.Printf("E! Error: Database %s not found and failed to recreate\n",
 						i.Database)
 				}

diff --git a/plugins/outputs/influxdb/influxdb_test.go b/plugins/outputs/influxdb/influxdb_test.go
@@ -8,9 +8,44 @@ import (
 
 	"github.com/influxdata/telegraf/testutil"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+func TestIdentQuoting(t *testing.T) {
+	var testCases = []struct {
+		database string
+		expected string
+	}{
+		{"x-y", `CREATE DATABASE "x-y"`},
+		{`x"y`, `CREATE DATABASE "x\"y"`},
+		{"x\ny", `CREATE DATABASE "x\ny"`},
+		{`x\y`, `CREATE DATABASE "x\\y"`},
+	}
+
+	for _, tc := range testCases {
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			r.ParseForm()
+			q := r.Form.Get("q")
+			assert.Equal(t, tc.expected, q)
+
+			w.WriteHeader(http.StatusOK)
+			w.Header().Set("Content-Type", "application/json")
+			fmt.Fprintln(w, `{"results":[{}]}`)
+		}))
+		defer ts.Close()
+
+		i := InfluxDB{
+			URLs:     []string{ts.URL},
+			Database: tc.database,
+		}
+
+		err := i.Connect()
+		require.NoError(t, err)
+		require.NoError(t, i.Close())
+	}
+}
+
 func TestUDPInflux(t *testing.T) {
 	i := InfluxDB{
 		URLs: []string{"udp://localhost:8089"},