From c7ce3e70daa632ff56c49c5f8633c776c8012b53 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Tue, 20 Dec 2022 20:51:17 +0200
Subject: [PATCH] feat: Add --raw-url to coder server postgres-builtin-*
 commands (#5478)

---
 cli/server.go      | 39 ++++++++++++++++++++++++++-------------
 cli/server_test.go | 13 +++++++++++++
 2 files changed, 39 insertions(+), 13 deletions(-)

diff --git a/cli/server.go b/cli/server.go
index cdb444e1dc..740782779a 100644
--- a/cli/server.go
+++ b/cli/server.go
@@ -921,7 +921,8 @@ func Server(vip *viper.Viper, newAPI func(context.Context, *coderd.Options) (*co
 		},
 	}
 
-	root.AddCommand(&cobra.Command{
+	var pgRawURL bool
+	postgresBuiltinURLCmd := &cobra.Command{
 		Use:   "postgres-builtin-url",
 		Short: "Output the connection URL for the built-in PostgreSQL deployment.",
 		RunE: func(cmd *cobra.Command, _ []string) error {
@@ -930,37 +931,49 @@ func Server(vip *viper.Viper, newAPI func(context.Context, *coderd.Options) (*co
 			if err != nil {
 				return err
 			}
-			_, _ = fmt.Fprintf(cmd.OutOrStdout(), "psql %q\n", url)
+			if pgRawURL {
+				_, _ = fmt.Fprintf(cmd.OutOrStdout(), "%s\n", url)
+			} else {
+				_, _ = fmt.Fprintf(cmd.OutOrStdout(), "%s\n", cliui.Styles.Code.Render(fmt.Sprintf("psql %q", url)))
+			}
 			return nil
 		},
-	})
-
-	root.AddCommand(&cobra.Command{
+	}
+	postgresBuiltinServeCmd := &cobra.Command{
 		Use:   "postgres-builtin-serve",
 		Short: "Run the built-in PostgreSQL deployment.",
 		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+
 			cfg := createConfig(cmd)
 			logger := slog.Make(sloghuman.Sink(cmd.ErrOrStderr()))
 			if ok, _ := cmd.Flags().GetBool(varVerbose); ok {
 				logger = logger.Leveled(slog.LevelDebug)
 			}
 
-			url, closePg, err := startBuiltinPostgres(cmd.Context(), cfg, logger)
+			ctx, cancel := signal.NotifyContext(ctx, InterruptSignals...)
+			defer cancel()
+
+			url, closePg, err := startBuiltinPostgres(ctx, cfg, logger)
 			if err != nil {
 				return err
 			}
 			defer func() { _ = closePg() }()
 
-			cmd.Println(cliui.Styles.Code.Render("psql \"" + url + "\""))
+			if pgRawURL {
+				_, _ = fmt.Fprintf(cmd.OutOrStdout(), "%s\n", url)
+			} else {
+				_, _ = fmt.Fprintf(cmd.OutOrStdout(), "%s\n", cliui.Styles.Code.Render(fmt.Sprintf("psql %q", url)))
+			}
 
-			stopChan := make(chan os.Signal, 1)
-			defer signal.Stop(stopChan)
-			signal.Notify(stopChan, os.Interrupt)
-
-			<-stopChan
+			<-ctx.Done()
 			return nil
 		},
-	})
+	}
+	postgresBuiltinURLCmd.Flags().BoolVar(&pgRawURL, "raw-url", false, "Output the raw connection URL instead of a psql command.")
+	postgresBuiltinServeCmd.Flags().BoolVar(&pgRawURL, "raw-url", false, "Output the raw connection URL instead of a psql command.")
+
+	root.AddCommand(postgresBuiltinURLCmd, postgresBuiltinServeCmd)
 
 	deployment.AttachFlags(root.Flags(), vip, false)
 
diff --git a/cli/server_test.go b/cli/server_test.go
index 5ff05c6b83..5a1f6ed3fa 100644
--- a/cli/server_test.go
+++ b/cli/server_test.go
@@ -118,6 +118,19 @@ func TestServer(t *testing.T) {
 
 		pty.ExpectMatch("psql")
 	})
+	t.Run("BuiltinPostgresURLRaw", func(t *testing.T) {
+		t.Parallel()
+		root, _ := clitest.New(t, "server", "postgres-builtin-url", "--raw-url")
+		pty := ptytest.New(t)
+		root.SetOutput(pty.Output())
+		err := root.Execute()
+		require.NoError(t, err)
+
+		got := pty.ReadLine()
+		if !strings.HasPrefix(got, "postgres://") {
+			t.Fatalf("expected postgres URL to start with \"postgres://\", got %q", got)
+		}
+	})
 
 	// Validate that a warning is printed that it may not be externally
 	// reachable.