reuse ConnectToPostgres

Emyrk · Emyrk · commit 1b284ecca7ca · 2024-12-30T15:33:35.000-06:00
diff --git a/cli/resetpassword.go b/cli/resetpassword.go
@@ -3,22 +3,27 @@
 package cli
 
 import (
-	"database/sql"
 	"fmt"
 
 	"golang.org/x/xerrors"
 
+	"cdr.dev/slog"
+	"cdr.dev/slog/sloggers/sloghuman"
+	"github.com/coder/coder/v2/coderd/database/awsiamrds"
+	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/pretty"
 	"github.com/coder/serpent"
 
 	"github.com/coder/coder/v2/cli/cliui"
 	"github.com/coder/coder/v2/coderd/database"
-	"github.com/coder/coder/v2/coderd/database/migrations"
 	"github.com/coder/coder/v2/coderd/userpassword"
 )
 
 func (*RootCmd) resetPassword() *serpent.Command {
-	var postgresURL string
+	var (
+		postgresURL  string
+		postgresAuth string
+	)
 
 	root := &serpent.Command{
 		Use:        "reset-password <username>",
@@ -27,20 +32,25 @@ func (*RootCmd) resetPassword() *serpent.Command {
 		Handler: func(inv *serpent.Invocation) error {
 			username := inv.Args[0]
 
-			sqlDB, err := sql.Open("postgres", postgresURL)
-			if err != nil {
-				return xerrors.Errorf("dial postgres: %w", err)
+			logger := slog.Make(sloghuman.Sink(inv.Stdout))
+			if ok, _ := inv.ParsedFlags().GetBool("verbose"); ok {
+				logger = logger.Leveled(slog.LevelDebug)
 			}
-			defer sqlDB.Close()
-			err = sqlDB.Ping()
-			if err != nil {
-				return xerrors.Errorf("ping postgres: %w", err)
+
+			sqlDriver := "postgres"
+			if codersdk.PostgresAuth(postgresAuth) == codersdk.PostgresAuthAWSIAMRDS {
+				var err error
+				sqlDriver, err = awsiamrds.Register(inv.Context(), sqlDriver)
+				if err != nil {
+					return xerrors.Errorf("register aws rds iam auth: %w", err)
+				}
 			}
 
-			err = migrations.EnsureClean(sqlDB)
+			sqlDB, err := ConnectToPostgres(inv.Context(), logger, false, sqlDriver, postgresURL)
 			if err != nil {
-				return xerrors.Errorf("database needs migration: %w", err)
+				return xerrors.Errorf("dial postgres: %w", err)
 			}
+
 			db := database.New(sqlDB)
 
 			user, err := db.GetUserByEmailOrUsername(inv.Context(), database.GetUserByEmailOrUsernameParams{
@@ -97,6 +107,14 @@ func (*RootCmd) resetPassword() *serpent.Command {
 			Env:         "CODER_PG_CONNECTION_URL",
 			Value:       serpent.StringOf(&postgresURL),
 		},
+		serpent.Option{
+			Name:        "Postgres Connection Auth",
+			Description: "Type of auth to use when connecting to postgres.",
+			Flag:        "postgres-connection-auth",
+			Env:         "CODER_PG_CONNECTION_AUTH",
+			Default:     "password",
+			Value:       serpent.EnumOf(&postgresAuth, codersdk.PostgresAuthDrivers...),
+		},
 	}
 
 	return root
