Merge pull request #233 from planetscale/nickvanw/connect-exec

connect: add experimental feature for having connect run a command after

Author: nickvanw

go.mod

@@ -15,6 +15,7 @@ require (
 	github.com/kataras/tablewriter v0.0.0-20180708051242-e063d29b7c23 // indirect
 	github.com/lensesio/tableprinter v0.0.0-20201125135848-89e81fc956e7
 	github.com/mattn/go-runewidth v0.0.9 // indirect
+	github.com/mattn/go-shellwords v1.0.11
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/pkg/browser v0.0.0-20201112035734-206646e67786
 	github.com/pkg/errors v0.9.1

go.sum

@@ -222,6 +222,8 @@ github.com/mattn/go-isatty v0.0.12 h1:wuysRhFDzyxgEmMf5xjvJ2M9dZoWAXNNr5LSBS7uHX
 github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU=
 github.com/mattn/go-runewidth v0.0.9 h1:Lm995f3rfxdpd6TSmuVCHVb/QhupuXlYr8sCI/QdE+0=
 github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI=
+github.com/mattn/go-shellwords v1.0.11 h1:vCoR9VPpsk/TZFW2JwK5I9S0xdrtUq2bph6/YjEPnaw=
+github.com/mattn/go-shellwords v1.0.11/go.mod h1:EZzvwXDESEeg03EKmM+RmDnNOPKG4lLtQsUlTZDWQ8Y=
 github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
 github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b h1:j7+1HpAFS1zy5+Q4qx1fWh90gTKwiN4QCGoY9TWyyO4=
 github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b/go.mod h1:01TrycV0kFyexm33Z7vhZRXopbI8J3TDReVlkTgMUxE=

internal/cmd/connect/connect.go

@@ -5,9 +5,11 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"os/exec"
 	"runtime"
 	"syscall"
 
+	"github.com/mattn/go-shellwords"
 	"github.com/planetscale/cli/internal/cmdutil"
 	"github.com/planetscale/cli/internal/printer"
 	"github.com/planetscale/cli/internal/promptutil"
@@ -21,8 +23,9 @@ import (
 
 func ConnectCmd(ch *cmdutil.Helper) *cobra.Command {
 	var flags struct {
-		localAddr  string
-		remoteAddr string
+		localAddr   string
+		remoteAddr  string
+		execCommand string
 	}
 
 	cmd := &cobra.Command{
@@ -47,10 +50,6 @@ argument:
 			ctx := context.Background()
 			database := args[0]
 
-			if !printer.IsTTY || ch.Printer.Format() != printer.Human {
-				return errors.New("pscale connect only works in interactive mode")
-			}
-
 			client, err := ch.Config.NewClientFromConfig()
 			if err != nil {
 				return err
@@ -97,12 +96,18 @@ argument:
 				Logger:     cmdutil.NewZapLogger(ch.Debug()),
 			}
 
-			err = runProxy(proxyOpts, database, branch)
+			proxyReady := make(chan string, 1)
+
+			if flags.execCommand != "" {
+				go runCommand(ctx, flags.execCommand, database, branch, proxyReady)
+			}
+
+			err = runProxy(proxyOpts, database, branch, proxyReady)
 			if err != nil {
 				if isAddrInUse(err) {
 					ch.Printer.Println("Tried address 127.0.0.1:3306, but it's already in use. Picking up a random port ...")
 					proxyOpts.LocalAddr = "127.0.0.1:0"
-					return runProxy(proxyOpts, database, branch)
+					return runProxy(proxyOpts, database, branch, proxyReady)
 				}
 				return err
 			}
@@ -117,18 +122,19 @@ argument:
 	cmd.PersistentFlags().StringVar(&flags.remoteAddr, "remote-addr", "",
 		"PlanetScale Database remote network address. By default the remote address is populated automatically from the PlanetScale API.")
 	cmd.MarkPersistentFlagRequired("org") // nolint:errcheck
+	cmd.PersistentFlags().StringVar(&flags.execCommand, "execute", "", "Run this command after successfully connecting to the database.")
 
 	return cmd
 }
 
-func runProxy(proxyOpts proxy.Options, database, branch string) error {
+func runProxy(proxyOpts proxy.Options, database, branch string, ready chan string) error {
 	ctx := context.Background()
 	p, err := proxy.NewClient(proxyOpts)
 	if err != nil {
 		return fmt.Errorf("couldn't create proxy client: %s", err)
 	}
 
-	go func() {
+	go func(ready chan string) {
 		// this is blocking and will only return once p.Run() below is
 		// invoked
 		addr, err := p.LocalAddr()
@@ -142,14 +148,51 @@ func runProxy(proxyOpts proxy.Options, database, branch string) error {
 			printer.BoldBlue(branch),
 			printer.BoldBlue(addr.String()),
 		)
-	}()
+		ready <- addr.String()
+	}(ready)
 
 	// TODO(fatih): replace with signal.NotifyContext once Go 1.16 is released
 	// https://go-review.googlesource.com/c/go/+/219640
 	ctx = sigutil.WithSignal(ctx, syscall.SIGINT, syscall.SIGTERM)
 	return p.Run(ctx)
 }
 
+func runCommand(ctx context.Context, command, database, branch string, ready chan string) {
+	args, err := shellwords.Parse(command)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "\nfailed to parse command, not running: %s", err)
+		return
+	}
+	addr := <-ready
+
+	ctx = sigutil.WithSignal(ctx, syscall.SIGINT, syscall.SIGTERM)
+	child := exec.CommandContext(ctx, args[0], args[1:]...)
+	child.Env = os.Environ()
+	child.Stdout = os.Stdout
+	child.Stderr = os.Stderr
+
+	connStr := fmt.Sprintf("DATABASE_URL=mysql2://root@%s/%s", addr, database)
+	child.Env = append(child.Env, connStr)
+
+	hostEnv := fmt.Sprintf("PLANETSCALE_DATABASE_HOST=%s", addr)
+	child.Env = append(child.Env, hostEnv)
+
+	dbName := fmt.Sprintf("PLANETSCALE_DATABASE_NAME=%s", database)
+	child.Env = append(child.Env, dbName)
+
+	branchName := fmt.Sprintf("PLANETSCALE_BRANCH_NAME=%s", branch)
+	child.Env = append(child.Env, branchName)
+
+	// todo(nickvanw): right now, this starts the process and then doesn't track what happens next
+	// if the process exits (non-zero or otherwise) the proxy will remain running
+	//
+	// the right behavior is probably to at least allow the user to configure that the proxy should
+	// exit when this process exits, and pass through the exit code
+	if err := child.Start(); err != nil {
+		fmt.Fprintf(os.Stderr, "\nfailed to execute command: %s", err)
+	}
+}
+
 // isAddrInUse returns an error if the error indicates that the given address
 // is already in use. Becaue different OS return different error messages, we
 // try to get the underlying error.

internal/cmd/root.go

@@ -20,6 +20,7 @@ import (
 	"io/fs"
 	"log"
 	"os"
+	"strings"
 
 	"github.com/planetscale/cli/internal/cmd/auth"
 	"github.com/planetscale/cli/internal/cmd/backup"
@@ -45,6 +46,7 @@ import (
 )
 
 var cfgFile string
+var replacer = strings.NewReplacer("-", "_", ".", "_")
 
 // rootCmd represents the base command when called without any subcommands
 var rootCmd = &cobra.Command{
@@ -119,8 +121,6 @@ func Execute(ver, commit, buildDate string) error {
 		"service-token-name", "", "The Service Token name for authenticating.")
 	rootCmd.PersistentFlags().StringVar(&cfg.ServiceToken,
 		"service-token", "", "Service Token for authenticating.")
-	_ = rootCmd.PersistentFlags().MarkHidden("service-token-name")
-	_ = rootCmd.PersistentFlags().MarkHidden("service-token")
 
 	// We don't want to show the default value
 	rootCmd.PersistentFlags().Lookup("api-token").DefValue = ""
@@ -184,6 +184,7 @@ func initConfig() {
 	}
 
 	viper.SetEnvPrefix("planetscale")
+	viper.SetEnvKeyReplacer(replacer)
 	viper.AutomaticEnv() // read in environment variables that match
 
 	// If a config file is found, read it in.