Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

database/restore: port to using psdbproxy #761

Merged
merged 1 commit into from
Dec 6, 2023
Merged

database/restore: port to using psdbproxy #761

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 62 additions & 35 deletions internal/cmd/database/restore.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.")
Expand All @@ -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,
Expand All @@ -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
Expand Down