From a9aaaad1fba65024010a0b61b1715bd37a472140 Mon Sep 17 00:00:00 2001
From: Adrian Macneil <adrian@adrianmacneil.com>
Date: Sat, 8 Aug 2020 11:19:33 -0700
Subject: [PATCH] Add --url flag (#150)

---
 README.md    |  9 +++++++++
 main.go      | 16 +++++++++++++---
 main_test.go | 35 +++++++++++++++++++----------------
 3 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/README.md b/README.md
index b536350..14889c9 100644
--- a/README.md
+++ b/README.md
@@ -312,6 +312,7 @@ Please note that the `wait` command does not verify whether your specified datab
 
 The following command line options are available with all commands. You must use command line arguments in the order `dbmate [global options] command [command options]`.
 
+* `--url, -u "protocol://host:port/dbname"` - specify the database url directly.
 * `--env, -e "DATABASE_URL"` - specify an environment variable to read the database connection URL from.
 * `--migrations-dir, -d "./db/migrations"` - where to keep the migration files.
 * `--schema-file, -s "./db/schema.sql"` - a path to keep the schema.sql file.
@@ -335,6 +336,14 @@ Creating: myapp_test
 Applying: 20151127184807_create_users_table.sql
 ```
 
+Alternatively, you can specify the url directly on the command line:
+
+```sh
+$ 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.
+
 ## FAQ
 
 **How do I use dbmate under Alpine linux?**
diff --git a/main.go b/main.go
index 39f4357..d05a12a 100644
--- a/main.go
+++ b/main.go
@@ -34,6 +34,11 @@ func NewApp() *cli.App {
 	app.Version = dbmate.Version
 
 	app.Flags = []cli.Flag{
+		&cli.StringFlag{
+			Name:    "url",
+			Aliases: []string{"u"},
+			Usage:   "specify the database URL",
+		},
 		&cli.StringFlag{
 			Name:    "env",
 			Aliases: []string{"e"},
@@ -220,10 +225,15 @@ func action(f func(*dbmate.DB, *cli.Context) error) cli.ActionFunc {
 	}
 }
 
-// getDatabaseURL returns the current environment database url
+// getDatabaseURL returns the current database url from cli flag or environment variable
 func getDatabaseURL(c *cli.Context) (u *url.URL, err error) {
-	env := c.String("env")
-	value := os.Getenv(env)
+	// check --url flag first
+	value := c.String("url")
+	if value == "" {
+		// if empty, default to --env or DATABASE_URL
+		env := c.String("env")
+		value = os.Getenv(env)
+	}
 
 	return url.Parse(value)
 }
diff --git a/main_test.go b/main_test.go
index 2aa3ab3..df1cab5 100644
--- a/main_test.go
+++ b/main_test.go
@@ -2,7 +2,6 @@ package main
 
 import (
 	"flag"
-	"net/url"
 	"os"
 	"testing"
 
@@ -10,30 +9,34 @@ import (
 	"github.com/urfave/cli/v2"
 )
 
-func testContext(t *testing.T, u *url.URL) *cli.Context {
-	err := os.Setenv("DATABASE_URL", u.String())
-	require.NoError(t, err)
+func TestGetDatabaseUrl(t *testing.T) {
+	// set environment variables
+	require.NoError(t, os.Setenv("DATABASE_URL", "foo://example.org/one"))
+	require.NoError(t, os.Setenv("CUSTOM_URL", "foo://example.org/two"))
 
 	app := NewApp()
 	flagset := flag.NewFlagSet(app.Name, flag.ContinueOnError)
 	for _, f := range app.Flags {
-		_ = f.Apply(flagset)
+		require.NoError(t, f.Apply(flagset))
 	}
+	ctx := cli.NewContext(app, flagset, nil)
 
-	return cli.NewContext(app, flagset, nil)
-}
-
-func TestGetDatabaseUrl(t *testing.T) {
-	envURL, err := url.Parse("foo://example.org/db")
-	require.NoError(t, err)
-	ctx := testContext(t, envURL)
-
+	// no flags defaults to DATABASE_URL
 	u, err := getDatabaseURL(ctx)
 	require.NoError(t, err)
+	require.Equal(t, "foo://example.org/one", u.String())
 
-	require.Equal(t, "foo", u.Scheme)
-	require.Equal(t, "example.org", u.Host)
-	require.Equal(t, "/db", u.Path)
+	// --env overwrites DATABASE_URL
+	require.NoError(t, ctx.Set("env", "CUSTOM_URL"))
+	u, err = getDatabaseURL(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)
+	require.NoError(t, err)
+	require.Equal(t, "foo://example.org/three", u.String())
 }
 
 func TestRedactLogString(t *testing.T) {