fix(librarian): update-image uses the newly detected image (#2830)

Fixes #2822

We update the docker request structs to contain an optional `Image`
field which will override the docker image used to run the command. When
running the generate command, we use this field to ensure we are using
the correct language container.

Author: chingor13

internal/docker/docker.go

@@ -79,6 +79,10 @@ type BuildRequest struct {
 	// State is a pointer to the [config.LibrarianState] struct, representing
 	// the overall state of the generation and release pipeline.
 	State *config.LibrarianState
+
+	// Image is the name of the docker image to use when running. If not
+	// specified, uses the default image configured for the client.
+	Image string
 }
 
 // ConfigureRequest contains all the information required for a language
@@ -106,6 +110,10 @@ type ConfigureRequest struct {
 	// State is a pointer to the [config.LibrarianState] struct, representing
 	// the overall state of the generation and release pipeline.
 	State *config.LibrarianState
+
+	// Image is the name of the docker image to use when running. If not
+	// specified, uses the default image configured for the client.
+	Image string
 }
 
 // GenerateRequest contains all the information required for a language
@@ -127,6 +135,10 @@ type GenerateRequest struct {
 	// State is a pointer to the [config.LibrarianState] struct, representing
 	// the overall state of the generation and release pipeline.
 	State *config.LibrarianState
+
+	// Image is the name of the docker image to use when running. If not
+	// specified, uses the default image configured for the client.
+	Image string
 }
 
 // ReleaseStageRequest contains all the information required for a language
@@ -164,6 +176,10 @@ type ReleaseStageRequest struct {
 	// State is a pointer to the [config.LibrarianState] struct, representing
 	// the overall state of the generation and release pipeline.
 	State *config.LibrarianState
+
+	// Image is the name of the docker image to use when running. If not
+	// specified, uses the default image configured for the client.
+	Image string
 }
 
 // DockerOptions contains optional configuration parameters for invoking
@@ -232,7 +248,9 @@ func (c *Docker) Generate(ctx context.Context, request *GenerateRequest) error {
 		fmt.Sprintf("%s:/output", request.Output),
 		fmt.Sprintf("%s:/source:ro", request.ApiRoot), // readonly volume
 	}
-	return c.runDocker(ctx, CommandGenerate, mounts, commandArgs)
+
+	image := c.resolveImage(request.Image)
+	return c.runDocker(ctx, image, CommandGenerate, mounts, commandArgs)
 }
 
 // Build builds the library with an ID of libraryID, as configured in
@@ -262,7 +280,8 @@ func (c *Docker) Build(ctx context.Context, request *BuildRequest) error {
 		"--repo=/repo",
 	}
 
-	return c.runDocker(ctx, CommandBuild, mounts, commandArgs)
+	image := c.resolveImage(request.Image)
+	return c.runDocker(ctx, image, CommandBuild, mounts, commandArgs)
 }
 
 // Configure configures an API within a repository, either adding it to an
@@ -307,7 +326,8 @@ func (c *Docker) Configure(ctx context.Context, request *ConfigureRequest) (stri
 		mounts = append(mounts, fmt.Sprintf("%s/%s:/repo/%s:ro", request.RepoDir, globalFile, globalFile))
 	}
 
-	if err := c.runDocker(ctx, CommandConfigure, mounts, commandArgs); err != nil {
+	image := c.resolveImage(request.Image)
+	if err := c.runDocker(ctx, image, CommandConfigure, mounts, commandArgs); err != nil {
 		return "", err
 	}
 
@@ -342,14 +362,15 @@ func (c *Docker) ReleaseStage(ctx context.Context, request *ReleaseStageRequest)
 		fmt.Sprintf("%s:/output", request.Output),
 	}
 
-	if err := c.runDocker(ctx, CommandReleaseStage, mounts, commandArgs); err != nil {
+	image := c.resolveImage(request.Image)
+	if err := c.runDocker(ctx, image, CommandReleaseStage, mounts, commandArgs); err != nil {
 		return err
 	}
 
 	return nil
 }
 
-func (c *Docker) runDocker(_ context.Context, command Command, mounts []string, commandArgs []string) (err error) {
+func (c *Docker) runDocker(_ context.Context, image string, command Command, mounts []string, commandArgs []string) (err error) {
 	mounts = maybeRelocateMounts(c.HostMount, mounts)
 	args := []string{
 		"run",
@@ -365,7 +386,7 @@ func (c *Docker) runDocker(_ context.Context, command Command, mounts []string,
 		args = append(args, "--user", fmt.Sprintf("%s:%s", c.uid, c.gid))
 	}
 
-	args = append(args, c.Image)
+	args = append(args, image)
 	args = append(args, string(command))
 	args = append(args, commandArgs...)
 	return c.run(args...)
@@ -389,6 +410,13 @@ func maybeRelocateMounts(hostMount string, mounts []string) []string {
 	return relocatedMounts
 }
 
+func (c *Docker) resolveImage(requestedImage string) string {
+	if requestedImage != "" {
+		return requestedImage
+	}
+	return c.Image
+}
+
 func (c *Docker) runCommand(cmdName string, args ...string) error {
 	cmd := exec.Command(cmdName, args...)
 	cmd.Stderr = os.Stderr

internal/docker/docker_test.go

@@ -144,6 +144,38 @@ func TestDockerRun(t *testing.T) {
 			wantErr:    true,
 			wantErrMsg: simulateDockerErrMsg,
 		},
+		{
+			name: "Generate runs in docker with image override",
+			docker: &Docker{
+				Image:     testImage,
+				HostMount: "hostDir:localDir",
+			},
+			runCommand: func(ctx context.Context, d *Docker) error {
+				generateRequest := &GenerateRequest{
+					State:     state,
+					RepoDir:   repoDir,
+					ApiRoot:   testAPIRoot,
+					Output:    "hostDir",
+					LibraryID: testLibraryID,
+					Image:     "custom-image:abcd123",
+				}
+
+				return d.Generate(ctx, generateRequest)
+			},
+			want: []string{
+				"run", "--rm",
+				"-v", fmt.Sprintf("%s/.librarian:/librarian", repoDir),
+				"-v", fmt.Sprintf("%s/.librarian/generator-input:/input", repoDir),
+				"-v", "localDir:/output",
+				"-v", fmt.Sprintf("%s:/source:ro", testAPIRoot),
+				"custom-image:abcd123",
+				string(CommandGenerate),
+				"--librarian=/librarian",
+				"--input=/input",
+				"--output=/output",
+				"--source=/source",
+			},
+		},
 		{
 			name: "Generate runs in docker",
 			docker: &Docker{
@@ -199,6 +231,32 @@ func TestDockerRun(t *testing.T) {
 				"--repo=/repo",
 			},
 		},
+		{
+			name: "Build runs in docker with image override",
+			docker: &Docker{
+				Image:     testImage,
+				HostMount: "hostDir:localDir",
+			},
+			runCommand: func(ctx context.Context, d *Docker) error {
+				buildRequest := &BuildRequest{
+					State:     state,
+					LibraryID: testLibraryID,
+					RepoDir:   repoDir,
+					Image:     "custom-image:abcd123",
+				}
+
+				return d.Build(ctx, buildRequest)
+			},
+			want: []string{
+				"run", "--rm",
+				"-v", fmt.Sprintf("%s/.librarian:/librarian", repoDir),
+				"-v", fmt.Sprintf("%s:/repo", repoDir),
+				"custom-image:abcd123",
+				string(CommandBuild),
+				"--librarian=/librarian",
+				"--repo=/repo",
+			},
+		},
 		{
 			name: "Build with invalid repo dir",
 			docker: &Docker{
@@ -273,6 +331,46 @@ func TestDockerRun(t *testing.T) {
 				"--source=/source",
 			},
 		},
+		{
+			name: "Configure runs in docker with image override",
+			docker: &Docker{
+				Image: testImage,
+			},
+			runCommand: func(ctx context.Context, d *Docker) error {
+				configureRequest := &ConfigureRequest{
+					State:     state,
+					LibraryID: testLibraryID,
+					RepoDir:   repoDir,
+					ApiRoot:   testAPIRoot,
+					Output:    testOutput,
+					GlobalFiles: []string{
+						"a/b/go.mod",
+						"go.mod",
+					},
+					Image: "custom-image:abcd123",
+				}
+
+				_, err := d.Configure(ctx, configureRequest)
+
+				return err
+			},
+			want: []string{
+				"run", "--rm",
+				"-v", fmt.Sprintf("%s/.librarian:/librarian", repoDir),
+				"-v", fmt.Sprintf("%s/.librarian/generator-input:/input", repoDir),
+				"-v", fmt.Sprintf("%s:/output", testOutput),
+				"-v", fmt.Sprintf("%s:/source:ro", testAPIRoot),
+				"-v", fmt.Sprintf("%s/a/b/go.mod:/repo/a/b/go.mod:ro", repoDir),
+				"-v", fmt.Sprintf("%s/go.mod:/repo/go.mod:ro", repoDir),
+				"custom-image:abcd123",
+				string(CommandConfigure),
+				"--librarian=/librarian",
+				"--input=/input",
+				"--output=/output",
+				"--repo=/repo",
+				"--source=/source",
+			},
+		},
 		{
 			name: "configure_with_nil_global_files",
 			docker: &Docker{
@@ -521,6 +619,41 @@ func TestDockerRun(t *testing.T) {
 				"--output=/output",
 			},
 		},
+		{
+			name: "Release stage runs in docker with image override",
+			docker: &Docker{
+				Image: testImage,
+			},
+			runCommand: func(ctx context.Context, d *Docker) error {
+				partialRepoDir := filepath.Join(repoDir, "release-stage-all-libraries")
+				if err := os.MkdirAll(filepath.Join(repoDir, config.LibrarianDir), 0755); err != nil {
+					t.Fatal(err)
+				}
+
+				releaseInitRequest := &ReleaseStageRequest{
+					State:           state,
+					Output:          testOutput,
+					LibrarianConfig: &config.LibrarianConfig{},
+					RepoDir:         partialRepoDir,
+					Image:           "custom-image:abcd123",
+				}
+
+				defer os.RemoveAll(partialRepoDir)
+
+				return d.ReleaseStage(ctx, releaseInitRequest)
+			},
+			want: []string{
+				"run", "--rm",
+				"-v", fmt.Sprintf("%s/.librarian:/librarian", filepath.Join(repoDir, "release-stage-all-libraries")),
+				"-v", fmt.Sprintf("%s:/repo:ro", filepath.Join(repoDir, "release-stage-all-libraries")),
+				"-v", fmt.Sprintf("%s:/output", testOutput),
+				"custom-image:abcd123",
+				string(CommandReleaseStage),
+				"--librarian=/librarian",
+				"--repo=/repo",
+				"--output=/output",
+			},
+		},
 		{
 			name: "Release stage returns error",
 			docker: &Docker{

internal/librarian/generate.go

@@ -48,6 +48,7 @@ func generateSingleLibrary(ctx context.Context, containerClient ContainerClient,
 		Output:    libraryOutputDir,
 		RepoDir:   repo.GetDir(),
 		State:     state,
+		Image:     state.Image,
 	}
 	slog.Info("performing generation for library", "id", libraryState.ID, "outputDir", libraryOutputDir)
 	if err := containerClient.Generate(ctx, generateRequest); err != nil {

internal/librarian/mocks_test.go

@@ -157,6 +157,8 @@ type mockContainerClient struct {
 	// Set this value if you want the configure-response
 	// has library source roots and remove regex.
 	configureLibraryPaths []string
+	// The last generation request
+	generateRequest *docker.GenerateRequest
 }
 
 func (m *mockContainerClient) Build(ctx context.Context, request *docker.BuildRequest) error {
@@ -241,6 +243,7 @@ func (m *mockContainerClient) Configure(ctx context.Context, request *docker.Con
 
 func (m *mockContainerClient) Generate(ctx context.Context, request *docker.GenerateRequest) error {
 	m.generateCalls++
+	m.generateRequest = request
 
 	if m.noGenerateResponse {
 		return m.generateErr

internal/librarian/update_image_test.go

@@ -210,6 +210,7 @@ func TestUpdateImageRunnerRun(t *testing.T) {
 		wantCreateIssueCalls       int
 		wantCommitMsg              string
 		checkoutError              error
+		wantImage                  string
 	}{
 		{
 			name:  "specific image",
@@ -234,6 +235,7 @@ func TestUpdateImageRunnerRun(t *testing.T) {
 			wantGenerateCalls:   1,
 			wantBuildCalls:      0, // no -build flag
 			wantCheckoutCalls:   2,
+			wantImage:           "some-image",
 		},
 		{
 			name:  "no change image",
@@ -258,6 +260,7 @@ func TestUpdateImageRunnerRun(t *testing.T) {
 			wantGenerateCalls:   1,
 			wantBuildCalls:      0,
 			wantCheckoutCalls:   2,
+			wantImage:           "gcr.io/test/image:v1.2.3",
 		},
 		{
 			name: "finds latest image",
@@ -283,6 +286,7 @@ func TestUpdateImageRunnerRun(t *testing.T) {
 			wantGenerateCalls:   1,
 			wantBuildCalls:      0, // no -build flag
 			wantCheckoutCalls:   2,
+			wantImage:           "gcr.io/test/image@sha256:abc123",
 		},
 		{
 			name: "finds image error",
@@ -761,6 +765,11 @@ func TestUpdateImageRunnerRun(t *testing.T) {
 			if diff := cmp.Diff(test.wantGenerateCalls, test.containerClient.generateCalls); diff != "" {
 				t.Errorf("%s: run() generateCalls mismatch (-want +got):%s", test.name, diff)
 			}
+			if test.wantImage != "" {
+				if diff := cmp.Diff(test.wantImage, test.containerClient.generateRequest.Image); diff != "" {
+					t.Errorf("%s: run() image mismatch (-want +got):%s", test.name, diff)
+				}
+			}
 			if diff := cmp.Diff(test.wantBuildCalls, test.containerClient.buildCalls); diff != "" {
 				t.Errorf("%s: run() buildCalls mismatch (-want +got):%s", test.name, diff)
 			}