Use persistent connection to postgresql database (#2701)

plugins/inputs/postgresql/service.go

@@ -0,0 +1,142 @@
+package postgresql
+
+import (
+	"database/sql"
+	"fmt"
+	"net"
+	"net/url"
+	"regexp"
+	"sort"
+	"strings"
+
+	"github.com/influxdata/telegraf"
+	"github.com/influxdata/telegraf/internal"
+)
+
+// pulled from lib/pq
+// ParseURL no longer needs to be used by clients of this library since supplying a URL as a
+// connection string to sql.Open() is now supported:
+//
+//	sql.Open("postgres", "postgres://bob:secret@1.2.3.4:5432/mydb?sslmode=verify-full")
+//
+// It remains exported here for backwards-compatibility.
+//
+// ParseURL converts a url to a connection string for driver.Open.
+// Example:
+//
+//	"postgres://bob:secret@1.2.3.4:5432/mydb?sslmode=verify-full"
+//
+// converts to:
+//
+//	"user=bob password=secret host=1.2.3.4 port=5432 dbname=mydb sslmode=verify-full"
+//
+// A minimal example:
+//
+//	"postgres://"
+//
+// This will be blank, causing driver.Open to use all of the defaults
+func parseURL(uri string) (string, error) {
+	u, err := url.Parse(uri)
+	if err != nil {
+		return "", err
+	}
+
+	if u.Scheme != "postgres" && u.Scheme != "postgresql" {
+		return "", fmt.Errorf("invalid connection protocol: %s", u.Scheme)
+	}
+
+	var kvs []string
+	escaper := strings.NewReplacer(` `, `\ `, `'`, `\'`, `\`, `\\`)
+	accrue := func(k, v string) {
+		if v != "" {
+			kvs = append(kvs, k+"="+escaper.Replace(v))
+		}
+	}
+
+	if u.User != nil {
+		v := u.User.Username()
+		accrue("user", v)
+
+		v, _ = u.User.Password()
+		accrue("password", v)
+	}
+
+	if host, port, err := net.SplitHostPort(u.Host); err != nil {
+		accrue("host", u.Host)
+	} else {
+		accrue("host", host)
+		accrue("port", port)
+	}
+
+	if u.Path != "" {
+		accrue("dbname", u.Path[1:])
+	}
+
+	q := u.Query()
+	for k := range q {
+		accrue(k, q.Get(k))
+	}
+
+	sort.Strings(kvs) // Makes testing easier (not a performance concern)
+	return strings.Join(kvs, " "), nil
+}
+
+// Service common functionality shared between the postgresql and postgresql_extensible
+// packages.
+type Service struct {
+	Address       string
+	Outputaddress string
+	MaxIdle       int
+	MaxOpen       int
+	MaxLifetime   internal.Duration
+	DB            *sql.DB
+}
+
+// Start starts the ServiceInput's service, whatever that may be
+func (p *Service) Start(telegraf.Accumulator) (err error) {
+	const localhost = "host=localhost sslmode=disable"
+
+	if p.Address == "" || p.Address == "localhost" {
+		p.Address = localhost
+	}
+
+	if p.DB, err = sql.Open("pgx", p.Address); err != nil {
+		return err
+	}
+
+	p.DB.SetMaxOpenConns(p.MaxOpen)
+	p.DB.SetMaxIdleConns(p.MaxIdle)
+	p.DB.SetConnMaxLifetime(p.MaxLifetime.Duration)
+
+	return nil
+}
+
+// Stop stops the services and closes any necessary channels and connections
+func (p *Service) Stop() {
+	p.DB.Close()
+}
+
+var kvMatcher, _ = regexp.Compile("(password|sslcert|sslkey|sslmode|sslrootcert)=\\S+ ?")
+
+// SanitizedAddress utility function to strip sensitive information from the connection string.
+func (p *Service) SanitizedAddress() (sanitizedAddress string, err error) {
+	var (
+		canonicalizedAddress string
+	)
+
+	if p.Outputaddress != "" {
+		return p.Outputaddress, nil
+	}
+
+	if strings.HasPrefix(p.Address, "postgres://") || strings.HasPrefix(p.Address, "postgresql://") {
+		if canonicalizedAddress, err = parseURL(p.Address); err != nil {
+			return sanitizedAddress, err
+		}
+	} else {
+		canonicalizedAddress = p.Address
+	}
+
+	sanitizedAddress = kvMatcher.ReplaceAllString(canonicalizedAddress, "")
+
+	return sanitizedAddress, err
+}