Skip to content

Commit

Permalink
Alternatively accept PG* env vars for database configuration
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@brandur
brandur committed Dec 24, 2024
1 parent e7ae173 commit 4c7a74b
Show file tree
Hide file tree
Showing 3 changed files with 113 additions and 60 deletions.
57 changes: 37 additions & 20 deletions cmd/river/rivercli/command.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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
Expand All @@ -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
}
Expand All @@ -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)
Expand All @@ -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) {
Expand Down Expand Up @@ -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") != ""
}
65 changes: 31 additions & 34 deletions cmd/river/rivercli/river_cli.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)")
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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")
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
}
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
51 changes: 45 additions & 6 deletions cmd/river/rivercli/river_cli_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,11 @@ package rivercli

import (
"bytes"
"cmp"
"context"
"fmt"
"net/url"
"os"
"runtime/debug"
"strings"
"testing"
Expand Down Expand Up @@ -84,9 +87,7 @@ func (p *TestDriverProcurer) ProcurePgxV5(pool *pgxpool.Pool) riverdriver.Driver

// High level integration tests that operate on the Cobra command directly. This
// isn't always appropriate because there's no way to inject a test transaction.
func TestBaseCommandSetIntegration(t *testing.T) {
t.Parallel()

func TestBaseCommandSetIntegration(t *testing.T) { //nolint:tparallel
type testBundle struct {
out *bytes.Buffer
}
Expand All @@ -113,16 +114,54 @@ 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("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("MigrateDownMissingDatabaseURL", func(t *testing.T) {
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("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 {
// Try a simpler alternative configuration. Requires a database on
// localhost that doesn't require authentication.
t.Setenv("PGDATABASE", "river_test")
t.Setenv("PGHOST", "localhost")
}

cmd.SetArgs([]string{"validate"})
require.NoError(t, cmd.Execute())
})

t.Run("VersionFlag", func(t *testing.T) {
Expand Down

0 comments on commit 4c7a74b

Please sign in to comment.