diff --git a/connection_string.go b/connection_string.go new file mode 100644 index 0000000..59e524c --- /dev/null +++ b/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 +} diff --git a/connection_string_test.go b/connection_string_test.go new file mode 100644 index 0000000..cdd3d2b --- /dev/null +++ b/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) +}