From 3f4375be063ca05e441025d561513b1c53f5144c Mon Sep 17 00:00:00 2001
From: Boris Popovschi <zyqsempai@mail.ru>
Date: Thu, 12 Sep 2019 23:38:35 +0300
Subject: [PATCH] Add ability to read query from file to postgresql_extensible
 input (#6361)

---
 .../inputs/postgresql_extensible/README.md    |  5 ++-
 .../postgresql_extensible.go                  | 36 ++++++++++++++++++-
 .../postgresql_extensible_test.go             | 27 +++++++++++++-
 .../postgresql_extensible/testdata/test.sql   |  1 +
 4 files changed, 66 insertions(+), 3 deletions(-)
 create mode 100644 plugins/inputs/postgresql_extensible/testdata/test.sql

diff --git a/plugins/inputs/postgresql_extensible/README.md b/plugins/inputs/postgresql_extensible/README.md
index 29c5e36d853e4..4621e46b54078 100644
--- a/plugins/inputs/postgresql_extensible/README.md
+++ b/plugins/inputs/postgresql_extensible/README.md
@@ -44,6 +44,9 @@ The example below has two queries are specified, with the following parameters:
   # Be careful that if the withdbname is set to false you don't have to define
   # the where clause (aka with the dbname)
   #
+  # The script option can be used to specify the .sql file path.
+  # If script and sqlquery options specified at same time, sqlquery will be used 
+  #
   # the tagvalue field is used to define custom tags (separated by comas).
   # the query is expected to return columns which match the names of the
   # defined tags. The values in these columns must be of a string-type,
@@ -61,7 +64,7 @@ The example below has two queries are specified, with the following parameters:
     withdbname=false
     tagvalue=""
   [[inputs.postgresql_extensible.query]]
-    sqlquery="SELECT * FROM pg_stat_bgwriter"
+    script="your_sql-filepath.sql"
     version=901
     withdbname=false
     tagvalue=""
diff --git a/plugins/inputs/postgresql_extensible/postgresql_extensible.go b/plugins/inputs/postgresql_extensible/postgresql_extensible.go
index c2bcb7b600efc..05e57583f83b3 100644
--- a/plugins/inputs/postgresql_extensible/postgresql_extensible.go
+++ b/plugins/inputs/postgresql_extensible/postgresql_extensible.go
@@ -3,7 +3,9 @@ package postgresql_extensible
 import (
 	"bytes"
 	"fmt"
+	"io/ioutil"
 	"log"
+	"os"
 	"strings"
 
 	// register in driver.
@@ -25,6 +27,7 @@ type Postgresql struct {
 
 type query []struct {
 	Sqlquery    string
+	Script      string
 	Version     int
 	Withdbname  bool
 	Tagvalue    string
@@ -75,7 +78,10 @@ var sampleConfig = `
   ## field is used to define custom tags (separated by commas)
   ## The optional "measurement" value can be used to override the default
   ## output measurement name ("postgresql").
-  #
+  ##
+  ## The script option can be used to specify the .sql file path.
+  ## If script and sqlquery options specified at same time, sqlquery will be used 
+  ##
   ## Structure :
   ## [[inputs.postgresql_extensible.query]]
   ##   sqlquery string
@@ -96,6 +102,19 @@ var sampleConfig = `
     tagvalue="postgresql.stats"
 `
 
+func (p *Postgresql) Init() error {
+	var err error
+	for i := range p.Query {
+		if p.Query[i].Sqlquery == "" {
+			p.Query[i].Sqlquery, err = ReadQueryFromFile(p.Query[i].Script)
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
 func (p *Postgresql) SampleConfig() string {
 	return sampleConfig
 }
@@ -108,6 +127,20 @@ func (p *Postgresql) IgnoredColumns() map[string]bool {
 	return ignoredColumns
 }
 
+func ReadQueryFromFile(filePath string) (string, error) {
+	file, err := os.Open(filePath)
+	if err != nil {
+		return "", err
+	}
+	defer file.Close()
+
+	query, err := ioutil.ReadAll(file)
+	if err != nil {
+		return "", err
+	}
+	return string(query), err
+}
+
 func (p *Postgresql) Gather(acc telegraf.Accumulator) error {
 	var (
 		err         error
@@ -131,6 +164,7 @@ func (p *Postgresql) Gather(acc telegraf.Accumulator) error {
 	for i := range p.Query {
 		sql_query = p.Query[i].Sqlquery
 		tag_value = p.Query[i].Tagvalue
+
 		if p.Query[i].Measurement != "" {
 			meas_name = p.Query[i].Measurement
 		} else {
diff --git a/plugins/inputs/postgresql_extensible/postgresql_extensible_test.go b/plugins/inputs/postgresql_extensible/postgresql_extensible_test.go
index 1ed62a1cd62c0..7fbc343028015 100644
--- a/plugins/inputs/postgresql_extensible/postgresql_extensible_test.go
+++ b/plugins/inputs/postgresql_extensible/postgresql_extensible_test.go
@@ -25,7 +25,7 @@ func queryRunner(t *testing.T, q query) *testutil.Accumulator {
 	}
 	var acc testutil.Accumulator
 	p.Start(&acc)
-
+	p.Init()
 	require.NoError(t, acc.GatherError(p.Gather))
 	return &acc
 }
@@ -201,6 +201,31 @@ func TestPostgresqlFieldOutput(t *testing.T) {
 	}
 }
 
+func TestPostgresqlSqlScript(t *testing.T) {
+	q := query{{
+		Script:     "testdata/test.sql",
+		Version:    901,
+		Withdbname: false,
+		Tagvalue:   "",
+	}}
+	p := &Postgresql{
+		Service: postgresql.Service{
+			Address: fmt.Sprintf(
+				"host=%s user=postgres sslmode=disable",
+				testutil.GetLocalHost(),
+			),
+			IsPgBouncer: false,
+		},
+		Databases: []string{"postgres"},
+		Query:     q,
+	}
+	var acc testutil.Accumulator
+	p.Start(&acc)
+	p.Init()
+
+	require.NoError(t, acc.GatherError(p.Gather))
+}
+
 func TestPostgresqlIgnoresUnwantedColumns(t *testing.T) {
 	if testing.Short() {
 		t.Skip("Skipping integration test in short mode")
diff --git a/plugins/inputs/postgresql_extensible/testdata/test.sql b/plugins/inputs/postgresql_extensible/testdata/test.sql
new file mode 100644
index 0000000000000..49ec02b258fe9
--- /dev/null
+++ b/plugins/inputs/postgresql_extensible/testdata/test.sql
@@ -0,0 +1 @@
+select * from pg_stat_database
\ No newline at end of file