Refactor connection string generator

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,32 +29,57 @@ 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
-	if !strings.HasPrefix(url, "postgres://") && !strings.HasPrefix(url, "postgresql://") {
-		return "", errors.New("Invalid URL. Valid format: postgres://user:password@host:port/db?sslmode=mode")
+	// Validate connection string prefix
+	if !hasValidPrefix(url) {
+		return "", formatError
 	}
 
-	// Special handling for local connections
-	if strings.Contains(url, "localhost") || strings.Contains(url, "127.0.0.1") {
-		if !strings.Contains(url, "?sslmode") {
-			if opts.Ssl == "" {
-				url += fmt.Sprintf("?sslmode=%s", "disable")
-			} else {
-				url += fmt.Sprintf("?sslmode=%s", opts.Ssl)
+	// Validate the URL
+	uri, err := neturl.Parse(url)
+	if err != nil {
+		return "", formatError
+	}
+
+	// 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
-	// in the connection string.
-	if !strings.Contains(url, "?sslmode") && opts.Ssl != "" {
-		url += fmt.Sprintf("?sslmode=%s", opts.Ssl)
+	// Rebuild query params
+	query := neturl.Values{}
+	for k, v := range params {
+		query.Add(k, v)
 	}
+	uri.RawQuery = query.Encode()
 
-	return url, nil
+	return uri.String(), nil
 }
 
 func IsBlank(opts command.Options) bool {

pkg/connection/connection_string_test.go

@@ -13,6 +13,7 @@ func Test_Invalid_Url(t *testing.T) {
 	opts := command.Options{}
 	examples := []string{
 		"postgre://foobar",
+		"tcp://blah",
 		"foobar",
 	}
 
@@ -48,14 +49,12 @@ func Test_Localhost_Url_And_No_Ssl_Flag(t *testing.T) {
 	str, err := BuildString(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{
 		Url: "postgres://127.0.0.1/database",
 	})
-
 	assert.Equal(t, nil, err)
 	assert.Equal(t, "postgres://127.0.0.1/database?sslmode=disable", str)
 }
@@ -65,7 +64,6 @@ func Test_Localhost_Url_And_Ssl_Flag(t *testing.T) {
 		Url: "postgres://localhost/database",
 		Ssl: "require",
 	})
-
 	assert.Equal(t, nil, err)
 	assert.Equal(t, "postgres://localhost/database?sslmode=require", str)
 
@@ -73,7 +71,6 @@ func Test_Localhost_Url_And_Ssl_Flag(t *testing.T) {
 		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)
 }
@@ -82,14 +79,12 @@ func Test_Localhost_Url_And_Ssl_Arg(t *testing.T) {
 	str, err := BuildString(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{
 		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)
 }