Allow CLI flags to be specified by environment variable (xataio#152)

Bind command line arguments to env vars. Allow the flags: * `--postgres-url` * `--schema` * `--pgroll-schema` To be set via these env vars: * `PGROLL_PG_URL` * `PGROLL_SCHEMA` * `PGROLL_STATE_SCHEMA` If a flag is specified both via a CLI flag and an env var the CLI flag takes priority.
go-baselib · Oct 3, 2023 · 3eae4ec · 3eae4ec
1 parent bdddd3c
commit 3eae4ec
Show file tree

Hide file tree

Showing 10 changed files with 572 additions and 36 deletions.
diff --git a/cmd/analyze.go b/cmd/analyze.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/spf13/cobra"
 
+	"github.com/xataio/pgroll/cmd/flags"
 	"github.com/xataio/pgroll/pkg/state"
 )
 
@@ -18,13 +19,13 @@ var analyzeCmd = &cobra.Command{
 	Args:   cobra.NoArgs,
 	RunE: func(cmd *cobra.Command, _ []string) error {
 		ctx := cmd.Context()
-		state, err := state.New(ctx, PGURL, StateSchema)
+		state, err := state.New(ctx, flags.PostgresURL(), flags.StateSchema())
 		if err != nil {
 			return err
 		}
 		defer state.Close()
 
-		schema, err := state.ReadSchema(ctx, Schema)
+		schema, err := state.ReadSchema(ctx, flags.Schema())
 		if err != nil {
 			return err
 		}

diff --git a/cmd/flags/flags.go b/cmd/flags/flags.go
@@ -0,0 +1,23 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package flags
+
+import (
+	"github.com/spf13/viper"
+)
+
+func PostgresURL() string {
+	return viper.GetString("PG_URL")
+}
+
+func Schema() string {
+	return viper.GetString("SCHEMA")
+}
+
+func StateSchema() string {
+	return viper.GetString("STATE_SCHEMA")
+}
+
+func LockTimeout() int {
+	return viper.GetInt("LOCK_TIMEOUT")
+}
diff --git a/cmd/root.go b/cmd/root.go
@@ -6,28 +6,28 @@ import (
 	"context"
 
 	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+	"github.com/xataio/pgroll/cmd/flags"
 	"github.com/xataio/pgroll/pkg/roll"
 	"github.com/xataio/pgroll/pkg/state"
 )
 
-var (
-	// PGURL is the Postgres URL to connect to
-	PGURL string
+// Version is the pgroll version
+var Version = "development"
 
-	// Schema is the schema to use for the migration
-	Schema string
-
-	// StateSchema is the Postgres schema where pgroll will store its state
-	StateSchema string
+func init() {
+	viper.SetEnvPrefix("PGROLL")
+	viper.AutomaticEnv()
 
-	// Version is the pgroll version
-	Version = "development"
-)
+	rootCmd.PersistentFlags().String("postgres-url", "postgres://postgres:postgres@localhost?sslmode=disable", "Postgres URL")
+	rootCmd.PersistentFlags().String("schema", "public", "Postgres schema to use for the migration")
+	rootCmd.PersistentFlags().String("pgroll-schema", "pgroll", "Postgres schema to use for pgroll internal state")
+	rootCmd.PersistentFlags().Int("lock-timeout", 500, "Postgres lock timeout in milliseconds for pgroll DDL operations")
 
-func init() {
-	rootCmd.PersistentFlags().StringVar(&PGURL, "postgres-url", "postgres://postgres:postgres@localhost?sslmode=disable", "Postgres URL")
-	rootCmd.PersistentFlags().StringVar(&Schema, "schema", "public", "Postgres schema to use for the migration")
-	rootCmd.PersistentFlags().StringVar(&StateSchema, "pgroll-schema", "pgroll", "Postgres schema in which the migration should be applied")
+	viper.BindPFlag("PG_URL", rootCmd.PersistentFlags().Lookup("postgres-url"))
+	viper.BindPFlag("SCHEMA", rootCmd.PersistentFlags().Lookup("schema"))
+	viper.BindPFlag("STATE_SCHEMA", rootCmd.PersistentFlags().Lookup("pgroll-schema"))
+	viper.BindPFlag("LOCK_TIMEOUT", rootCmd.PersistentFlags().Lookup("lock-timeout"))
 }
 
 var rootCmd = &cobra.Command{
@@ -37,12 +37,17 @@ var rootCmd = &cobra.Command{
 }
 
 func NewRoll(ctx context.Context) (*roll.Roll, error) {
-	state, err := state.New(ctx, PGURL, StateSchema)
+	pgURL := flags.PostgresURL()
+	schema := flags.Schema()
+	stateSchema := flags.StateSchema()
+	lockTimeout := flags.LockTimeout()
+
+	state, err := state.New(ctx, pgURL, stateSchema)
 	if err != nil {
 		return nil, err
 	}
 
-	return roll.New(ctx, PGURL, Schema, state)
+	return roll.New(ctx, pgURL, schema, lockTimeout, state)
 }
 
 // Execute executes the root command.

diff --git a/cmd/start.go b/cmd/start.go
@@ -8,6 +8,7 @@ import (
 	"strings"
 
 	"github.com/pterm/pterm"
+	"github.com/xataio/pgroll/cmd/flags"
 	"github.com/xataio/pgroll/pkg/migrations"
 	"github.com/xataio/pgroll/pkg/roll"
 
@@ -54,7 +55,7 @@ func startCmd() *cobra.Command {
 			}
 
 			version := strings.TrimSuffix(filepath.Base(fileName), filepath.Ext(fileName))
-			viewName := roll.VersionedSchemaName(Schema, version)
+			viewName := roll.VersionedSchemaName(flags.Schema(), version)
 			msg := fmt.Sprintf("New version of the schema available under the postgres %q schema", viewName)
 			sp.Success(msg)
 

diff --git a/cmd/status.go b/cmd/status.go
@@ -7,6 +7,7 @@ import (
 	"encoding/json"
 	"fmt"
 
+	"github.com/xataio/pgroll/cmd/flags"
 	"github.com/xataio/pgroll/pkg/state"
 
 	"github.com/spf13/cobra"
@@ -23,13 +24,13 @@ var statusCmd = &cobra.Command{
 	Short: "Show pgroll status",
 	RunE: func(cmd *cobra.Command, _ []string) error {
 		ctx := cmd.Context()
-		state, err := state.New(ctx, PGURL, StateSchema)
+		state, err := state.New(ctx, flags.PostgresURL(), flags.StateSchema())
 		if err != nil {
 			return err
 		}
 		defer state.Close()
 
-		statusLine, err := statusForSchema(ctx, state, Schema)
+		statusLine, err := statusForSchema(ctx, state, flags.StateSchema())
 		if err != nil {
 			return err
 		}

diff --git a/go.mod b/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/lib/pq v1.10.9
 	github.com/pterm/pterm v0.12.69
 	github.com/spf13/cobra v1.7.0
+	github.com/spf13/viper v1.16.0
 	github.com/stretchr/testify v1.8.4
 	github.com/testcontainers/testcontainers-go v0.23.0
 	github.com/testcontainers/testcontainers-go/modules/postgres v0.23.0
@@ -28,36 +29,45 @@ require (
 	github.com/docker/docker v24.0.5+incompatible // indirect
 	github.com/docker/go-connections v0.4.0 // indirect
 	github.com/docker/go-units v0.5.0 // indirect
+	github.com/fsnotify/fsnotify v1.6.0 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang/protobuf v1.5.3 // indirect
 	github.com/google/uuid v1.3.0 // indirect
 	github.com/gookit/color v1.5.4 // indirect
+	github.com/hashicorp/hcl v1.0.0 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/klauspost/compress v1.16.0 // indirect
 	github.com/lithammer/fuzzysearch v1.1.8 // indirect
 	github.com/magiconair/properties v1.8.7 // indirect
 	github.com/mattn/go-runewidth v0.0.15 // indirect
+	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/moby/patternmatcher v0.5.0 // indirect
 	github.com/moby/sys/sequential v0.5.0 // indirect
 	github.com/moby/term v0.5.0 // indirect
 	github.com/morikuni/aec v1.0.0 // indirect
 	github.com/opencontainers/go-digest v1.0.0 // indirect
 	github.com/opencontainers/image-spec v1.1.0-rc4 // indirect
 	github.com/opencontainers/runc v1.1.5 // indirect
+	github.com/pelletier/go-toml/v2 v2.0.8 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/rivo/uniseg v0.4.4 // indirect
 	github.com/sirupsen/logrus v1.9.0 // indirect
+	github.com/spf13/afero v1.9.5 // indirect
+	github.com/spf13/cast v1.5.1 // indirect
+	github.com/spf13/jwalterweatherman v1.1.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/subosito/gotenv v1.4.2 // indirect
 	github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
 	golang.org/x/mod v0.9.0 // indirect
-	golang.org/x/net v0.9.0 // indirect
+	golang.org/x/net v0.10.0 // indirect
 	golang.org/x/sys v0.11.0 // indirect
 	golang.org/x/term v0.11.0 // indirect
 	golang.org/x/text v0.12.0 // indirect
 	golang.org/x/tools v0.7.0 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20230525234030-28d5490b6b19 // indirect
 	google.golang.org/grpc v1.57.0 // indirect
 	google.golang.org/protobuf v1.30.0 // indirect
+	gopkg.in/ini.v1 v1.67.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )