From ece5d3cf0e1702d6970843b59f738058899c3a0b Mon Sep 17 00:00:00 2001
From: Adrian Macneil <adrian@adrianmacneil.com>
Date: Wed, 25 Nov 2015 11:32:13 -0800
Subject: [PATCH] Allow database URL to be specified via command line

---
 commands.go                      | 13 ++++++++-----
 commands_test.go                 | 17 +++++++++++++++--
 driver/postgres/postgres_test.go |  2 +-
 main.go                          | 13 ++++++++++++-
 4 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/commands.go b/commands.go
index 6575bea..57773f9 100644
--- a/commands.go
+++ b/commands.go
@@ -16,7 +16,7 @@ import (
 
 // CreateCommand creates the current database
 func CreateCommand(ctx *cli.Context) error {
-	u, err := GetDatabaseURL()
+	u, err := GetDatabaseURL(ctx)
 	if err != nil {
 		return err
 	}
@@ -31,7 +31,7 @@ func CreateCommand(ctx *cli.Context) error {
 
 // DropCommand drops the current database (if it exists)
 func DropCommand(ctx *cli.Context) error {
-	u, err := GetDatabaseURL()
+	u, err := GetDatabaseURL(ctx)
 	if err != nil {
 		return err
 	}
@@ -86,8 +86,11 @@ func NewCommand(ctx *cli.Context) error {
 }
 
 // GetDatabaseURL returns the current environment database url
-func GetDatabaseURL() (u *url.URL, err error) {
-	return url.Parse(os.Getenv("DATABASE_URL"))
+func GetDatabaseURL(ctx *cli.Context) (u *url.URL, err error) {
+	env := ctx.GlobalString("env")
+	value := os.Getenv(env)
+
+	return url.Parse(value)
 }
 
 func doTransaction(db *sql.DB, txFunc func(shared.Transaction) error) error {
@@ -116,7 +119,7 @@ func MigrateCommand(ctx *cli.Context) error {
 		return fmt.Errorf("No migration files found.")
 	}
 
-	u, err := GetDatabaseURL()
+	u, err := GetDatabaseURL(ctx)
 	if err != nil {
 		return err
 	}
diff --git a/commands_test.go b/commands_test.go
index daab091..65ddd85 100644
--- a/commands_test.go
+++ b/commands_test.go
@@ -1,16 +1,29 @@
 package main_test
 
 import (
+	"flag"
 	"github.com/adrianmacneil/dbmate"
+	"github.com/codegangsta/cli"
 	"github.com/stretchr/testify/require"
 	"os"
 	"testing"
 )
 
-func TestGetDatabaseUrl(t *testing.T) {
+func newContext() *cli.Context {
+	app := main.NewApp()
+	flagset := flag.NewFlagSet(app.Name, flag.ContinueOnError)
+	for _, f := range app.Flags {
+		f.Apply(flagset)
+	}
+
+	return cli.NewContext(app, flagset, nil)
+}
+
+func TestGetDatabaseUrl_Default(t *testing.T) {
 	os.Setenv("DATABASE_URL", "postgres://example.org/db")
 
-	u, err := main.GetDatabaseURL()
+	ctx := newContext()
+	u, err := main.GetDatabaseURL(ctx)
 	require.Nil(t, err)
 
 	require.Equal(t, "postgres", u.Scheme)
diff --git a/driver/postgres/postgres_test.go b/driver/postgres/postgres_test.go
index 0efa9b7..e352d5a 100644
--- a/driver/postgres/postgres_test.go
+++ b/driver/postgres/postgres_test.go
@@ -11,7 +11,7 @@ import (
 
 func testURL(t *testing.T) *url.URL {
 	str := os.Getenv("POSTGRES_PORT")
-	require.NotEmpty(t, str)
+	require.NotEmpty(t, str, "missing POSTGRES_PORT environment variable")
 
 	u, err := url.Parse(str)
 	require.Nil(t, err)
diff --git a/main.go b/main.go
index 0f4e20e..2f1d2c7 100644
--- a/main.go
+++ b/main.go
@@ -11,6 +11,12 @@ import (
 func main() {
 	loadDotEnv()
 
+	app := NewApp()
+	app.Run(os.Args)
+}
+
+// NewApp creates a new command line app
+func NewApp() *cli.App {
 	app := cli.NewApp()
 	app.Name = "dbmate"
 	app.Usage = "A lightweight, framework-independent database migration tool."
@@ -21,6 +27,11 @@ func main() {
 			Value: "./db/migrations",
 			Usage: "specify the directory containing migration files",
 		},
+		cli.StringFlag{
+			Name:  "env, e",
+			Value: "DATABASE_URL",
+			Usage: "specify an environment variable containing the database URL",
+		},
 	}
 
 	app.Commands = []cli.Command{
@@ -54,7 +65,7 @@ func main() {
 		},
 	}
 
-	app.Run(os.Args)
+	return app
 }
 
 type command func(*cli.Context) error