snapshot: add pg_dump command builder with 'snapshot databases' (#889)

Add simple utility for generating commands that can then be run/modified to create database dumps.

Right now, it's mostly for reference - as we build experience with different setups we can see if this can be made more flexible in the right ways. I believe a templating command like this is important to have, however, since we want to make sure all databases are exported and the correct flags are used with pg_dump, and documentation + manual copy-paste is error-prone.

Co-authored-by: Michael Lin <[email protected]>

Author: bobheadxi

cmd/src/snapshot_databases.go

@@ -0,0 +1,173 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"os"
+	"strings"
+
+	"github.com/sourcegraph/sourcegraph/lib/errors"
+	"github.com/sourcegraph/sourcegraph/lib/output"
+	"gopkg.in/yaml.v3"
+
+	"github.com/sourcegraph/src-cli/internal/pgdump"
+)
+
+func init() {
+	usage := `'src snapshot databases' generates commands to export Sourcegraph database dumps.
+Note that these commands are intended for use as reference - you may need to adjust the commands for your deployment.
+
+USAGE
+	src [-v] snapshot databases <pg_dump|docker|kubectl> [--targets=<docker|k8s|"targets.yaml">]
+
+TARGETS FILES
+	Predefined targets are available based on default Sourcegraph configurations ('docker', 'k8s').
+	Custom targets configuration can be provided in YAML format with '--targets=target.yaml', e.g.
+
+		primary:
+			target: ...   # the DSN of the database deployment, e.g. in docker, the name of the database container
+			dbname: ...   # name of database
+			username: ... # username for database access
+			password: ... # password for database access - only include password if it is non-sensitive
+		codeintel:
+			# same as above
+		codeinsights:
+			# same as above
+
+	See the pgdump.Targets type for more details.
+`
+	flagSet := flag.NewFlagSet("databases", flag.ExitOnError)
+	targetsKeyFlag := flagSet.String("targets", "auto", "predefined targets ('docker' or 'k8s'), or a custom targets.yaml file")
+
+	snapshotCommands = append(snapshotCommands, &command{
+		flagSet: flagSet,
+		handler: func(args []string) error {
+			if err := flagSet.Parse(args); err != nil {
+				return err
+			}
+			out := output.NewOutput(flagSet.Output(), output.OutputOpts{Verbose: *verbose})
+
+			var builder string
+			if len(args) > 0 {
+				builder = args[0]
+			}
+
+			targetKey := "docker"
+			var commandBuilder pgdump.CommandBuilder
+			switch builder {
+			case "pg_dump", "":
+				targetKey = "local"
+				commandBuilder = func(t pgdump.Target) (string, error) {
+					cmd := pgdump.Command(t)
+					if t.Target != "" {
+						return fmt.Sprintf("%s --host=%s", cmd, t.Target), nil
+					}
+					return cmd, nil
+				}
+			case "docker":
+				commandBuilder = func(t pgdump.Target) (string, error) {
+					return fmt.Sprintf("docker exec -it %s sh -c '%s'", t.Target, pgdump.Command(t)), nil
+				}
+			case "kubectl":
+				targetKey = "k8s"
+				commandBuilder = func(t pgdump.Target) (string, error) {
+					return fmt.Sprintf("kubectl exec -it %s -- bash -c '%s'", t.Target, pgdump.Command(t)), nil
+				}
+			default:
+				return errors.Newf("unknown or invalid template type %q", builder)
+			}
+			if *targetsKeyFlag != "auto" {
+				targetKey = *targetsKeyFlag
+			}
+
+			targets, ok := predefinedDatabaseDumpTargets[targetKey]
+			if !ok {
+				out.WriteLine(output.Emojif(output.EmojiInfo, "Using targets defined in targets file %q", targetKey))
+				f, err := os.Open(targetKey)
+				if err != nil {
+					return errors.Wrapf(err, "invalid targets file %q", targetKey)
+				}
+				if err := yaml.NewDecoder(f).Decode(&targets); err != nil {
+					return errors.Wrapf(err, "invalid targets file %q", targetKey)
+				}
+			} else {
+				out.WriteLine(output.Emojif(output.EmojiInfo, "Using predefined targets for %s environments", targetKey))
+			}
+
+			commands, err := pgdump.BuildCommands(commandBuilder, targets)
+			if err != nil {
+				return errors.Wrap(err, "failed to build commands")
+			}
+
+			b := out.Block(output.Emoji(output.EmojiSuccess, "Commands generated - run them all to generate required database dumps:"))
+			b.Write("\n" + strings.Join(commands, "\n"))
+			b.Close()
+
+			out.WriteLine(output.Styledf(output.StyleSuggestion, "Note that you may need to do some additional setup, such as authentication, beforehand."))
+
+			return nil
+		},
+		usageFunc: func() { fmt.Fprint(flag.CommandLine.Output(), usage) },
+	})
+}
+
+// predefinedDatabaseDumpTargets is based on default Sourcegraph configurations.
+var predefinedDatabaseDumpTargets = map[string]pgdump.Targets{
+	"local": {
+		Primary: pgdump.Target{
+			DBName:   "sg",
+			Username: "sg",
+			Password: "sg",
+		},
+		CodeIntel: pgdump.Target{
+			DBName:   "sg",
+			Username: "sg",
+			Password: "sg",
+		},
+		CodeInsights: pgdump.Target{
+			DBName:   "postgres",
+			Username: "postgres",
+			Password: "password",
+		},
+	},
+	"docker": { // based on deploy-sourcegraph-managed
+		Primary: pgdump.Target{
+			Target:   "pgsql",
+			DBName:   "sg",
+			Username: "sg",
+			Password: "sg",
+		},
+		CodeIntel: pgdump.Target{
+			Target:   "codeintel-db",
+			DBName:   "sg",
+			Username: "sg",
+			Password: "sg",
+		},
+		CodeInsights: pgdump.Target{
+			Target:   "codeinsights-db",
+			DBName:   "postgres",
+			Username: "postgres",
+			Password: "password",
+		},
+	},
+	"k8s": { // based on deploy-sourcegraph-helm
+		Primary: pgdump.Target{
+			Target:   "statefulset/pgsql",
+			DBName:   "sg",
+			Username: "sg",
+			Password: "sg",
+		},
+		CodeIntel: pgdump.Target{
+			Target:   "statefulset/codeintel-db",
+			DBName:   "sg",
+			Username: "sg",
+			Password: "sg",
+		},
+		CodeInsights: pgdump.Target{
+			Target:   "statefulset/codeinsights-db",
+			DBName:   "postgres",
+			Username: "postgres",
+			Password: "password",
+		},
+	},
+}

internal/pgdump/pgdump.go

@@ -0,0 +1,68 @@
+package pgdump
+
+import (
+	"fmt"
+
+	"github.com/sourcegraph/sourcegraph/lib/errors"
+)
+
+// Targets represents configuration for each of Sourcegraph's databases.
+type Targets struct {
+	Primary      Target `yaml:"primary"`
+	CodeIntel    Target `yaml:"codeintel"`
+	CodeInsights Target `yaml:"codeinsights"`
+}
+
+// Target represents a database for pg_dump to export.
+type Target struct {
+	// Target is the DSN of the database deployment:
+	//
+	// - in docker, the name of the database container, e.g. pgsql, codeintel-db, codeinsights-db
+	// - in k8s, the name of the deployment or statefulset, e.g. deploy/pgsql, sts/pgsql
+	// - in plain pg_dump, the server host or socket directory
+	Target string `yaml:"target"`
+
+	DBName   string `yaml:"dbname"`
+	Username string `yaml:"username"`
+
+	// Only include password if non-sensitive
+	Password string `yaml:"password"`
+}
+
+// Command generates a pg_dump command that can be used for on-prem-to-Cloud migrations.
+func Command(t Target) string {
+	dump := fmt.Sprintf("pg_dump --no-owner --format=p --no-acl --username=%s --dbname=%s",
+		t.Username, t.DBName)
+	if t.Password == "" {
+		return dump
+	}
+	return fmt.Sprintf("PGPASSWORD=%s %s", t.Password, dump)
+}
+
+type CommandBuilder func(Target) (string, error)
+
+// BuildCommands generates commands that output Postgres dumps and sends them to predefined
+// files for each target database.
+func BuildCommands(commandBuilder CommandBuilder, targets Targets) ([]string, error) {
+	var commands []string
+	for _, t := range []struct {
+		Output string
+		Target Target
+	}{{
+		Output: "primary.sql",
+		Target: targets.Primary,
+	}, {
+		Output: "codeintel.sql",
+		Target: targets.CodeIntel,
+	}, {
+		Output: "codeinsights.sql",
+		Target: targets.CodeInsights,
+	}} {
+		c, err := commandBuilder(t.Target)
+		if err != nil {
+			return nil, errors.Wrapf(err, "generating command for %q", t.Output)
+		}
+		commands = append(commands, fmt.Sprintf("%s > %s", c, t.Output))
+	}
+	return commands, nil
+}