Merge pull request #684 from sysadmind/probe-dsn

Update multi-target handler to use new DSN type

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://") {
+	db, err := sql.Open("postgres", dsn.GetConnectionString())
-		dsn = fmt.Sprintf("postgres://%s", dsn)
-	}
-
-	db, err := sql.Open("postgres", dsn)
 	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) {
+func (m AuthModule) ConfigureTarget(target string) (DSN, error) {
-	// ip:port urls do not parse properly and that is the typical way users interact with postgres
+	dsn, err := dsnFromString(target)
-	t := fmt.Sprintf("exporter://%s", target)
-	u, err := url.Parse(t)
 	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()
+	return dsn, nil
-	trim := strings.TrimPrefix(parsed, "exporter://")
-
-	return trim, nil
 }

config/dsn.go

@@ -0,0 +1,238 @@
+// Copyright 2022 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package config
+
+import (
+	"fmt"
+	"net/url"
+	"regexp"
+	"strings"
+	"unicode"
+)
+
+// 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    url.Values
+}
+
+// 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.Encode())
+	}
+
+	if d.username != "" {
+		return fmt.Sprintf("%s://%s@%s%s?%s", d.scheme, d.username, d.host, d.path, d.query.Encode())
+	}
+
+	return fmt.Sprintf("%s://%s%s?%s", d.scheme, d.host, d.path, d.query.Encode())
+}
+
+// GetConnectionString returns the URL to pass to the driver for database connections. This value should not be logged.
+func (d DSN) GetConnectionString() string {
+	u := url.URL{
+		Scheme:   d.scheme,
+		Host:     d.host,
+		Path:     d.path,
+		RawQuery: d.query.Encode(),
+	}
+
+	// Username and Password
+	if d.username != "" {
+		u.User = url.UserPassword(d.username, d.password)
+	}
+
+	return u.String()
+}
+
+// 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
+	}
+
+	// Parse the string as a URL, with the scheme prefixed
+	d, err = dsnFromURL(fmt.Sprintf("postgresql://%s", 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,
+	}
+
+	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
+
+	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
+}

config/dsn_test.go

@@ -11,9 +11,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package main
+package config
 
 import (
+	"net/url"
 	"reflect"
 	"testing"
 )
@@ -28,7 +29,7 @@ func Test_dsn_String(t *testing.T) {
 		password string
 		host     string
 		path     string
-		query    string
+		query    url.Values
 	}
 	tests := []struct {
 		name   string
@@ -41,7 +42,7 @@ func Test_dsn_String(t *testing.T) {
 				scheme:   "postgresql",
 				username: "test",
 				host:     "localhost:5432",
-				query:    "",
+				query:    url.Values{},
 			},
 			want: "postgresql://test@localhost:5432?",
 		},
@@ -52,7 +53,7 @@ func Test_dsn_String(t *testing.T) {
 				username: "test",
 				password: "supersecret",
 				host:     "localhost:5432",
-				query:    "",
+				query:    url.Values{},
 			},
 			want: "postgresql://test:******@localhost:5432?",
 		},
@@ -63,7 +64,9 @@ func Test_dsn_String(t *testing.T) {
 				username: "test",
 				password: "supersecret",
 				host:     "localhost:5432",
-				query:    "ssldisable=true",
+				query: url.Values{
+					"ssldisable": []string{"true"},
+				},
 			},
 			want: "postgresql://test:******@localhost:5432?ssldisable=true",
 		},
@@ -75,14 +78,16 @@ func Test_dsn_String(t *testing.T) {
 				password: "supersecret",
 				host:     "localhost:5432",
 				path:     "/somevalue",
-				query:    "ssldisable=true",
+				query: url.Values{
+					"ssldisable": []string{"true"},
+				},
 			},
 			want: "postgresql://test:******@localhost:5432/somevalue?ssldisable=true",
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			d := dsn{
+			d := DSN{
 				scheme:   tt.fields.scheme,
 				username: tt.fields.username,
 				password: tt.fields.password,
@@ -105,61 +110,65 @@ func Test_dsnFromString(t *testing.T) {
 	tests := []struct {
 		name    string
 		input   string
-		want    dsn
+		want    DSN
 		wantErr bool
 	}{
 		{
 			name:  "Key value with password",
 			input: "host=host.example.com user=postgres port=5432 password=s3cr3t",
-			want: dsn{
+			want: DSN{
 				scheme:   "postgresql",
 				host:     "host.example.com:5432",
 				username: "postgres",
 				password: "s3cr3t",
+				query:    url.Values{},
 			},
 			wantErr: false,
 		},
 		{
 			name:  "Key value with quoted password and space",
 			input: "host=host.example.com user=postgres port=5432 password=\"s3cr 3t\"",
-			want: dsn{
+			want: DSN{
 				scheme:   "postgresql",
 				host:     "host.example.com:5432",
 				username: "postgres",
 				password: "s3cr 3t",
+				query:    url.Values{},
 			},
 			wantErr: false,
 		},
 		{
 			name:  "Key value with different order",
 			input: "password=abcde host=host.example.com user=postgres port=5432",
-			want: dsn{
+			want: DSN{
 				scheme:   "postgresql",
 				host:     "host.example.com:5432",
 				username: "postgres",
 				password: "abcde",
+				query:    url.Values{},
 			},
 			wantErr: false,
 		},
 		{
 			name:  "Key value with different order, quoted password, duplicate password",
 			input: "password=abcde host=host.example.com user=postgres port=5432 password=\"s3cr 3t\"",
-			want: dsn{
+			want: DSN{
 				scheme:   "postgresql",
 				host:     "host.example.com:5432",
 				username: "postgres",
 				password: "s3cr 3t",
+				query:    url.Values{},
 			},
 			wantErr: false,
 		},
 		{
 			name:  "URL with user in query string",
 			input: "postgresql://host.example.com:5432/tsdb?user=postgres",
-			want: dsn{
+			want: DSN{
 				scheme:   "postgresql",
 				host:     "host.example.com:5432",
 				path:     "/tsdb",
-				query:    "",
+				query:    url.Values{},
 				username: "postgres",
 			},
 			wantErr: false,
@@ -167,11 +176,11 @@ func Test_dsnFromString(t *testing.T) {
 		{
 			name:  "URL with user and password",
 			input: "postgresql://user:s3cret@host.example.com:5432/tsdb?user=postgres",
-			want: dsn{
+			want: DSN{
 				scheme:   "postgresql",
 				host:     "host.example.com:5432",
 				path:     "/tsdb",
-				query:    "",
+				query:    url.Values{},
 				username: "user",
 				password: "s3cret",
 			},
@@ -180,11 +189,11 @@ func Test_dsnFromString(t *testing.T) {
 		{
 			name:  "URL with user and password in query string",
 			input: "postgresql://host.example.com:5432/tsdb?user=postgres&password=s3cr3t",
-			want: dsn{
+			want: DSN{
 				scheme:   "postgresql",
 				host:     "host.example.com:5432",
 				path:     "/tsdb",
-				query:    "",
+				query:    url.Values{},
 				username: "postgres",
 				password: "s3cr3t",
 			},