snapshot: add snapshot contents uploader

Author: bobheadxi

.gitignore

@@ -6,3 +6,4 @@ release
 .idea
 .env
 .envrc
+src-snapshot/

cmd/src/snapshot_databases.go

@@ -94,7 +94,7 @@ TARGETS FILES
 				out.WriteLine(output.Emojif(output.EmojiInfo, "Using predefined targets for %s environments", targetKey))
 			}
 
-			commands, err := pgdump.BuildCommands(commandBuilder, targets)
+			commands, err := pgdump.BuildCommands(srcSnapshotDir, commandBuilder, targets)
 			if err != nil {
 				return errors.Wrap(err, "failed to build commands")
 			}

cmd/src/snapshot_summary.go

@@ -18,10 +18,9 @@ func init() {
 
 USAGE
 	src login # site-admin authentication required
-	src [-v] snapshot summary [-summary-path="./src-snapshot-summary.json"]
+	src [-v] snapshot summary
 `
 	flagSet := flag.NewFlagSet("summary", flag.ExitOnError)
-	snapshotPath := flagSet.String("summary-path", "./src-snapshot-summary.json", "path to write snapshot summary to")
 	apiFlags := api.NewFlags(flagSet)
 
 	snapshotCommands = append(snapshotCommands, &command{
@@ -39,7 +38,7 @@ USAGE
 				return err
 			}
 
-			f, err := os.OpenFile(*snapshotPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, os.ModePerm)
+			f, err := os.OpenFile(snapshotSummaryPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, os.ModePerm)
 			if err != nil {
 				return errors.Wrap(err, "open snapshot file")
 			}

cmd/src/snapshot_upload.go

@@ -0,0 +1,158 @@
+package main
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"io"
+	"io/fs"
+	"os"
+	"path"
+
+	"cloud.google.com/go/storage"
+	"github.com/sourcegraph/sourcegraph/lib/errors"
+	"github.com/sourcegraph/sourcegraph/lib/group"
+	"github.com/sourcegraph/sourcegraph/lib/output"
+	"google.golang.org/api/option"
+
+	"github.com/sourcegraph/src-cli/internal/pgdump"
+)
+
+const srcSnapshotDir = "./src-snapshot"
+
+var snapshotSummaryPath = path.Join(srcSnapshotDir, "summary.json")
+
+// https://pkg.go.dev/cloud.google.com/go/storage#section-readme
+
+func init() {
+	usage := `'src snapshot upload' uploads snapshot contents.
+
+USAGE
+	src
+`
+	flagSet := flag.NewFlagSet("upload", flag.ExitOnError)
+	bucketName := flagSet.String("bucket", "", "destination Cloud Storage bucket name")
+	credentialsPath := flagSet.String("credentials", "", "JSON credentials file for Google Cloud")
+
+	snapshotCommands = append(snapshotCommands, &command{
+		flagSet: flagSet,
+		handler: func(args []string) error {
+			if err := flagSet.Parse(args); err != nil {
+				return err
+			}
+
+			if *bucketName == "" {
+				return errors.New("-bucket required")
+			}
+			if *credentialsPath == "" {
+				return errors.New("-credentials required")
+			}
+
+			out := output.NewOutput(flagSet.Output(), output.OutputOpts{Verbose: *verbose})
+			ctx := context.Background()
+			c, err := storage.NewClient(ctx, option.WithCredentialsFile(*credentialsPath))
+			if err != nil {
+				return errors.Wrap(err, "create Cloud Storage client")
+			}
+
+			type upload struct {
+				file *os.File
+				stat os.FileInfo
+			}
+			var (
+				uploads      []upload             // index aligned with progressBars
+				progressBars []output.ProgressBar // index aligned with uploads
+			)
+
+			// Open snapshot summary
+			if f, err := os.Open(snapshotSummaryPath); err != nil {
+				return errors.Wrap(err, "failed to open snapshot summary - generate one with 'src snapshot summary'")
+			} else {
+				stat, err := f.Stat()
+				if err != nil {
+					return errors.Wrap(err, "get file size")
+				}
+				uploads = append(uploads, upload{
+					file: f,
+					stat: stat,
+				})
+				progressBars = append(progressBars, output.ProgressBar{
+					Label: stat.Name(),
+					Max:   float64(stat.Size()),
+				})
+			}
+
+			// Open database dumps
+			for _, o := range pgdump.Outputs(srcSnapshotDir, pgdump.Targets{}) {
+				if f, err := os.Open(o.Output); err != nil {
+					return errors.Wrap(err, "failed to database dump - generate one with 'src snapshot databases'")
+				} else {
+					stat, err := f.Stat()
+					if err != nil {
+						return errors.Wrap(err, "get file size")
+					}
+					uploads = append(uploads, upload{
+						file: f,
+						stat: stat,
+					})
+					progressBars = append(progressBars, output.ProgressBar{
+						Label: stat.Name(),
+						Max:   float64(stat.Size()),
+					})
+				}
+			}
+
+			// Start uploads
+			progress := out.Progress(progressBars, nil)
+			progress.WriteLine(output.Emoji(output.EmojiHourglass, "Starting uploads..."))
+			bucket := c.Bucket(*bucketName)
+			g := group.New().WithErrors().WithContext(ctx)
+			for i, u := range uploads {
+				i := i
+				u := u
+				g.Go(func(ctx context.Context) error {
+					progressFn := func(p int64) { progress.SetValue(i, float64(p)) }
+
+					if err := copyToBucket(ctx, u.file, u.stat, bucket, progressFn); err != nil {
+						return errors.Wrap(err, u.stat.Name())
+					}
+
+					return nil
+				})
+			}
+
+			// Finalize
+			errs := g.Wait()
+			progress.Complete()
+			if errs != nil {
+				out.WriteLine(output.Line(output.EmojiFailure, output.StyleFailure, "Some snapshot contents failed to upload."))
+				return errs
+			}
+
+			out.WriteLine(output.Emoji(output.EmojiSuccess, "Summary contents uploaded!"))
+			return nil
+		},
+		usageFunc: func() { fmt.Fprint(flag.CommandLine.Output(), usage) },
+	})
+}
+
+func copyToBucket(ctx context.Context, src io.Reader, stat fs.FileInfo, dst *storage.BucketHandle, progressFn func(int64)) error {
+	writer := dst.Object(stat.Name()).NewWriter(ctx)
+	writer.ProgressFunc = progressFn
+	defer writer.Close()
+
+	// io.Copy is the best way to copy from a reader to writer in Go, and storage.Writer
+	// has its own chunking mechanisms internally.
+	written, err := io.Copy(writer, src)
+	if err != nil {
+		return err
+	}
+
+	// Validate we have sent all data
+	if size := stat.Size(); written != size {
+		return errors.Newf("expected to write %d bytes, but actually wrote %d bytes",
+			size, written)
+	}
+
+	return nil
+}

go.mod

@@ -3,13 +3,14 @@ module github.com/sourcegraph/src-cli
 go 1.18
 
 require (
+	cloud.google.com/go/storage v1.28.0
 	github.com/Masterminds/semver v1.5.0
 	github.com/creack/goselect v0.1.2
 	github.com/derision-test/glock v0.0.0-20210316032053-f5b74334bb29
 	github.com/dineshappavoo/basex v0.0.0-20170425072625-481a6f6dc663
 	github.com/dustin/go-humanize v1.0.0
 	github.com/gobwas/glob v0.2.3
-	github.com/google/go-cmp v0.5.8
+	github.com/google/go-cmp v0.5.9
 	github.com/grafana/regexp v0.0.0-20220304100321-149c8afcd6cb
 	github.com/hexops/autogold v1.3.0
 	github.com/jedib0t/go-pretty/v6 v6.3.7
@@ -22,19 +23,27 @@ require (
 	github.com/sourcegraph/go-diff v0.6.1
 	github.com/sourcegraph/jsonx v0.0.0-20200629203448-1a936bd500cf
 	github.com/sourcegraph/scip v0.2.3
-	github.com/sourcegraph/sourcegraph/lib v0.0.0-20221004162410-237d5855fa13
+	github.com/sourcegraph/sourcegraph/lib v0.0.0-20221128221602-4b48d5d69f5d
 	github.com/stretchr/testify v1.8.0
-	golang.org/x/net v0.0.0-20220722155237-a158d28d115b
-	golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4
-	google.golang.org/protobuf v1.28.0
+	golang.org/x/net v0.2.0
+	golang.org/x/sync v0.1.0
+	google.golang.org/api v0.102.0
+	google.golang.org/protobuf v1.28.1
 	gopkg.in/yaml.v3 v3.0.1
 	jaytaylor.com/html2text v0.0.0-20200412013138-3577fbdbcff7
 )
 
 require (
+	cloud.google.com/go v0.104.0 // indirect
+	cloud.google.com/go/compute v1.12.1 // indirect
+	cloud.google.com/go/compute/metadata v0.2.1 // indirect
+	cloud.google.com/go/iam v0.5.0 // indirect
 	github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 // indirect
 	github.com/Masterminds/sprig v2.15.0+incompatible // indirect
 	github.com/alecthomas/chroma v0.10.0 // indirect
+	github.com/alecthomas/kingpin v2.2.6+incompatible // indirect
+	github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751 // indirect
+	github.com/alecthomas/units v0.0.0-20211218093645-b94a6e3cc137 // indirect
 	github.com/aokoli/goutils v1.0.1 // indirect
 	github.com/aymerick/douceur v0.2.0 // indirect
 	github.com/bufbuild/buf v1.4.0 // indirect
@@ -43,7 +52,9 @@ require (
 	github.com/cockroachdb/logtags v0.0.0-20211118104740-dabe8e521a4f // indirect
 	github.com/cockroachdb/redact v1.1.3 // indirect
 	github.com/cpuguy83/go-md2man/v2 v2.0.2 // indirect
+	github.com/dave/jennifer v1.6.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/derision-test/go-mockgen v1.3.7 // indirect
 	github.com/dlclark/regexp2 v1.4.0 // indirect
 	github.com/envoyproxy/protoc-gen-validate v0.3.0-java // indirect
 	github.com/fatih/color v1.13.0 // indirect
@@ -55,6 +66,8 @@ require (
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
 	github.com/golang/protobuf v1.5.2 // indirect
 	github.com/google/uuid v1.3.0 // indirect
+	github.com/googleapis/enterprise-certificate-proxy v0.2.0 // indirect
+	github.com/googleapis/gax-go/v2 v2.6.0 // indirect
 	github.com/gorilla/css v1.0.0 // indirect
 	github.com/hexops/gotextdiff v1.0.3 // indirect
 	github.com/hexops/valast v1.4.1 // indirect
@@ -72,6 +85,7 @@ require (
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-runewidth v0.0.13 // indirect
 	github.com/microcosm-cc/bluemonday v1.0.17 // indirect
+	github.com/mitchellh/go-wordwrap v1.0.1 // indirect
 	github.com/moby/term v0.0.0-20210619224110-3f7ff695adc6 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
@@ -104,13 +118,16 @@ require (
 	go.uber.org/multierr v1.8.0 // indirect
 	go.uber.org/zap v1.23.0 // indirect
 	golang.org/x/crypto v0.0.0-20210921155107-089bfa567519 // indirect
-	golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4 // indirect
-	golang.org/x/sys v0.0.0-20220928140112-f11e5e49a4ec // indirect
-	golang.org/x/term v0.0.0-20220411215600-e5f449aeb171 // indirect
-	golang.org/x/text v0.3.7 // indirect
-	golang.org/x/tools v0.1.12 // indirect
-	google.golang.org/genproto v0.0.0-20220414192740-2d67ff6cf2b4 // indirect
-	google.golang.org/grpc v1.45.0 // indirect
+	golang.org/x/mod v0.7.0 // indirect
+	golang.org/x/oauth2 v0.0.0-20221014153046-6fdb5e3db783 // indirect
+	golang.org/x/sys v0.2.0 // indirect
+	golang.org/x/term v0.2.0 // indirect
+	golang.org/x/text v0.4.0 // indirect
+	golang.org/x/tools v0.3.0 // indirect
+	golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect
+	google.golang.org/appengine v1.6.7 // indirect
+	google.golang.org/genproto v0.0.0-20221024183307-1bc688fe9f3e // indirect
+	google.golang.org/grpc v1.50.1 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	mvdan.cc/gofumpt v0.2.1 // indirect
 )