Postgres: support a connection string DSN

README.md

@@ -196,6 +196,8 @@ $ dbmate -u "postgres://postgres@127.0.0.1:5432/myapp_test?sslmode=disable" up
 
 The only advantage of using `dbmate -e TEST_DATABASE_URL` over `dbmate -u $TEST_DATABASE_URL` is that the former takes advantage of dbmate's automatic `.env` file loading.
 
+Some drivers (currently just Postgres) support specifying a DSN string instead of a url by using `DATABASE_DSN` instead of `DATABBASE_URL`. When using `DATABASE_DSN` you can specify the driver with `DBMATE_DRIVER`.
+
 #### PostgreSQL
 
 When connecting to Postgres, you may need to add the `sslmode=disable` option to your connection string, as dbmate by default requires a TLS connection (some other frameworks/languages allow unencrypted connections by default).

main.go

@@ -11,6 +11,7 @@ import (
 	"github.com/urfave/cli/v2"
 
 	"github.com/amacneil/dbmate/v2/pkg/dbmate"
+	"github.com/amacneil/dbmate/v2/pkg/dbutil"
 	_ "github.com/amacneil/dbmate/v2/pkg/driver/clickhouse"
 	_ "github.com/amacneil/dbmate/v2/pkg/driver/mysql"
 	_ "github.com/amacneil/dbmate/v2/pkg/driver/postgres"
@@ -225,11 +226,11 @@ func loadDotEnv() {
 // action wraps a cli.ActionFunc with dbmate initialization logic
 func action(f func(*dbmate.DB, *cli.Context) error) cli.ActionFunc {
 	return func(c *cli.Context) error {
-		u, err := getDatabaseURL(c)
+		u, dsn, err := getDatabaseConfig(c)
 		if err != nil {
 			return err
 		}
-		db := dbmate.New(u)
+		db := dbmate.NewWithDSN(u, dsn)
 		db.AutoDumpSchema = !c.Bool("no-dump-schema")
 		db.MigrationsDir = c.String("migrations-dir")
 		db.MigrationsTableName = c.String("migrations-table")
@@ -245,7 +246,7 @@ func action(f func(*dbmate.DB, *cli.Context) error) cli.ActionFunc {
 }
 
 // getDatabaseURL returns the current database url from cli flag or environment variable
-func getDatabaseURL(c *cli.Context) (u *url.URL, err error) {
+func getDatabaseConfig(c *cli.Context) (*url.URL, *dbutil.DSN, error) {
 	// check --url flag first
 	value := c.String("url")
 	if value == "" {
@@ -254,7 +255,16 @@ func getDatabaseURL(c *cli.Context) (u *url.URL, err error) {
 		value = os.Getenv(env)
 	}
 
-	return url.Parse(value)
+	if value == "" {
+		if dsn := os.Getenv("DATABASE_DSN"); dsn != "" {
+			driver := os.Getenv("DBMATE_DRIVER")
+			dsn, err := dbmate.NewDSN(driver, dsn)
+			return nil, &dsn, err
+		}
+	}
+
+	u, err := url.Parse(value)
+	return u, nil, err
 }
 
 // redactLogString attempts to redact passwords from errors

main_test.go

@@ -22,19 +22,19 @@ func TestGetDatabaseUrl(t *testing.T) {
 	ctx := cli.NewContext(app, flagset, nil)
 
 	// no flags defaults to DATABASE_URL
-	u, err := getDatabaseURL(ctx)
+	u, _, err := getDatabaseConfig(ctx)
 	require.NoError(t, err)
 	require.Equal(t, "foo://example.org/one", u.String())
 
 	// --env overwrites DATABASE_URL
 	require.NoError(t, ctx.Set("env", "CUSTOM_URL"))
-	u, err = getDatabaseURL(ctx)
+	u, _, err = getDatabaseConfig(ctx)
 	require.NoError(t, err)
 	require.Equal(t, "foo://example.org/two", u.String())
 
 	// --url takes precedence over preceding two options
 	require.NoError(t, ctx.Set("url", "foo://example.org/three"))
-	u, err = getDatabaseURL(ctx)
+	u, _, err = getDatabaseConfig(ctx)
 	require.NoError(t, err)
 	require.Equal(t, "foo://example.org/three", u.String())
 }

pkg/dbmate/db.go

@@ -11,6 +11,7 @@ import (
 	"path/filepath"
 	"regexp"
 	"sort"
+	"strings"
 	"time"
 
 	"github.com/amacneil/dbmate/v2/pkg/dbutil"
@@ -30,6 +31,14 @@ var (
 	ErrCreateDirectory       = errors.New("unable to create directory")
 )
 
+func NewDSN(driver string, input string) (dbutil.DSN, error) {
+	dsn, err := dbutil.NewDSN(driver, input)
+	if errors.Is(err, dbutil.ErrDriverUnset) {
+		return dsn, fmt.Errorf("%s: expected DBMATE_DRIVER to be set or driver= in DSN: %w", err.Error(), ErrUnsupportedDriver)
+	}
+	return dsn, err
+}
+
 // migrationFileRegexp pattern for valid migration files
 var migrationFileRegexp = regexp.MustCompile(`^(\d+).*\.sql$`)
 
@@ -39,6 +48,8 @@ type DB struct {
 	AutoDumpSchema bool
 	// DatabaseURL is the database connection string
 	DatabaseURL *url.URL
+	// DatabaseDSN is the database connection dsn
+	DatabaseDSN *dbutil.DSN
 	// FS specifies the filesystem, or nil for OS filesystem
 	FS fs.FS
 	// Log is the interface to write stdout
@@ -67,9 +78,15 @@ type StatusResult struct {
 
 // New initializes a new dbmate database
 func New(databaseURL *url.URL) *DB {
+	return NewWithDSN(databaseURL, nil)
+}
+
+// NewWithDSN initializes a new dbmate database with either a url or a dsn
+func NewWithDSN(databaseURL *url.URL, dsn *dbutil.DSN) *DB {
 	return &DB{
 		AutoDumpSchema:      true,
 		DatabaseURL:         databaseURL,
+		DatabaseDSN:         dsn,
 		FS:                  nil,
 		Log:                 os.Stdout,
 		MigrationsDir:       "./db/migrations",
@@ -84,17 +101,25 @@ func New(databaseURL *url.URL) *DB {
 
 // Driver initializes the appropriate database driver
 func (db *DB) Driver() (Driver, error) {
+	var scheme string
 	if db.DatabaseURL == nil || db.DatabaseURL.Scheme == "" {
-		return nil, ErrInvalidURL
+		if db.DatabaseDSN != nil {
+			scheme = strings.ToLower(db.DatabaseDSN.Driver())
+		} else {
+			return nil, ErrInvalidURL
+		}
+	} else {
+		scheme = db.DatabaseURL.Scheme
 	}
 
-	driverFunc := drivers[db.DatabaseURL.Scheme]
+	driverFunc := drivers[scheme]
 	if driverFunc == nil {
-		return nil, fmt.Errorf("%w: %s", ErrUnsupportedDriver, db.DatabaseURL.Scheme)
+		return nil, fmt.Errorf("%w: %s", ErrUnsupportedDriver, scheme)
 	}
 
 	config := DriverConfig{
 		DatabaseURL:         db.DatabaseURL,
+		DatabaseDSN:         db.DatabaseDSN,
 		Log:                 db.Log,
 		MigrationsTableName: db.MigrationsTableName,
 	}

pkg/dbmate/driver.go

@@ -26,6 +26,7 @@ type Driver interface {
 // DriverConfig holds configuration passed to driver constructors
 type DriverConfig struct {
 	DatabaseURL         *url.URL
+	DatabaseDSN         *dbutil.DSN
 	Log                 io.Writer
 	MigrationsTableName string
 }

pkg/dbutil/dbutil.go

@@ -166,3 +166,79 @@ func MustUnescapePath(s string) string {
 
 	return path
 }
+
+type DSN struct {
+	connectionString string
+	driver           string
+}
+
+var ErrDriverUnset = errors.New("driver not set")
+
+func NewDSN(driver string, input string) (DSN, error) {
+	// We support some manipulation of the DSN.
+	// We Don't actually parse the DSN because would bring in a dependency
+	// This means there will be edge cases
+	// However, the DSN manipulation we do is very limited
+	dsn := DSN{
+		connectionString: input,
+	}
+	if driver == "" {
+		driver = dsn.DeleteKey("driver")
+		if driver == "" {
+			return DSN{}, ErrDriverUnset
+		}
+	}
+	dsn.driver = driver
+	return dsn, nil
+}
+
+func (dsn *DSN) Driver() string {
+	return dsn.driver
+}
+
+func (dsn *DSN) ConnectionString() string {
+	return dsn.connectionString
+}
+
+// This function will not work properly if the key value is quoting a space
+func (dsn *DSN) GetKey(key string) string {
+	items := strings.Split(dsn.connectionString, " ")
+	for _, item := range items {
+		if strings.HasPrefix(item, key+"=") {
+			keySplit := strings.Split(item, "=")
+			if len(keySplit) == 2 {
+				return keySplit[1]
+			}
+		}
+	}
+	return ""
+}
+
+// return the original value
+// This will not work properly if the key value is quoting a space
+func (dsn *DSN) DeleteKey(key string) string {
+	return dsn.SetKey(key, "")
+}
+
+// Return the original value
+// This function will not work properly if the key value is quoting a space
+func (dsn *DSN) SetKey(key string, newValue string) (original string) {
+	// delete key (and its value) from the DSN
+	items := strings.Split(dsn.connectionString, " ")
+	newItems := make([]string, 0, len(items))
+	for _, item := range items {
+		if strings.HasPrefix(item, key+"=") {
+			kv := strings.Split(item, "=")
+			if len(kv) == 2 {
+				original = kv[1]
+			}
+		} else {
+			newItems = append(newItems, item)
+		}
+	}
+	if newValue != "" {
+		newItems = append(newItems, newValue)
+	}
+	dsn.connectionString = strings.Join(newItems, " ")
+	return
+}

pkg/dbutil/dbutil_test.go

@@ -56,3 +56,42 @@ func TestQueryValue(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, "7", val)
 }
+
+func TestDSNDriverRequired(t *testing.T) {
+	driver := "postgres"
+	connStr := "user=User dbname=DBName"
+	_, err := dbutil.NewDSN(driver, connStr)
+	require.NoError(t, err)
+	_, err = dbutil.NewDSN("", connStr)
+	require.Error(t, err)
+	_, err = dbutil.NewDSN("", connStr+" driver=postgres")
+	require.NoError(t, err)
+}
+
+func TestDSNGetKey(t *testing.T) {
+	driver := "postgres"
+	connStr := "user=User dbname=DBName"
+	dsn, err := dbutil.NewDSN(driver, connStr)
+	require.NoError(t, err)
+	require.Equal(t, connStr, dsn.ConnectionString())
+
+	userValue := dsn.GetKey("user")
+	require.Equal(t, "User", userValue)
+	dbnameValue := dsn.GetKey("dbname")
+	require.Equal(t, "DBName", dbnameValue)
+}
+
+func TestDSNDeleteKey(t *testing.T) {
+	driver := "postgres"
+	connStr := "user=User dbname=DBName"
+	dsn, err := dbutil.NewDSN(driver, connStr)
+	require.NoError(t, err)
+	require.Equal(t, connStr, dsn.ConnectionString())
+
+	userValue := dsn.DeleteKey("user")
+	require.Equal(t, "User", userValue)
+	require.Equal(t, "dbname=DBName", dsn.ConnectionString())
+	dbnameValue := dsn.DeleteKey("dbname")
+	require.Equal(t, "DBName", dbnameValue)
+	require.Equal(t, "", dsn.ConnectionString())
+}