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 {