snapshot_upload.go

package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path"

	"cloud.google.com/go/storage"
	"github.com/sourcegraph/conc/pool"
	"github.com/sourcegraph/sourcegraph/lib/errors"
	"github.com/sourcegraph/sourcegraph/lib/output"
	"google.golang.org/api/option"

	"github.com/sourcegraph/src-cli/internal/pgdump"
)

const srcSnapshotDir = "./src-snapshot"

var srcSnapshotSummaryPath = path.Join(srcSnapshotDir, "summary.json")

// https://pkg.go.dev/cloud.google.com/go/storage#section-readme

func init() {
	usage := `'src snapshot upload' uploads instance snapshot contents generated by 'src snapshot databases' and 'src snapshot summary' to the designated bucket.

USAGE
	src snapshot upload -bucket=$BUCKET -credentials=$CREDENTIALS_FILE

BUCKET
	In general, a Google Cloud Storage bucket and relevant credentials will be provided by Sourcegraph when using this functionality to share a snapshot with Sourcegraph.
`
	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 service account")
	trimExtensions := flagSet.Bool("trim-extensions", true, "trim EXTENSION statements from database dumps for import to Google Cloud SQL")

	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
				trimExtensions bool
			}
			var (
				uploads      []upload             // index aligned with progressBars
				progressBars []output.ProgressBar // index aligned with uploads
			)

			// Open snapshot summary
			if f, err := os.Open(srcSnapshotSummaryPath); 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,
					trimExtensions: false, // not a database dump
				})
				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,
						trimExtensions: *trimExtensions,
					})
					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 := pool.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 := copyDumpToBucket(ctx, u.file, u.stat, bucket, progressFn, u.trimExtensions); 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 copyDumpToBucket(ctx context.Context, src io.ReadSeeker, stat fs.FileInfo, dst *storage.BucketHandle, progressFn func(int64), trimExtensions bool) error {
	// Set up object to write to
	object := dst.Object(stat.Name()).NewWriter(ctx)
	object.ProgressFunc = progressFn
	defer object.Close()

	// To assert against actual file size
	var totalWritten int64

	// Do a partial copy, that filters out incompatible statements
	if trimExtensions {
		written, err := pgdump.FilterInvalidLines(object, src, progressFn)
		if err != nil {
			return errors.Wrap(err, "filter out incompatible statements and upload")
		}
		totalWritten += written
	}

	// io.Copy is the best way to copy from a reader to writer in Go,
	// storage.Writer has its own chunking mechanisms internally.
	// io.Reader is stateful, so this copy will just continue from where FilterInvalidLines left off, if used
	written, err := io.Copy(object, src)
	if err != nil {
		return errors.Wrap(err, "upload")
	}
	totalWritten += written

	// Progress is not called on completion of io.Copy,
	// so we call it manually after to update our pretty progress bars.
	progressFn(written)

	// Validate we have sent all data.
	// FilterInvalidLines may add some bytes, so the check is not a strict equality.
	size := stat.Size()
	if totalWritten < size {
		return errors.Newf("expected to write %d bytes, but actually wrote %d bytes (diff: %d bytes)",
			size, totalWritten, totalWritten-size)
	}

	return nil
}