sg: Add migration add-log command (#30863)

cmd/migrator/main.go

@@ -56,6 +56,7 @@ func mainErr(ctx context.Context, args []string) error {
 			cliutil.UpTo(appName, runnerFactory, out),
 			cliutil.DownTo(appName, runnerFactory, out),
 			cliutil.Validate(appName, runnerFactory, out),
+			cliutil.AddLog(appName, runnerFactory, out),
 		},
 	}
 
@@ -84,6 +85,6 @@ func newRunnerFactory() func(ctx context.Context, schemaNames []string) (cliutil
 			return connections.NewStoreShim(store.NewWithDB(db, migrationsTable, operations))
 		}
 
-		return connections.RunnerFromDSNs(dsns, appName, storeFactory), nil
+		return cliutil.NewShim(connections.RunnerFromDSNs(dsns, appName, storeFactory)), nil
 	}
 }

dev/sg/sg_migration.go

@@ -47,6 +47,7 @@ var (
 	UndoCommand     = cliutil.Undo("sg migration", makeRunner, stdout.Out)
 	downToCommand   = cliutil.DownTo("sg migration", makeRunner, stdout.Out)
 	validateCommand = cliutil.Validate("sg validate", makeRunner, stdout.Out)
+	addLogCommand   = cliutil.AddLog("sg migration", makeRunner, stdout.Out)
 
 	leavesFlagSet = flag.NewFlagSet("sg migration leaves", flag.ExitOnError)
 	leavesCommand = &ffcli.Command{
@@ -86,6 +87,7 @@ var (
 			UndoCommand,
 			downToCommand,
 			validateCommand,
+			addLogCommand,
 			leavesCommand,
 			squashCommand,
 		},
@@ -97,7 +99,7 @@ func makeRunner(ctx context.Context, schemaNames []string) (cliutil.Runner, erro
 		return connections.NewStoreShim(store.NewWithDB(db, migrationsTable, store.NewOperations(&observation.TestContext)))
 	}
 
-	return connections.RunnerFromDSNs(postgresdsn.RawDSNsBySchema(schemaNames), "sg", storeFactory), nil
+	return cliutil.NewShim(connections.RunnerFromDSNs(postgresdsn.RawDSNsBySchema(schemaNames), "sg", storeFactory)), nil
 }
 
 func addExec(ctx context.Context, args []string) error {

internal/database/migration/cliutil/addlog.go

@@ -0,0 +1,68 @@
+package cliutil
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"strconv"
+
+	"github.com/peterbourgon/ff/v3/ffcli"
+
+	"github.com/sourcegraph/sourcegraph/internal/database/migration/definition"
+	"github.com/sourcegraph/sourcegraph/lib/output"
+)
+
+func AddLog(commandName string, factory RunnerFactory, out *output.Output) *ffcli.Command {
+	var (
+		flagSet        = flag.NewFlagSet(fmt.Sprintf("%s add-log", commandName), flag.ExitOnError)
+		schemaNameFlag = flagSet.String("db", "", `The target schema to modify.`)
+		versionFlag    = flagSet.String("version", "", "The migration version.")
+		upFlag         = flagSet.Bool("up", true, "The migration direction.")
+	)
+
+	exec := func(ctx context.Context, args []string) error {
+		if len(args) != 0 {
+			out.WriteLine(output.Linef("", output.StyleWarning, "ERROR: too many arguments"))
+			return flag.ErrHelp
+		}
+
+		if *schemaNameFlag == "" {
+			out.WriteLine(output.Linef("", output.StyleWarning, "ERROR: supply a schema via -db"))
+			return flag.ErrHelp
+		}
+
+		if *versionFlag == "" {
+			out.WriteLine(output.Linef("", output.StyleWarning, "ERROR: supply a migration version via -version"))
+			return flag.ErrHelp
+		}
+		version, err := strconv.Atoi(*versionFlag)
+		if err != nil {
+			out.WriteLine(output.Linef("", output.StyleWarning, "ERROR: invalid migration version %q", *versionFlag))
+			return flag.ErrHelp
+		}
+
+		r, err := factory(ctx, []string{*schemaNameFlag})
+		if err != nil {
+			return err
+		}
+		store, err := r.Store(ctx, *schemaNameFlag)
+		if err != nil {
+			return err
+		}
+
+		return store.WithMigrationLog(ctx, definition.Definition{ID: version}, *upFlag, noop)
+	}
+
+	return &ffcli.Command{
+		Name:       "add-log",
+		ShortUsage: fmt.Sprintf("%s add-log -db=<schema> -version=<version> [-up=true|false]", commandName),
+		ShortHelp:  "Add an entry to the migration log",
+		FlagSet:    flagSet,
+		Exec:       exec,
+		LongHelp:   ConstructLongHelp(),
+	}
+}
+
+func noop() error {
+	return nil
+}

internal/database/migration/cliutil/iface.go

@@ -3,12 +3,30 @@ package cliutil
 import (
 	"context"
 
+	"github.com/sourcegraph/sourcegraph/internal/database/migration/definition"
 	"github.com/sourcegraph/sourcegraph/internal/database/migration/runner"
 )
 
 type Runner interface {
 	Run(ctx context.Context, options runner.Options) error
 	Validate(ctx context.Context, schemaNames ...string) error
+	Store(ctx context.Context, schemaName string) (Store, error)
+}
+
+type Store interface {
+	WithMigrationLog(ctx context.Context, definition definition.Definition, up bool, f func() error) error
 }
 
 type RunnerFactory func(ctx context.Context, schemaNames []string) (Runner, error)
+
+type runnerShim struct {
+	*runner.Runner
+}
+
+func NewShim(runner *runner.Runner) Runner {
+	return &runnerShim{Runner: runner}
+}
+
+func (r *runnerShim) Store(ctx context.Context, schemaName string) (Store, error) {
+	return r.Runner.Store(ctx, schemaName)
+}

internal/database/migration/runner/runner.go

@@ -36,6 +36,16 @@ type schemaVersion struct {
 
 type visitFunc func(ctx context.Context, schemaContext schemaContext) error
 
+// Store returns the store associated with the given schema.
+func (r *Runner) Store(ctx context.Context, schemaName string) (Store, error) {
+	if factory, ok := r.storeFactories[schemaName]; ok {
+		return factory(ctx)
+
+	}
+
+	return nil, errors.Newf("unknown store %q", schemaName)
+}
+
 // forEachSchema invokes the given function once for each schema in the given list, with
 // store instances initialized for each given schema name. Each function invocation occurs
 // concurrently. Errors from each invocation are collected and returned. An error from one