Merge pull request #684 from sysadmind/probe-dsn

Update multi-target handler to use new DSN type

Author: sysadmind

cmd/postgres_exporter/datasource.go

@@ -20,7 +20,6 @@ import (
 	"os"
 	"regexp"
 	"strings"
-	"unicode"
 
 	"github.com/go-kit/log/level"
 	"github.com/prometheus/client_golang/prometheus"
@@ -173,196 +172,3 @@ func getDataSources() ([]string, error) {
 
 	return []string{dsn}, nil
 }
-
-// dsn represents a parsed datasource. It contains fields for the individual connection components.
-type dsn struct {
-	scheme   string
-	username string
-	password string
-	host     string
-	path     string
-	query    string
-}
-
-// String makes a dsn safe to print by excluding any passwords. This allows dsn to be used in
-// strings and log messages without needing to call a redaction function first.
-func (d dsn) String() string {
-	if d.password != "" {
-		return fmt.Sprintf("%s://%s:******@%s%s?%s", d.scheme, d.username, d.host, d.path, d.query)
-	}
-
-	if d.username != "" {
-		return fmt.Sprintf("%s://%s@%s%s?%s", d.scheme, d.username, d.host, d.path, d.query)
-	}
-
-	return fmt.Sprintf("%s://%s%s?%s", d.scheme, d.host, d.path, d.query)
-}
-
-// dsnFromString parses a connection string into a dsn. It will attempt to parse the string as
-// a URL and as a set of key=value pairs. If both attempts fail, dsnFromString will return an error.
-func dsnFromString(in string) (dsn, error) {
-	if strings.HasPrefix(in, "postgresql://") {
-		return dsnFromURL(in)
-	}
-
-	// Try to parse as key=value pairs
-	d, err := dsnFromKeyValue(in)
-	if err == nil {
-		return d, nil
-	}
-
-	return dsn{}, fmt.Errorf("could not understand DSN")
-}
-
-// dsnFromURL parses the input as a URL and returns the dsn representation.
-func dsnFromURL(in string) (dsn, error) {
-	u, err := url.Parse(in)
-	if err != nil {
-		return dsn{}, err
-	}
-	pass, _ := u.User.Password()
-	user := u.User.Username()
-
-	query := u.Query()
-
-	if queryPass := query.Get("password"); queryPass != "" {
-		if pass == "" {
-			pass = queryPass
-		}
-	}
-	query.Del("password")
-
-	if queryUser := query.Get("user"); queryUser != "" {
-		if user == "" {
-			user = queryUser
-		}
-	}
-	query.Del("user")
-
-	d := dsn{
-		scheme:   u.Scheme,
-		username: user,
-		password: pass,
-		host:     u.Host,
-		path:     u.Path,
-		query:    query.Encode(),
-	}
-
-	return d, nil
-}
-
-// dsnFromKeyValue parses the input as a set of key=value pairs and returns the dsn representation.
-func dsnFromKeyValue(in string) (dsn, error) {
-	// Attempt to confirm at least one key=value pair before starting the rune parser
-	connstringRe := regexp.MustCompile(`^ *[a-zA-Z0-9]+ *= *[^= ]+`)
-	if !connstringRe.MatchString(in) {
-		return dsn{}, fmt.Errorf("input is not a key-value DSN")
-	}
-
-	// Anything other than known fields should be part of the querystring
-	query := url.Values{}
-
-	pairs, err := parseKeyValue(in)
-	if err != nil {
-		return dsn{}, fmt.Errorf("failed to parse key-value DSN: %v", err)
-	}
-
-	// Build the dsn from the key=value pairs
-	d := dsn{
-		scheme: "postgresql",
-	}
-
-	hostname := ""
-	port := ""
-
-	for k, v := range pairs {
-		switch k {
-		case "host":
-			hostname = v
-		case "port":
-			port = v
-		case "user":
-			d.username = v
-		case "password":
-			d.password = v
-		default:
-			query.Set(k, v)
-		}
-	}
-
-	if hostname == "" {
-		hostname = "localhost"
-	}
-
-	if port == "" {
-		d.host = hostname
-	} else {
-		d.host = fmt.Sprintf("%s:%s", hostname, port)
-	}
-
-	d.query = query.Encode()
-
-	return d, nil
-}
-
-// parseKeyValue is a key=value parser. It loops over each rune to split out keys and values
-// and attempting to honor quoted values. parseKeyValue will return an error if it is unable
-// to properly parse the input.
-func parseKeyValue(in string) (map[string]string, error) {
-	out := map[string]string{}
-
-	inPart := false
-	inQuote := false
-	part := []rune{}
-	key := ""
-	for _, c := range in {
-		switch {
-		case unicode.In(c, unicode.Quotation_Mark):
-			if inQuote {
-				inQuote = false
-			} else {
-				inQuote = true
-			}
-		case unicode.In(c, unicode.White_Space):
-			if inPart {
-				if inQuote {
-					part = append(part, c)
-				} else {
-					// Are we finishing a key=value?
-					if key == "" {
-						return out, fmt.Errorf("invalid input")
-					}
-					out[key] = string(part)
-					inPart = false
-					part = []rune{}
-				}
-			} else {
-				// Are we finishing a key=value?
-				if key == "" {
-					return out, fmt.Errorf("invalid input")
-				}
-				out[key] = string(part)
-				inPart = false
-				part = []rune{}
-				// Do something with the value
-			}
-		case c == '=':
-			if inPart {
-				inPart = false
-				key = string(part)
-				part = []rune{}
-			} else {
-				return out, fmt.Errorf("invalid input")
-			}
-		default:
-			inPart = true
-			part = append(part, c)
-		}
-	}
-
-	if key != "" && len(part) > 0 {
-		out[key] = string(part)
-	}
-
-	return out, nil
-}

collector/probe.go

@@ -16,11 +16,10 @@ package collector
 import (
 	"context"
 	"database/sql"
-	"fmt"
-	"strings"
 	"sync"
 
 	"github.com/go-kit/log"
+	"github.com/prometheus-community/postgres_exporter/config"
 	"github.com/prometheus/client_golang/prometheus"
 )
 
@@ -31,7 +30,7 @@ type ProbeCollector struct {
 	db         *sql.DB
 }
 
-func NewProbeCollector(logger log.Logger, registry *prometheus.Registry, dsn string) (*ProbeCollector, error) {
+func NewProbeCollector(logger log.Logger, registry *prometheus.Registry, dsn config.DSN) (*ProbeCollector, error) {
 	collectors := make(map[string]Collector)
 	initiatedCollectorsMtx.Lock()
 	defer initiatedCollectorsMtx.Unlock()
@@ -55,11 +54,7 @@ func NewProbeCollector(logger log.Logger, registry *prometheus.Registry, dsn str
 		}
 	}
 
-	if !strings.HasPrefix(dsn, "postgres://") {
-		dsn = fmt.Sprintf("postgres://%s", dsn)
-	}
-
-	db, err := sql.Open("postgres", dsn)
+	db, err := sql.Open("postgres", dsn.GetConnectionString())
 	if err != nil {
 		return nil, err
 	}

config/config.go

@@ -15,9 +15,7 @@ package config
 
 import (
 	"fmt"
-	"net/url"
 	"os"
-	"strings"
 	"sync"
 
 	"github.com/go-kit/log"
@@ -97,26 +95,26 @@ func (ch *ConfigHandler) ReloadConfig(f string, logger log.Logger) error {
 	return nil
 }
 
-func (m AuthModule) ConfigureTarget(target string) (string, error) {
-	// ip:port urls do not parse properly and that is the typical way users interact with postgres
-	t := fmt.Sprintf("exporter://%s", target)
-	u, err := url.Parse(t)
+func (m AuthModule) ConfigureTarget(target string) (DSN, error) {
+	dsn, err := dsnFromString(target)
 	if err != nil {
-		return "", err
+		return DSN{}, err
 	}
 
+	// Set the credentials from the authentication module
+	// TODO(@sysadmind): What should the order of precedence be?
 	if m.Type == "userpass" {
-		u.User = url.UserPassword(m.UserPass.Username, m.UserPass.Password)
+		if m.UserPass.Username != "" {
+			dsn.username = m.UserPass.Username
+		}
+		if m.UserPass.Password != "" {
+			dsn.password = m.UserPass.Password
+		}
 	}
 
-	query := u.Query()
 	for k, v := range m.Options {
-		query.Set(k, v)
+		dsn.query.Set(k, v)
 	}
-	u.RawQuery = query.Encode()
-
-	parsed := u.String()
-	trim := strings.TrimPrefix(parsed, "exporter://")
 
-	return trim, nil
+	return dsn, nil
 }