Merge pull request #378 from sosedoff/connection-string-fixes

Fix issues with connection string builder

pkg/cli/cli.go

@@ -38,7 +38,7 @@ func initClientUsingBookmark(bookmarkPath, bookmarkName string) (*client.Client,
 	if opt.Url != "" { // if the bookmark has url set, use it
 		connStr = opt.Url
 	} else {
-		connStr, err = connection.BuildString(opt)
+		connStr, err = connection.BuildStringFromOptions(opt)
 		if err != nil {
 			return nil, fmt.Errorf("error building connection string: %v", err)
 		}

pkg/client/client.go

@@ -55,7 +55,7 @@ func getSchemaAndTable(str string) (string, string) {
 }
 
 func New() (*Client, error) {
-	str, err := connection.BuildString(command.Opts)
+	str, err := connection.BuildStringFromOptions(command.Opts)
 
 	if command.Opts.Debug && str != "" {
 		fmt.Println("Creating a new client for:", str)

pkg/connection/connection_string.go

@@ -11,6 +11,10 @@ import (
 	"github.com/sosedoff/pgweb/pkg/command"
 )
 
+var (
+	formatError = errors.New("Invalid URL. Valid format: postgres://user:password@host:port/db?sslmode=mode")
+)
+
 func currentUser() (string, error) {
 	u, err := user.Current()
 	if err == nil {
@@ -25,39 +29,66 @@ func currentUser() (string, error) {
 	return "", errors.New("Unable to detect OS user")
 }
 
+// Check if connection url has a correct postgres prefix
+func hasValidPrefix(str string) bool {
+	return strings.HasPrefix(str, "postgres://") || strings.HasPrefix(str, "postgresql://")
+}
+
+// Extract all query vals and return as a map
+func valsFromQuery(vals neturl.Values) map[string]string {
+	result := map[string]string{}
+	for k, v := range vals {
+		result[strings.ToLower(k)] = v[0]
+	}
+	return result
+}
+
 func FormatUrl(opts command.Options) (string, error) {
 	url := opts.Url
 
-	// Make sure to only accept urls in a standard format
+	// Validate connection string prefix
-	if !strings.HasPrefix(url, "postgres://") && !strings.HasPrefix(url, "postgresql://") {
+	if !hasValidPrefix(url) {
-		return "", errors.New("Invalid URL. Valid format: postgres://user:password@host:port/db?sslmode=mode")
+		return "", formatError
 	}
 
-	// Special handling for local connections
+	// Validate the URL
-	if strings.Contains(url, "localhost") || strings.Contains(url, "127.0.0.1") {
+	uri, err := neturl.Parse(url)
-		if !strings.Contains(url, "?sslmode") {
+	if err != nil {
-			if opts.Ssl == "" {
+		return "", formatError
-				url += fmt.Sprintf("?sslmode=%s", "disable")
+	}
-			} else {
+
-				url += fmt.Sprintf("?sslmode=%s", opts.Ssl)
+	// Get query params
+	params := valsFromQuery(uri.Query())
+
+	// Determine if we need to specify sslmode if it's missing
+	if params["sslmode"] == "" {
+		if opts.Ssl == "" {
+			// Only modify sslmode for local connections
+			if strings.Contains(uri.Host, "localhost") || strings.Contains(uri.Host, "127.0.0.1") {
+				params["sslmode"] = "disable"
 			}
+		} else {
+			params["sslmode"] = opts.Ssl
 		}
 	}
 
-	// Append sslmode parameter only if its defined as a flag and not present
+	// Rebuild query params
-	// in the connection string.
+	query := neturl.Values{}
-	if !strings.Contains(url, "?sslmode") && opts.Ssl != "" {
+	for k, v := range params {
-		url += fmt.Sprintf("?sslmode=%s", opts.Ssl)
+		query.Add(k, v)
 	}
+	uri.RawQuery = query.Encode()
 
-	return url, nil
+	return uri.String(), nil
 }
 
 func IsBlank(opts command.Options) bool {
 	return opts.Host == "" && opts.User == "" && opts.DbName == "" && opts.Url == ""
 }
 
-func BuildString(opts command.Options) (string, error) {
+// Build a new database connection string for the CLI options
+func BuildStringFromOptions(opts command.Options) (string, error) {
+	// If connection string is provided we just use that
 	if opts.Url != "" {
 		return FormatUrl(opts)
 	}
@@ -71,31 +102,22 @@ func BuildString(opts command.Options) (string, error) {
 	}
 
 	// Disable ssl for localhost connections, most users have it disabled
-	if opts.Host == "localhost" || opts.Host == "127.0.0.1" {
+	if opts.Ssl == "" && (opts.Host == "localhost" || opts.Host == "127.0.0.1") {
-		if opts.Ssl == "" {
+		opts.Ssl = "disable"
-			opts.Ssl = "disable"
-		}
-	}
-
-	url := "postgres://"
-
-	if opts.User != "" {
-		url += opts.User
-	}
-
-	if opts.Pass != "" {
-		url += fmt.Sprintf(":%s", neturl.QueryEscape(opts.Pass))
-	}
-
-	url += fmt.Sprintf("@%s:%d", opts.Host, opts.Port)
-
-	if opts.DbName != "" {
-		url += fmt.Sprintf("/%s", opts.DbName)
 	}
 
+	query := neturl.Values{}
 	if opts.Ssl != "" {
-		url += fmt.Sprintf("?sslmode=%s", opts.Ssl)
+		query.Add("sslmode", opts.Ssl)
 	}
 
-	return url, nil
+	url := neturl.URL{
+		Scheme:   "postgres",
+		Host:     fmt.Sprintf("%v:%v", opts.Host, opts.Port),
+		User:     neturl.UserPassword(opts.User, opts.Pass),
+		Path:     fmt.Sprintf("/%s", opts.DbName),
+		RawQuery: query.Encode(),
+	}
+
+	return url.String(), nil
 }

pkg/connection/connection_string_test.go

@@ -2,6 +2,7 @@ package connection
 
 import (
 	"fmt"
+	"net/url"
 	"os/user"
 	"testing"
 
@@ -13,12 +14,13 @@ func Test_Invalid_Url(t *testing.T) {
 	opts := command.Options{}
 	examples := []string{
 		"postgre://foobar",
+		"tcp://blah",
 		"foobar",
 	}
 
 	for _, val := range examples {
 		opts.Url = val
-		str, err := BuildString(opts)
+		str, err := BuildStringFromOptions(opts)
 
 		assert.Equal(t, "", str)
 		assert.Error(t, err)
@@ -28,14 +30,14 @@ func Test_Invalid_Url(t *testing.T) {
 
 func Test_Valid_Url(t *testing.T) {
 	url := "postgres://myhost/database"
-	str, err := BuildString(command.Options{Url: url})
+	str, err := BuildStringFromOptions(command.Options{Url: url})
 
 	assert.Equal(t, nil, err)
 	assert.Equal(t, url, str)
 }
 
 func Test_Url_And_Ssl_Flag(t *testing.T) {
-	str, err := BuildString(command.Options{
+	str, err := BuildStringFromOptions(command.Options{
 		Url: "postgres://myhost/database",
 		Ssl: "disable",
 	})
@@ -45,57 +47,51 @@ func Test_Url_And_Ssl_Flag(t *testing.T) {
 }
 
 func Test_Localhost_Url_And_No_Ssl_Flag(t *testing.T) {
-	str, err := BuildString(command.Options{
+	str, err := BuildStringFromOptions(command.Options{
 		Url: "postgres://localhost/database",
 	})
-
 	assert.Equal(t, nil, err)
 	assert.Equal(t, "postgres://localhost/database?sslmode=disable", str)
 
-	str, err = BuildString(command.Options{
+	str, err = BuildStringFromOptions(command.Options{
 		Url: "postgres://127.0.0.1/database",
 	})
-
 	assert.Equal(t, nil, err)
 	assert.Equal(t, "postgres://127.0.0.1/database?sslmode=disable", str)
 }
 
 func Test_Localhost_Url_And_Ssl_Flag(t *testing.T) {
-	str, err := BuildString(command.Options{
+	str, err := BuildStringFromOptions(command.Options{
 		Url: "postgres://localhost/database",
 		Ssl: "require",
 	})
-
 	assert.Equal(t, nil, err)
 	assert.Equal(t, "postgres://localhost/database?sslmode=require", str)
 
-	str, err = BuildString(command.Options{
+	str, err = BuildStringFromOptions(command.Options{
 		Url: "postgres://127.0.0.1/database",
 		Ssl: "require",
 	})
-
 	assert.Equal(t, nil, err)
 	assert.Equal(t, "postgres://127.0.0.1/database?sslmode=require", str)
 }
 
 func Test_Localhost_Url_And_Ssl_Arg(t *testing.T) {
-	str, err := BuildString(command.Options{
+	str, err := BuildStringFromOptions(command.Options{
 		Url: "postgres://localhost/database?sslmode=require",
 	})
-
 	assert.Equal(t, nil, err)
 	assert.Equal(t, "postgres://localhost/database?sslmode=require", str)
 
-	str, err = BuildString(command.Options{
+	str, err = BuildStringFromOptions(command.Options{
 		Url: "postgres://127.0.0.1/database?sslmode=require",
 	})
-
 	assert.Equal(t, nil, err)
 	assert.Equal(t, "postgres://127.0.0.1/database?sslmode=require", str)
 }
 
 func Test_Flag_Args(t *testing.T) {
-	str, err := BuildString(command.Options{
+	str, err := BuildStringFromOptions(command.Options{
 		Host:   "host",
 		Port:   5432,
 		User:   "user",
@@ -116,12 +112,12 @@ func Test_Localhost(t *testing.T) {
 		DbName: "db",
 	}
 
-	str, err := BuildString(opts)
+	str, err := BuildStringFromOptions(opts)
 	assert.Equal(t, nil, err)
 	assert.Equal(t, "postgres://user:password@localhost:5432/db?sslmode=disable", str)
 
 	opts.Host = "127.0.0.1"
-	str, err = BuildString(opts)
+	str, err = BuildStringFromOptions(opts)
 	assert.Equal(t, nil, err)
 	assert.Equal(t, "postgres://user:password@127.0.0.1:5432/db?sslmode=disable", str)
 }
@@ -136,7 +132,7 @@ func Test_Localhost_And_Ssl(t *testing.T) {
 		Ssl:    "require",
 	}
 
-	str, err := BuildString(opts)
+	str, err := BuildStringFromOptions(opts)
 	assert.Equal(t, nil, err)
 	assert.Equal(t, "postgres://user:password@localhost:5432/db?sslmode=require", str)
 }
@@ -144,18 +140,19 @@ func Test_Localhost_And_Ssl(t *testing.T) {
 func Test_No_User(t *testing.T) {
 	opts := command.Options{Host: "host", Port: 5432, DbName: "db"}
 	u, _ := user.Current()
-	str, err := BuildString(opts)
+	str, err := BuildStringFromOptions(opts)
+	userAndPass := url.UserPassword(u.Username, "").String()
 
 	assert.Equal(t, nil, err)
-	assert.Equal(t, fmt.Sprintf("postgres://%s@host:5432/db", u.Username), str)
+	assert.Equal(t, fmt.Sprintf("postgres://%s@host:5432/db", userAndPass), str)
 }
 
 func Test_Port(t *testing.T) {
 	opts := command.Options{Host: "host", User: "user", Port: 5000, DbName: "db"}
-	str, err := BuildString(opts)
+	str, err := BuildStringFromOptions(opts)
 
 	assert.Equal(t, nil, err)
-	assert.Equal(t, "postgres://user@host:5000/db", str)
+	assert.Equal(t, "postgres://user:@host:5000/db", str)
 }
 
 func Test_Blank(t *testing.T) {