Output plugin to PostgreSQL (Copy, Batch insert)

plugins/outputs/all/all.go

@@ -25,6 +25,8 @@ import (
 	_ "github.com/influxdata/telegraf/plugins/outputs/nats"
 	_ "github.com/influxdata/telegraf/plugins/outputs/nsq"
 	_ "github.com/influxdata/telegraf/plugins/outputs/opentsdb"
+	_ "github.com/influxdata/telegraf/plugins/outputs/postgresql_batch"
+	_ "github.com/influxdata/telegraf/plugins/outputs/postgresql_copy"
 	_ "github.com/influxdata/telegraf/plugins/outputs/prometheus_client"
 	_ "github.com/influxdata/telegraf/plugins/outputs/riemann"
 	_ "github.com/influxdata/telegraf/plugins/outputs/riemann_legacy"

plugins/outputs/postgresql_batch/README.md

@@ -0,0 +1,13 @@
+# PostgreSQL Output Plugin
+
+This output plugin writes all metrics to PostgreSQL in batch.
+
+Obs: Currently, you should create your PostgreSQL tables first
+
+### Configuration:
+
+```toml
+# Send metrics to PostgreSQL using batch (grouped insert)
+[[outputs.postgresql_batch]]
+  address = "host=localhost user=postgres sslmode=verify-full"
+```

plugins/outputs/postgresql_batch/postgresql_batch.go

@@ -0,0 +1,131 @@
+package postgresql_batch
+
+import (
+	"database/sql"
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/jackc/pgx"
+
+	"github.com/influxdata/telegraf"
+	"github.com/influxdata/telegraf/plugins/outputs"
+)
+
+type PostgresqlBatch struct {
+	db      *sql.DB
+	Address string
+	Inserts map[string]string
+	Columns map[string][]string
+}
+
+func (p *PostgresqlBatch) Connect() error {
+	db, err := sql.Open("pgx", p.Address)
+	if err != nil {
+		return err
+	}
+	p.db = db
+
+	p.Inserts = make(map[string]string)
+	p.Columns = make(map[string][]string)
+
+	return nil
+}
+
+func (p *PostgresqlBatch) Close() error {
+	return p.db.Close()
+}
+
+func quoteIdent(name string) string {
+	return pgx.Identifier{name}.Sanitize()
+}
+
+var sampleConfig = `
+  # Send metrics to PostgreSQL using COPY
+  address = "host=localhost user=postgres sslmode=verify-full"
+`
+
+func (p *PostgresqlBatch) SampleConfig() string { return sampleConfig }
+func (p *PostgresqlBatch) Description() string  { return "Send metrics to Postgresql in batch" }
+
+func (p *PostgresqlBatch) generateInsert(tablename string, columns []string) string {
+	var quoted []string
+	for _, column := range columns {
+		quoted = append(quoted, quoteIdent(column))
+	}
+
+	return fmt.Sprintf("INSERT INTO %s(%s) VALUES ", quoteIdent(tablename), strings.Join(quoted, ","))
+}
+
+func (p *PostgresqlBatch) buildTableInsert(metric telegraf.Metric) {
+	table := metric.Name()
+	if p.Inserts[table] == "" {
+		for key := range metric.Tags() {
+			p.Columns[table] = append(p.Columns[table], key)
+		}
+		for key := range metric.Fields() {
+			p.Columns[table] = append(p.Columns[table], key)
+		}
+		p.Inserts[table] = p.generateInsert(table, append(p.Columns[table], "time"))
+	}
+}
+
+func quoted(value interface{}) interface{} {
+	switch value.(type) {
+	case string:
+		return "'" + value.(string) + "'"
+	case time.Time:
+		return quoted(value.(time.Time).Format("2006-01-02 15:04:05"))
+	default:
+		return value
+	}
+}
+
+func joinValues(values []interface{}) string {
+	strs := make([]string, len(values))
+	for i, value := range values {
+		strs[i] = fmt.Sprintf("%v", value)
+	}
+	return strings.Join(strs, ", ")
+}
+
+func buildValues(metric telegraf.Metric, columns []string) string {
+	var values []interface{}
+	mapString := metric.Tags()
+	for key, value := range metric.Fields() {
+		mapString[key] = fmt.Sprintf("%v", value)
+	}
+	for _, column := range columns {
+		values = append(values, quoted(mapString[column]))
+	}
+	values = append(values, quoted(metric.Time()))
+	return "(" + joinValues(values) + ")"
+}
+
+func (p *PostgresqlBatch) Write(metrics []telegraf.Metric) error {
+	values := make(map[string][]string)
+	for _, metric := range metrics {
+		p.buildTableInsert(metric)
+		table := metric.Name()
+		values[table] = append(values[table], buildValues(metric, p.Columns[table]))
+	}
+	for table, values := range values {
+		if len(values) == 0 {
+			continue
+		}
+		sql := p.Inserts[table] + strings.Join(values, ",")
+		_, err := p.db.Exec(sql)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func init() {
+	outputs.Add("postgresql_batch", func() telegraf.Output { return newPostgresqlBatch() })
+}
+
+func newPostgresqlBatch() *PostgresqlBatch {
+	return &PostgresqlBatch{}
+}

plugins/outputs/postgresql_batch/postgresql_batch_test.go

@@ -0,0 +1,54 @@
+package postgresql_batch
+
+import (
+	"testing"
+	"time"
+
+	"github.com/influxdata/telegraf/metric"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBuildInsert(t *testing.T) {
+	table := "cpu_usage"
+	timestamp := time.Date(2010, time.November, 10, 23, 0, 0, 0, time.UTC)
+	tags := map[string]string{"host": "address", "zone": "west"}
+	fields := map[string]interface{}{"cpu_perc": float64(0.2)}
+	m, _ := metric.New(table, tags, fields, timestamp)
+
+	p := newPostgresqlBatch()
+	p.Inserts = make(map[string]string)
+	assert.Empty(t, p.Inserts[table])
+	p.Columns = make(map[string][]string)
+	assert.Empty(t, p.Columns[table])
+
+	p.buildTableInsert(m)
+	assert.Equal(t, len(p.Columns[table]), 3)
+	assert.Contains(t, p.Columns[table], "host")
+	assert.Contains(t, p.Columns[table], "zone")
+	assert.Contains(t, p.Columns[table], "cpu_perc")
+	assert.Contains(t, p.Inserts[table], "INSERT INTO \""+table+"\"(")
+	assert.Contains(t, p.Inserts[table], "\"host\"")
+	assert.Contains(t, p.Inserts[table], "\"zone\"")
+	assert.Contains(t, p.Inserts[table], "\"cpu_perc\"")
+	assert.Contains(t, p.Inserts[table], ",\"time\") VALUES ")
+}
+
+func TestBuildValues(t *testing.T) {
+	table := "cpu_usage"
+	timestamp := time.Date(2010, time.November, 10, 23, 0, 0, 0, time.UTC)
+	tags := map[string]string{"host": "address", "zone": "west"}
+	fields := map[string]interface{}{"cpu_perc": float64(0.2)}
+	m, _ := metric.New(table, tags, fields, timestamp)
+
+	p := newPostgresqlBatch()
+	p.Inserts = make(map[string]string)
+	p.Columns = make(map[string][]string)
+
+	p.buildTableInsert(m)
+	values := buildValues(m, p.Columns[table])
+	assert.Contains(t, values, "'address'")
+	assert.Contains(t, values, "'west'")
+	assert.Contains(t, values, "'0.2'")
+	assert.Contains(t, values, "'2010-11-10")
+}

plugins/outputs/postgresql_copy/README.md

@@ -0,0 +1,14 @@
+# PostgreSQL Output Plugin
+
+This output plugin writes all metrics to PostgreSQL using CopyIn.
+
+Obs: Currently, you should create your PostgreSQL tables first
+
+### Configuration:
+
+```toml
+# Send metrics to PostgreSQL using CopyIn
+[[outputs.postgresql_copy]]
+  address = "postgres://USER:PWD@HOST:PORT/DATABASE?sslmode=disable"
+  ignore_insert_errors = false
+```

plugins/outputs/postgresql_copy/postgresql_copy.go

@@ -0,0 +1,143 @@
+package postgresql_copy
+
+import (
+	"database/sql"
+	"log"
+
+	"github.com/influxdata/telegraf"
+	"github.com/influxdata/telegraf/plugins/outputs"
+	"github.com/lib/pq"
+)
+
+type PostgresqlCopy struct {
+	db                 *sql.DB
+	Address            string
+	Columns            map[string][]string
+	IgnoreInsertErrors bool
+}
+
+func (p *PostgresqlCopy) Connect() error {
+	db, err := sql.Open("postgres", p.Address)
+	if err != nil {
+		return err
+	}
+	p.db = db
+	return nil
+}
+
+func (p *PostgresqlCopy) Close() error {
+	return p.db.Close()
+}
+
+var sampleConfig = `
+  # Send metrics to PostgreSQL using COPY
+  address = "postgres://USER:PWD@HOST:PORT/DATABASE?sslmode=disable"
+`
+
+func (p *PostgresqlCopy) SampleConfig() string { return sampleConfig }
+func (p *PostgresqlCopy) Description() string  { return "Send metrics to Postgres using Copy" }
+
+func (p *PostgresqlCopy) buildColumns(metrics []telegraf.Metric) {
+	table_columns := make(map[string]map[string]bool)
+	for _, metric := range metrics {
+		table := metric.Name()
+		if table_columns[table] == nil {
+			table_columns[table] = map[string]bool{}
+		}
+		for key := range metric.Fields() {
+			table_columns[table][key] = true
+		}
+		for key := range metric.Tags() {
+			table_columns[table][key] = true
+		}
+	}
+
+	p.Columns = make(map[string][]string)
+	for table, columns := range table_columns {
+		for column := range columns {
+			p.Columns[table] = append(p.Columns[table], column)
+		}
+	}
+}
+
+func buildValues(metric telegraf.Metric, columns []string) []interface{} {
+	var values []interface{}
+	all_metric := metric.Fields()
+	for key, value := range metric.Tags() {
+		all_metric[key] = value
+	}
+	for _, column := range columns {
+		values = append(values, all_metric[column])
+	}
+	values = append(values, metric.Time())
+	return values
+}
+
+func (p *PostgresqlCopy) Write(metrics []telegraf.Metric) error {
+	p.buildColumns(metrics)
+	tables := make(map[string][][]interface{})
+	for _, metric := range metrics {
+		table := metric.Name()
+		tables[table] = append(tables[table], buildValues(metric, p.Columns[table]))
+	}
+
+	txn, err := p.db.Begin()
+	if err != nil {
+		return err
+	}
+	for table, values := range tables {
+		if len(values) == 0 {
+			continue
+		}
+
+		columns := append(p.Columns[table], "time")
+		stmt, err := txn.Prepare(pq.CopyIn(table, columns...))
+		if err != nil {
+			if p.IgnoreInsertErrors {
+				log.Printf("E! Error in stmt execute %s", err)
+				txn, err = p.db.Begin()
+				continue
+			} else {
+				return err
+			}
+		}
+
+		for _, value := range values {
+			_, err = stmt.Exec(value...)
+			if err == nil {
+				continue
+			}
+
+			if p.IgnoreInsertErrors {
+				log.Printf("E! Could not insert into %s: %s", table, values)
+				continue
+			}
+			return err
+		}
+		_, err = stmt.Exec()
+		if err == nil {
+			continue
+		}
+		if p.IgnoreInsertErrors {
+			log.Printf("E! Error in stmt execute %s", err)
+			continue
+		}
+	}
+	err = txn.Commit()
+	if err == nil {
+		return nil
+	}
+	if p.IgnoreInsertErrors {
+		log.Printf("E! Error in commit %s", err)
+		return nil
+	}
+	return err
+}
+
+func init() {
+	outputs.Add("postgresql_copy", func() telegraf.Output { return newPostgresqlCopy() })
+}
+
+func newPostgresqlCopy() *PostgresqlCopy {
+	return &PostgresqlCopy{}
+}