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