From 495079f0bbd2c8460297a6c655457dda5d349daf Mon Sep 17 00:00:00 2001
From: Matt Robenolt <matt@ydekproductions.com>
Date: Mon, 27 Nov 2023 14:49:02 -0800
Subject: [PATCH] database/restore: port to using psdbproxy

---
 internal/cmd/database/restore.go | 97 ++++++++++++++++++++------------
 1 file changed, 62 insertions(+), 35 deletions(-)

diff --git a/internal/cmd/database/restore.go b/internal/cmd/database/restore.go
index 75066b51..48739284 100644
--- a/internal/cmd/database/restore.go
+++ b/internal/cmd/database/restore.go
@@ -4,23 +4,25 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net"
 	"time"
 
 	"github.com/planetscale/cli/internal/cmdutil"
 	"github.com/planetscale/cli/internal/dumper"
+	"github.com/planetscale/cli/internal/passwordutil"
 	"github.com/planetscale/cli/internal/printer"
 	"github.com/planetscale/cli/internal/proxyutil"
 	ps "github.com/planetscale/planetscale-go/planetscale"
-	"github.com/planetscale/sql-proxy/proxy"
 
 	"github.com/spf13/cobra"
 )
 
 type restoreFlags struct {
-	localAddr string
-	dir       string
-	overwrite bool
-	threads   int
+	localAddr  string
+	remoteAddr string
+	dir        string
+	overwrite  bool
+	threads    int
 }
 
 // RestoreCmd encapsulates the commands for restore a database
@@ -35,6 +37,8 @@ func RestoreCmd(ch *cmdutil.Helper) *cobra.Command {
 
 	cmd.PersistentFlags().StringVar(&f.localAddr, "local-addr",
 		"", "Local address to bind and listen for connections. By default the proxy binds to 127.0.0.1 with a random port.")
+	cmd.PersistentFlags().StringVar(&f.remoteAddr, "remote-addr", "",
+		"PlanetScale Database remote network address. By default the remote address is populated automatically from the PlanetScale API. (format: `hostname:port`)")
 	cmd.PersistentFlags().StringVar(&f.dir, "dir", "",
 		"Directory containing the files to be used for the restore (required)")
 	cmd.PersistentFlags().BoolVar(&f.overwrite, "overwrite-tables", false, "If true, will attempt to DROP TABLE before restoring.")
@@ -59,31 +63,6 @@ func restore(ch *cmdutil.Helper, cmd *cobra.Command, flags *restoreFlags, args [
 		return err
 	}
 
-	const localProxyAddr = "127.0.0.1"
-	localAddr := localProxyAddr + ":0"
-	if flags.localAddr != "" {
-		localAddr = flags.localAddr
-	}
-
-	proxyOpts := proxy.Options{
-		CertSource: proxyutil.NewRemoteCertSource(client, cmdutil.AdministratorRole),
-		LocalAddr:  localAddr,
-		Instance:   fmt.Sprintf("%s/%s/%s", ch.Config.Organization, database, branch),
-		Logger:     cmdutil.NewZapLogger(ch.Debug()),
-	}
-
-	p, err := proxy.NewClient(proxyOpts)
-	if err != nil {
-		return fmt.Errorf("couldn't create proxy client: %s", err)
-	}
-
-	go func() {
-		err := p.Run(ctx)
-		if err != nil {
-			ch.Printer.Println("proxy error: ", err)
-		}
-	}()
-
 	dbBranch, err := client.DatabaseBranches.Get(ctx, &ps.GetDatabaseBranchRequest{
 		Organization: ch.Config.Organization,
 		Database:     database,
@@ -103,16 +82,64 @@ func restore(ch *cmdutil.Helper, cmd *cobra.Command, flags *restoreFlags, args [
 		return errors.New("database branch is not ready yet, please try again in a few minutes")
 	}
 
-	addr, err := p.LocalAddr()
+	pw, err := passwordutil.New(ctx, client, passwordutil.Options{
+		Organization: ch.Config.Organization,
+		Database:     database,
+		Branch:       branch,
+		Role:         cmdutil.AdministratorRole,
+		Name:         passwordutil.GenerateName("pscale-cli-restore"),
+		TTL:          6 * time.Hour, // TODO: use shorter TTL, but implement refreshing
+	})
 	if err != nil {
-		return err
+		return cmdutil.HandleError(err)
 	}
+	defer func() {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		if err := pw.Cleanup(ctx); err != nil {
+			ch.Printer.Println("failed to delete credentials: ", err)
+		}
+	}()
+
+	localAddr := "127.0.0.1:0"
+	if flags.localAddr != "" {
+		localAddr = flags.localAddr
+	}
+
+	remoteAddr := flags.remoteAddr
+	if remoteAddr == "" {
+		remoteAddr = pw.Password.Hostname
+	}
+
+	proxy := proxyutil.New(proxyutil.Config{
+		Logger:       cmdutil.NewZapLogger(ch.Debug()),
+		UpstreamAddr: remoteAddr,
+		Username:     pw.Password.Username,
+		Password:     pw.Password.PlainText,
+	})
+	defer proxy.Close()
+
+	l, err := net.Listen("tcp", localAddr)
+	if err != nil {
+		return cmdutil.HandleError(err)
+	}
+	defer l.Close()
+
+	go func() {
+		if err := proxy.Serve(l); err != nil {
+			ch.Printer.Println("proxy error: ", err)
+		}
+	}()
+
+	addr := l.Addr()
 
 	cfg := dumper.NewDefaultConfig()
 	cfg.Threads = flags.threads
-	cfg.User = "root"
-	// NOTE: the password is a placeholder, replace once we get rid of the proxy
-	cfg.Password = "root"
+	// NOTE(mattrobenolt): credentials are needed even though they aren't used,
+	// otherwise, dumper will complain.
+	cfg.User = "nobody"
+	cfg.Password = "nobody"
 	cfg.Address = addr.String()
 	cfg.Debug = ch.Debug()
 	cfg.IntervalMs = 10 * 1000