From 6cbf038cc306e25b36728e3636405a7e5ba3917a Mon Sep 17 00:00:00 2001 From: Brandur Date: Mon, 23 Dec 2024 21:31:43 -0700 Subject: [PATCH] Alternatively accept PG* env vars for database configuration This one's related to #676, in which it was reported that database URLs work somewhat suboptimally when dealing with passwords that contain characters that traditionally require encoding in URLs (e.g. various types of brackets). I was originally going down a path wherein we'd offer an alternative that'd involve adding additional CLI parameters that could work as an alternative to `--database-url` based on these ones from psql: Connection options: -h, --host=HOSTNAME database server host or socket directory (default: "local socket") -p, --port=PORT database server port (default: "5432") -U, --username=USERNAME database user name (default: "brandur") -w, --no-password never prompt for password -W, --password force password prompt (should happen automatically) However, as I implemented it, I found that pgx somewhat surprisingly doesn't allow you to instantiate your own config struct. If not parsed from a URL, it requires you to assemble a more traditional connection string like `user=XX pass=YY database=ZZ` and parse that, which introduces its own special character encoding challenges when using something like a space, equal, or quotation mark. Digging further I realized that pgx automatically supports the whole set of PG* environmental variables like `PGHOST`, `PGPORT`, `PGUSER`, and supporting these would add a solid way of using the River CLI without a database URL using a common interface that's also supported by psql and a variety of other pieces of related software. Even better, since we expect most people to be using a database URL most of the time, env vars give us a way to provide an alternative, but one which doesn't add any complication to the existing CLI interface. In general fewer parameters is better because it keeps the help docs easy to understand. Another advantage is that this will technically let us support more complex connection setups involving things like TLS certificates. All of these vars are supported, and none of which had a great answer with a database URL: * `PGSSLCERT` * `PGSSLKEY` * `PGSSLROOTCERT` * `PGSSLPASSWORD` Fixes #676. --- cmd/river/rivercli/command.go | 57 ++++++++++++++-------- cmd/river/rivercli/river_cli.go | 65 ++++++++++++------------- cmd/river/rivercli/river_cli_test.go | 71 ++++++++++++++++++++++++++-- 3 files changed, 136 insertions(+), 57 deletions(-) diff --git a/cmd/river/rivercli/command.go b/cmd/river/rivercli/command.go index d8c6a09b..31f5283a 100644 --- a/cmd/river/rivercli/command.go +++ b/cmd/river/rivercli/command.go @@ -14,6 +14,7 @@ import ( "github.com/riverqueue/river/cmd/river/riverbench" "github.com/riverqueue/river/rivermigrate" + "github.com/riverqueue/river/rivershared/util/ptrutil" ) const ( @@ -77,7 +78,7 @@ type RunCommandBundle struct { } // RunCommand bootstraps and runs a River CLI subcommand. -func RunCommand[TOpts CommandOpts](ctx context.Context, bundle *RunCommandBundle, command Command[TOpts], opts TOpts) { +func RunCommand[TOpts CommandOpts](ctx context.Context, bundle *RunCommandBundle, command Command[TOpts], opts TOpts) error { procureAndRun := func() (bool, error) { if err := opts.Validate(); err != nil { return false, err @@ -89,16 +90,33 @@ func RunCommand[TOpts CommandOpts](ctx context.Context, bundle *RunCommandBundle Out: bundle.OutStd, } + var databaseURL *string + switch { - // If database URL is still nil after Validate check, then assume this - // command doesn't take one. - case bundle.DatabaseURL == nil: - commandBase.GetBenchmarker = func() BenchmarkerInterface { panic("databaseURL was not set") } - commandBase.GetMigrator = func(config *rivermigrate.Config) (MigratorInterface, error) { panic("databaseURL was not set") } - - case strings.HasPrefix(*bundle.DatabaseURL, uriScheme) || - strings.HasPrefix(*bundle.DatabaseURL, uriSchemeAlias): - dbPool, err := openPgxV5DBPool(ctx, *bundle.DatabaseURL) + case pgEnvConfigured(): + databaseURL = ptrutil.Ptr("") + + case bundle.DatabaseURL != nil: + if !strings.HasPrefix(*bundle.DatabaseURL, uriScheme) && + !strings.HasPrefix(*bundle.DatabaseURL, uriSchemeAlias) { + return false, fmt.Errorf( + "unsupported database URL (`%s`); try one with a `%s` or `%s` scheme/prefix", + *bundle.DatabaseURL, + uriSchemeAlias, + uriScheme, + ) + } + + databaseURL = bundle.DatabaseURL + } + + if databaseURL == nil { + commandBase.GetBenchmarker = func() BenchmarkerInterface { panic("neither PG* env nor databaseURL was not set") } + commandBase.GetMigrator = func(config *rivermigrate.Config) (MigratorInterface, error) { + panic("neither PG* env nor databaseURL was not set") + } + } else { + dbPool, err := openPgxV5DBPool(ctx, *databaseURL) if err != nil { return false, err } @@ -108,14 +126,6 @@ func RunCommand[TOpts CommandOpts](ctx context.Context, bundle *RunCommandBundle commandBase.GetBenchmarker = func() BenchmarkerInterface { return riverbench.NewBenchmarker(driver, commandBase.Logger) } commandBase.GetMigrator = func(config *rivermigrate.Config) (MigratorInterface, error) { return rivermigrate.New(driver, config) } - - default: - return false, fmt.Errorf( - "unsupported database URL (`%s`); try one with a `%s` or `%s` scheme/prefix", - *bundle.DatabaseURL, - uriSchemeAlias, - uriScheme, - ) } command.SetCommandBase(commandBase) @@ -125,11 +135,12 @@ func RunCommand[TOpts CommandOpts](ctx context.Context, bundle *RunCommandBundle ok, err := procureAndRun() if err != nil { - fmt.Fprintf(os.Stderr, "failed: %s\n", err) + return err } - if err != nil || !ok { + if !ok { os.Exit(1) } + return nil } func openPgxV5DBPool(ctx context.Context, databaseURL string) (*pgxpool.Pool, error) { @@ -164,3 +175,9 @@ func openPgxV5DBPool(ctx context.Context, databaseURL string) (*pgxpool.Pool, er return dbPool, nil } + +func pgEnvConfigured() bool { + // Consider these two vars the bare minimum for an env-based configuration. + // return os.Getenv("PGHOST") != "" && os.Getenv("PGDATABASE") != "" + return os.Getenv("PGDATABASE") != "" +} diff --git a/cmd/river/rivercli/river_cli.go b/cmd/river/rivercli/river_cli.go index 66728be1..a5b387f7 100644 --- a/cmd/river/rivercli/river_cli.go +++ b/cmd/river/rivercli/river_cli.go @@ -101,13 +101,20 @@ func (c *CLI) BaseCommandSet() *cobra.Command { Short: "Provides command line facilities for the River job queue", Long: strings.TrimSpace(` Provides command line facilities for the River job queue. + +Commands that need database access will take a --database-url argument, but can +also accept Postgres configuration through its standard set of environment +variables like PGHOST, PGPORT, PGDATABASE, PGUSER, PGPASSWORD, and PGSSLMODE, +with a minimum of PGDATABASE required to use this strategy. --database-url will +take precedence of PG* vars if it's been specified. `), - Run: func(cmd *cobra.Command, args []string) { + RunE: func(cmd *cobra.Command, args []string) error { if rootOpts.Version { - RunCommand(ctx, makeCommandBundle(nil), &version{}, &versionOpts{Name: c.name}) - } else { - _ = cmd.Usage() + return RunCommand(ctx, makeCommandBundle(nil), &version{}, &versionOpts{Name: c.name}) } + + _ = cmd.Usage() + return nil }, } rootCmd.SetOut(c.out) @@ -119,17 +126,8 @@ Provides command line facilities for the River job queue. rootCmd.Flags().BoolVar(&rootOpts.Version, "version", false, "print version information") } - mustMarkFlagRequired := func(cmd *cobra.Command, name string) { - // We just panic here because this will never happen outside of an error - // in development. - if err := cmd.MarkFlagRequired(name); err != nil { - panic(err) - } - } - addDatabaseURLFlag := func(cmd *cobra.Command, databaseURL *string) { cmd.Flags().StringVar(databaseURL, "database-url", "", "URL of the database (should look like `postgres://...`") - mustMarkFlagRequired(cmd, "database-url") } addLineFlag := func(cmd *cobra.Command, line *string) { cmd.Flags().StringVar(line, "line", "", "migration line to operate on (default: main)") @@ -155,8 +153,8 @@ before starting the client, and works until all jobs are finished. The database in --database-url will have its jobs table truncated, so make sure to use a development database only. `), - Run: func(cmd *cobra.Command, args []string) { - RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL), &bench{}, &opts) + RunE: func(cmd *cobra.Command, args []string) error { + return RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL), &bench{}, &opts) }, } addDatabaseURLFlag(cmd, &opts.DatabaseURL) @@ -194,8 +192,8 @@ SQL being run can be output using --show-sql, and executing real database operations can be prevented with --dry-run. Combine --show-sql and --dry-run to dump prospective migrations that would be applied to stdout. `), - Run: func(cmd *cobra.Command, args []string) { - RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL), &migrateDown{}, &opts) + RunE: func(cmd *cobra.Command, args []string) error { + return RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL), &migrateDown{}, &opts) }, } addMigrateFlags(cmd, &opts) @@ -232,8 +230,8 @@ framework, which aren't necessary if using an external framework: river migrate-get --all --exclude-version 1 --up > river_all.up.sql river migrate-get --all --exclude-version 1 --down > river_all.down.sql `), - Run: func(cmd *cobra.Command, args []string) { - RunCommand(ctx, makeCommandBundle(nil), &migrateGet{}, &opts) + RunE: func(cmd *cobra.Command, args []string) error { + return RunCommand(ctx, makeCommandBundle(nil), &migrateGet{}, &opts) }, } cmd.Flags().BoolVar(&opts.All, "all", false, "print all migrations; down migrations are printed in descending order") @@ -259,8 +257,8 @@ framework, which aren't necessary if using an external framework: Long: strings.TrimSpace(` TODO `), - Run: func(cmd *cobra.Command, args []string) { - RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL), &migrateList{}, &opts) + RunE: func(cmd *cobra.Command, args []string) error { + return RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL), &migrateList{}, &opts) }, } addDatabaseURLFlag(cmd, &opts.DatabaseURL) @@ -285,8 +283,8 @@ SQL being run can be output using --show-sql, and executing real database operations can be prevented with --dry-run. Combine --show-sql and --dry-run to dump prospective migrations that would be applied to stdout. `), - Run: func(cmd *cobra.Command, args []string) { - RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL), &migrateUp{}, &opts) + RunE: func(cmd *cobra.Command, args []string) error { + return RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL), &migrateUp{}, &opts) }, } addMigrateFlags(cmd, &opts) @@ -307,12 +305,11 @@ are outstanding migrations that still need to be run. Can be paired with river migrate-up --dry-run --show-sql to dump information on migrations that need to be run, but without running them. `), - Run: func(cmd *cobra.Command, args []string) { - RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL), &validate{}, &opts) + RunE: func(cmd *cobra.Command, args []string) error { + return RunCommand(ctx, makeCommandBundle(&opts.DatabaseURL), &validate{}, &opts) }, } addDatabaseURLFlag(cmd, &opts.DatabaseURL) - mustMarkFlagRequired(cmd, "database-url") cmd.Flags().StringVar(&opts.Line, "line", "", "migration line to operate on (default: main)") rootCmd.AddCommand(cmd) } @@ -325,8 +322,8 @@ migrations that need to be run, but without running them. Long: strings.TrimSpace(` Print River and Go version information. `), - Run: func(cmd *cobra.Command, args []string) { - RunCommand(ctx, makeCommandBundle(nil), &version{}, &versionOpts{Name: c.name}) + RunE: func(cmd *cobra.Command, args []string) error { + return RunCommand(ctx, makeCommandBundle(nil), &version{}, &versionOpts{Name: c.name}) }, } rootCmd.AddCommand(cmd) @@ -347,8 +344,8 @@ type benchOpts struct { } func (o *benchOpts) Validate() error { - if o.DatabaseURL == "" { - return errors.New("database URL cannot be empty") + if o.DatabaseURL == "" && !pgEnvConfigured() { + return errors.New("either PG* env vars or --database-url must be set") } return nil @@ -375,8 +372,8 @@ type migrateOpts struct { } func (o *migrateOpts) Validate() error { - if o.DatabaseURL == "" { - return errors.New("database URL cannot be empty") + if o.DatabaseURL == "" && !pgEnvConfigured() { + return errors.New("either PG* env vars or --database-url must be set") } return nil @@ -604,8 +601,8 @@ type validateOpts struct { } func (o *validateOpts) Validate() error { - if o.DatabaseURL == "" { - return errors.New("database URL cannot be empty") + if o.DatabaseURL == "" && !pgEnvConfigured() { + return errors.New("either PG* env vars or --database-url must be set") } return nil diff --git a/cmd/river/rivercli/river_cli_test.go b/cmd/river/rivercli/river_cli_test.go index 02337815..f7530d46 100644 --- a/cmd/river/rivercli/river_cli_test.go +++ b/cmd/river/rivercli/river_cli_test.go @@ -2,8 +2,11 @@ package rivercli import ( "bytes" + "cmp" "context" "fmt" + "net/url" + "os" "runtime/debug" "strings" "testing" @@ -113,16 +116,25 @@ func TestBaseCommandSetIntegration(t *testing.T) { cmd, _ := setup(t) cmd.SetArgs([]string{"--debug", "--verbose"}) - require.EqualError(t, cmd.Execute(), `if any flags in the group [debug verbose] are set none of the others can be; [debug verbose] were all set`) + require.EqualError(t, cmd.Execute(), "if any flags in the group [debug verbose] are set none of the others can be; [debug verbose] were all set") }) - t.Run("MigrateDownMissingDatabaseURL", func(t *testing.T) { + t.Run("DatabaseURLWithInvalidPrefix", func(t *testing.T) { + t.Parallel() + + cmd, _ := setup(t) + + cmd.SetArgs([]string{"migrate-down", "--database-url", "post://"}) + require.EqualError(t, cmd.Execute(), "unsupported database URL (`post://`); try one with a `postgres://` or `postgresql://` scheme/prefix") + }) + + t.Run("MissingDatabaseURLAndPGEnv", func(t *testing.T) { t.Parallel() cmd, _ := setup(t) cmd.SetArgs([]string{"migrate-down"}) - require.EqualError(t, cmd.Execute(), `required flag(s) "database-url" not set`) + require.EqualError(t, cmd.Execute(), "either PG* env vars or --database-url must be set") }) t.Run("VersionFlag", func(t *testing.T) { @@ -158,6 +170,59 @@ Built with %s }) } +// Same as the above, but non-parallel so tests can use `t.Setenv`. +func TestBaseCommandSetNonParallel(t *testing.T) { + type testBundle struct { + out *bytes.Buffer + } + + setup := func(t *testing.T) (*cobra.Command, *testBundle) { + t.Helper() + + cli := NewCLI(&Config{ + DriverProcurer: &TestDriverProcurer{}, + Name: "River", + }) + + var out bytes.Buffer + cli.SetOut(&out) + + return cli.BaseCommandSet(), &testBundle{ + out: &out, + } + } + + t.Run("PGEnvWithoutDatabaseURL", func(t *testing.T) { + cmd, _ := setup(t) + + t.Logf("TEST_DATABASE_URL = %s", os.Getenv("TEST_DATABASE_URL")) + + // Deconstruct a database URL into its PG* components. This path is the + // one that gets taken in CI, but could work locally as well. + if databaseURL := os.Getenv("TEST_DATABASE_URL"); databaseURL != "" { + databaseU, err := url.Parse(databaseURL) + require.NoError(t, err) + + t.Setenv("PGDATABASE", databaseU.Path[1:]) + t.Setenv("PGHOST", databaseU.Hostname()) + pass, _ := databaseU.User.Password() + t.Setenv("PGPASSWORD", pass) + t.Setenv("PGPORT", cmp.Or(databaseU.Port(), "5432")) + t.Setenv("PGSSLMODE", databaseU.Query().Get("sslmode")) + t.Setenv("PGUSER", databaseU.User.Username()) + } else { + // With no `TEST_DATABASE_URL` available, try a simpler alternative + // configuration. Requires a database on localhost that doesn't + // require authentication, which should exist from testdbman. + t.Setenv("PGDATABASE", "river_test") + t.Setenv("PGHOST", "localhost") + } + + cmd.SetArgs([]string{"validate"}) + require.NoError(t, cmd.Execute()) + }) +} + func TestMigrateList(t *testing.T) { t.Parallel()