refactor(internal): pass context through to command.Run (#3192)

`command.Run` now takes a `context.Context`, so callers can cancel work
in subprocesses. This makes it so that in parallel flows subprocesses
will stop as soon as the context is cancelled, instead of continuing to
run in the background.

Author: julieqiu

internal/command/command.go

@@ -16,14 +16,15 @@
 package command
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"os/exec"
 )
 
 // Run executes a program (with arguments) and captures any error output.
-func Run(command string, arg ...string) error {
-	cmd := exec.Command(command, arg...)
+func Run(ctx context.Context, command string, arg ...string) error {
+	cmd := exec.CommandContext(ctx, command, arg...)
 	fmt.Fprintf(os.Stderr, "Running: %s\n", cmd.String())
 	if output, err := cmd.CombinedOutput(); err != nil {
 		return fmt.Errorf("%v: %v\n%s", cmd, err, output)

internal/command/command_test.go

@@ -20,13 +20,13 @@ import (
 )
 
 func TestRun(t *testing.T) {
-	if err := Run("go", "version"); err != nil {
+	if err := Run(t.Context(), "go", "version"); err != nil {
 		t.Fatal(err)
 	}
 }
 
 func TestRunError(t *testing.T) {
-	err := Run("go", "invalid-subcommand-bad-bad-bad")
+	err := Run(t.Context(), "go", "invalid-subcommand-bad-bad-bad")
 	if err == nil {
 		t.Fatal("expected error, got nil")
 	}

internal/librarian/generate.go

@@ -79,7 +79,7 @@ func runGenerate(ctx context.Context, all bool, libraryName string) error {
 	if err != nil {
 		return err
 	}
-	return formatLibrary(cfg.Language, lib)
+	return formatLibrary(ctx, cfg.Language, lib)
 }
 
 func generateAll(ctx context.Context, cfg *config.Config) error {
@@ -98,7 +98,7 @@ func generateAll(ctx context.Context, cfg *config.Config) error {
 		if err != nil {
 			return err
 		}
-		if err := formatLibrary(cfg.Language, lib); err != nil {
+		if err := formatLibrary(ctx, cfg.Language, lib); err != nil {
 			return err
 		}
 	}
@@ -253,12 +253,12 @@ func generate(ctx context.Context, language string, library *config.Library, sou
 	return library, nil
 }
 
-func formatLibrary(language string, library *config.Library) error {
+func formatLibrary(ctx context.Context, language string, library *config.Library) error {
 	switch language {
 	case "testhelper":
 		return nil
 	case "rust":
-		return rust.Format(library)
+		return rust.Format(ctx, library)
 	}
 	return fmt.Errorf("format not implemented for %q", language)
 }

internal/librarian/internal/rust/generate.go

@@ -41,7 +41,7 @@ func Generate(ctx context.Context, library *config.Library, sources *config.Sour
 		return err
 	}
 	if library.Veneer {
-		return generateVeneer(library, googleapisDir)
+		return generateVeneer(ctx, library, googleapisDir)
 	}
 
 	if len(library.Channels) != 1 {
@@ -57,7 +57,7 @@ func Generate(ctx context.Context, library *config.Library, sources *config.Sour
 	if err != nil {
 		return err
 	}
-	if err := sidekickrust.Generate(model, library.Output, sidekickConfig); err != nil {
+	if err := sidekickrust.Generate(ctx, model, library.Output, sidekickConfig); err != nil {
 		return err
 	}
 	return nil
@@ -66,17 +66,17 @@ func Generate(ctx context.Context, library *config.Library, sources *config.Sour
 // Format formats a generated Rust library. Must be called sequentially;
 // parallel calls cause race conditions as cargo fmt runs cargo metadata,
 // which competes for locks on the workspace Cargo.toml and Cargo.lock.
-func Format(library *config.Library) error {
-	if err := command.Run("taplo", "fmt", filepath.Join(library.Output, "Cargo.toml")); err != nil {
+func Format(ctx context.Context, library *config.Library) error {
+	if err := command.Run(ctx, "taplo", "fmt", filepath.Join(library.Output, "Cargo.toml")); err != nil {
 		return err
 	}
-	if err := command.Run("cargo", "fmt", "-p", library.Name); err != nil {
+	if err := command.Run(ctx, "cargo", "fmt", "-p", library.Name); err != nil {
 		return err
 	}
 	return nil
 }
 
-func generateVeneer(library *config.Library, googleapisDir string) error {
+func generateVeneer(ctx context.Context, library *config.Library, googleapisDir string) error {
 	if library.Rust == nil || len(library.Rust.Modules) == 0 {
 		return fmt.Errorf("veneer %q has no modules defined", library.Name)
 	}
@@ -88,9 +88,9 @@ func generateVeneer(library *config.Library, googleapisDir string) error {
 		}
 		switch sidekickConfig.General.Language {
 		case "rust":
-			err = sidekickrust.Generate(model, module.Output, sidekickConfig)
+			err = sidekickrust.Generate(ctx, model, module.Output, sidekickConfig)
 		case "rust+prost":
-			err = rust_prost.Generate(model, module.Output, sidekickConfig)
+			err = rust_prost.Generate(ctx, model, module.Output, sidekickConfig)
 		default:
 			err = fmt.Errorf("unknown language: %s", sidekickConfig.General.Language)
 		}

internal/librarian/internal/rust/generate_test.go

@@ -180,7 +180,7 @@ func TestGenerate(t *testing.T) {
 	if err := Generate(t.Context(), library, sources); err != nil {
 		t.Fatal(err)
 	}
-	if err := Format(library); err != nil {
+	if err := Format(t.Context(), library); err != nil {
 		t.Fatal(err)
 	}
 

internal/sidekick/codec_sample/generate.go

@@ -16,6 +16,7 @@
 package codec_sample
 
 import (
+	"context"
 	"embed"
 
 	"github.com/googleapis/librarian/internal/sidekick/api"
@@ -27,7 +28,7 @@ import (
 var templates embed.FS
 
 // Generate generates code from the model.
-func Generate(model *api.API, outdir string, cfg *config.Config) error {
+func Generate(ctx context.Context, model *api.API, outdir string, cfg *config.Config) error {
 	// A template provide converts a template name into the contents.
 	provider := func(name string) (string, error) {
 		contents, err := templates.ReadFile(name)

internal/sidekick/codec_sample/generate_test.go

@@ -57,7 +57,7 @@ func TestFromProtobuf(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if err := Generate(model, outDir, cfg); err != nil {
+	if err := Generate(t.Context(), model, outDir, cfg); err != nil {
 		t.Fatal(err)
 	}
 	filename := path.Join(outDir, "README.md")

internal/sidekick/dart/dart.go

@@ -16,6 +16,7 @@
 package dart
 
 import (
+	"context"
 	"fmt"
 	"regexp"
 	"strings"
@@ -256,8 +257,8 @@ func shouldGenerateMethod(m *api.Method) bool {
 	return m.PathInfo.Bindings[0].PathTemplate != nil
 }
 
-func formatDirectory(dir string) error {
-	if err := command.Run("dart", "format", dir); err != nil {
+func formatDirectory(ctx context.Context, dir string) error {
+	if err := command.Run(ctx, "dart", "format", dir); err != nil {
 		return fmt.Errorf("got an error trying to run `dart format`; perhaps try https://dart.dev/get-dart (%w)", err)
 	}
 	return nil

internal/sidekick/dart/generate.go

@@ -15,6 +15,7 @@
 package dart
 
 import (
+	"context"
 	"embed"
 	"path/filepath"
 
@@ -27,7 +28,7 @@ import (
 var dartTemplates embed.FS
 
 // Generate generates Dart code from the model.
-func Generate(model *api.API, outdir string, config *config.Config) error {
+func Generate(ctx context.Context, model *api.API, outdir string, config *config.Config) error {
 	annotate := newAnnotateModel(model)
 	if err := annotate.annotateModel(config.Codec); err != nil {
 		return err
@@ -39,7 +40,7 @@ func Generate(model *api.API, outdir string, config *config.Config) error {
 		// Check if we're configured to skip formatting.
 		skipFormat := config.Codec["skip-format"]
 		if skipFormat != "true" {
-			err = formatDirectory(outdir)
+			err = formatDirectory(ctx, outdir)
 		}
 	}
 	return err

internal/sidekick/dart/generate_test.go

@@ -65,7 +65,7 @@ func TestFromProtobuf(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if err := Generate(model, outDir, cfg); err != nil {
+	if err := Generate(t.Context(), model, outDir, cfg); err != nil {
 		t.Fatal(err)
 	}
 	for _, expected := range []string{"pubspec.yaml", "lib/secretmanager.dart", "README.md"} {