Add connection string builder

connection_string.go

@@ -0,0 +1,81 @@
+package main
+
+import (
+	"errors"
+	"fmt"
+	"os/user"
+	"strings"
+)
+
+func formatConnectionUrl(opts Options) (string, error) {
+	url := opts.Url
+
+	// Make sure to only accept urls in a standard format
+	if !strings.Contains(url, "postgres://") {
+		return "", errors.New("Invalid URL. Valid format: postgres://user:password@host:port/db?sslmode=mode")
+	}
+
+	// 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)
+			}
+		}
+	}
+
+	// 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)
+	}
+
+	return url, nil
+}
+
+func buildConnectionString(opts Options) (string, error) {
+	if opts.Url != "" {
+		return formatConnectionUrl(opts)
+	}
+
+	// Try to detect user from current OS user
+	// TODO: remove os/user dependency
+	if opts.User == "" {
+		user, err := user.Current()
+
+		if err == nil {
+			opts.User = user.Username
+		}
+	}
+
+	// Disable ssl for localhost connections, most users have it disabled
+	if opts.Host == "localhost" || opts.Host == "127.0.0.1" {
+		if opts.Ssl == "" {
+			opts.Ssl = "disable"
+		}
+	}
+
+	url := "postgres://"
+
+	if opts.User != "" {
+		url += opts.User
+	}
+
+	if opts.Pass != "" {
+		url += fmt.Sprintf(":%s", opts.Pass)
+	}
+
+	url += fmt.Sprintf("@%s:%d", opts.Host, opts.Port)
+
+	if opts.DbName != "" {
+		url += fmt.Sprintf("/%s", opts.DbName)
+	}
+
+	if opts.Ssl != "" {
+		url += fmt.Sprintf("?sslmode=%s", opts.Ssl)
+	}
+
+	return url, nil
+}

connection_string_test.go

@@ -0,0 +1,158 @@
+package main
+
+import (
+	"fmt"
+	"os/user"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_Invalid_Url(t *testing.T) {
+	opts := Options{}
+	examples := []string{
+		"postgresql://foobar",
+		"foobar",
+	}
+
+	for _, val := range examples {
+		opts.Url = val
+		str, err := buildConnectionString(opts)
+
+		assert.Equal(t, "", str)
+		assert.Error(t, err)
+		assert.Equal(t, "Invalid URL. Valid format: postgres://user:password@host:port/db?sslmode=mode", err.Error())
+	}
+}
+
+func Test_Valid_Url(t *testing.T) {
+	url := "postgres://myhost/database"
+	str, err := buildConnectionString(Options{Url: url})
+
+	assert.Equal(t, nil, err)
+	assert.Equal(t, url, str)
+}
+
+func Test_Url_And_Ssl_Flag(t *testing.T) {
+	str, err := buildConnectionString(Options{
+		Url: "postgres://myhost/database",
+		Ssl: "disable",
+	})
+
+	assert.Equal(t, nil, err)
+	assert.Equal(t, "postgres://myhost/database?sslmode=disable", str)
+}
+
+func Test_Localhost_Url_And_No_Ssl_Flag(t *testing.T) {
+	str, err := buildConnectionString(Options{
+		Url: "postgres://localhost/database",
+	})
+
+	assert.Equal(t, nil, err)
+	assert.Equal(t, "postgres://localhost/database?sslmode=disable", str)
+
+	str, err = buildConnectionString(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 := buildConnectionString(Options{
+		Url: "postgres://localhost/database",
+		Ssl: "require",
+	})
+
+	assert.Equal(t, nil, err)
+	assert.Equal(t, "postgres://localhost/database?sslmode=require", str)
+
+	str, err = buildConnectionString(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 := buildConnectionString(Options{
+		Url: "postgres://localhost/database?sslmode=require",
+	})
+
+	assert.Equal(t, nil, err)
+	assert.Equal(t, "postgres://localhost/database?sslmode=require", str)
+
+	str, err = buildConnectionString(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 := buildConnectionString(Options{
+		Host:   "host",
+		Port:   5432,
+		User:   "user",
+		Pass:   "password",
+		DbName: "db",
+	})
+
+	assert.Equal(t, nil, err)
+	assert.Equal(t, "postgres://user:password@host:5432/db", str)
+}
+
+func Test_Localhost(t *testing.T) {
+	opts := Options{
+		Host:   "localhost",
+		Port:   5432,
+		User:   "user",
+		Pass:   "password",
+		DbName: "db",
+	}
+
+	str, err := buildConnectionString(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 = buildConnectionString(opts)
+	assert.Equal(t, nil, err)
+	assert.Equal(t, "postgres://user:password@127.0.0.1:5432/db?sslmode=disable", str)
+}
+
+func Test_Localhost_And_Ssl(t *testing.T) {
+	opts := Options{
+		Host:   "localhost",
+		Port:   5432,
+		User:   "user",
+		Pass:   "password",
+		DbName: "db",
+		Ssl:    "require",
+	}
+
+	str, err := buildConnectionString(opts)
+	assert.Equal(t, nil, err)
+	assert.Equal(t, "postgres://user:password@localhost:5432/db?sslmode=require", str)
+}
+
+func Test_No_User(t *testing.T) {
+	opts := Options{Host: "host", Port: 5432, DbName: "db"}
+	u, _ := user.Current()
+	str, err := buildConnectionString(opts)
+
+	assert.Equal(t, nil, err)
+	assert.Equal(t, fmt.Sprintf("postgres://%s@host:5432/db", u.Username), str)
+}
+
+func Test_Port(t *testing.T) {
+	opts := Options{Host: "host", User: "user", Port: 5000, DbName: "db"}
+	str, err := buildConnectionString(opts)
+
+	assert.Equal(t, nil, err)
+	assert.Equal(t, "postgres://user@host:5000/db", str)
+}