feat: add update-image CLI command (#2580)

chingor13 · web-flow · commit 90e0f6e52bf0 · 2025-10-16T14:41:44.000-05:00
Adds new `librarian update-image` CLI command. Towards #2342 Fixes #2546
diff --git a/cmd/librarian/doc.go b/cmd/librarian/doc.go
@@ -269,6 +269,67 @@ Flags:
 	  	is configured as a language repository.
 	-v	enables verbose logging
 
+# update-image
+
+The 'update-image' command is used to update the 'image' SHA
+of the language container for a language repository.
+
+This command's primary responsibilities are to:
+
+  - Update the 'image' field in '.librarian/state.yaml'
+  - Regenerate each library with the new language container using googleapis'
+    proto definitions at the 'last_generated_commit'
+
+Examples:
+
+	# Create a PR that updates the language container to latest image.
+	librarian update-image --commit --push
+
+	# Create a PR that updates the language container to the specified image.
+	librarian update-image --commit --push --image=<some-image-with-sha>
+
+Usage:
+
+	librarian update-image [flags]
+
+Flags:
+
+	-api-source string
+	  	The location of an API specification repository.
+	  	Can be a remote URL or a local file path. (default "https://github.com/googleapis/googleapis")
+	-branch string
+	  	The branch to use with remote code repositories. This is used to specify
+	  	which branch to clone and which branch to use as the base for a pull
+	  	request. (default "main")
+	-build
+	  	If true, Librarian will build each generated library by invoking the
+	  	language-specific container.
+	-commit
+	  	If true, librarian will create a commit for the change but not create
+	  	a pull request. This flag is ignored if push is set to true.
+	-host-mount string
+	  	For use when librarian is running in a container. A mapping of a
+	  	directory from the host to the container, in the format
+	  	<host-mount>:<local-mount>.
+	-image string
+	  	Language specific image used to invoke code generation and releasing.
+	  	If not specified, the image configured in the state.yaml is used.
+	-output string
+	  	Working directory root. When this is not specified, a working directory
+	  	will be created in /tmp.
+	-push
+	  	If true, Librarian will create a commit,
+	  	push and create a pull request for the changes.
+	  	A GitHub token with push access must be provided via the
+	  	LIBRARIAN_GITHUB_TOKEN environment variable.
+	-repo string
+	  	Code repository where the generated code will reside. Can be a remote
+	  	in the format of a remote URL such as https://github.com/{owner}/{repo} or a
+	  	local file path like /path/to/repo. Both absolute and relative paths are
+	  	supported. If not specified, will try to detect if the current working directory
+	  	is configured as a language repository.
+	-v	enables verbose logging
+
 # version
 
 Version prints version information for the librarian binary.
diff --git a/internal/librarian/command.go b/internal/librarian/command.go
@@ -53,6 +53,7 @@ const (
 	pullRequestOnboard
 	pullRequestGenerate
 	pullRequestRelease
+	pullRequestUpdateImage
 )
 
 // String returns the string representation of a pullRequestType.
@@ -63,6 +64,7 @@ func (t pullRequestType) String() string {
 		pullRequestOnboard:     "onboard",
 		pullRequestGenerate:    "generate",
 		pullRequestRelease:     "release",
+		pullRequestUpdateImage: "update image",
 	}
 	if name, ok := names[t]; ok {
 		return name
@@ -97,17 +99,29 @@ type ContainerClient interface {
 }
 
 type commitInfo struct {
-	branch            string
-	commit            bool
-	commitMessage     string
-	ghClient          GitHubClient
-	prType            pullRequestType
+	// branch is the base branch of the created pull request.
+	branch string
+	// commit declares whether or not to create a commit.
+	commit bool
+	// commitMessage is used as the message on the actual git commit.
+	commitMessage string
+	// ghClient is used to interact with the GitHub API.
+	ghClient GitHubClient
+	// prType is an enum for which type of librarian pull request we are creating.
+	prType pullRequestType
+	// pullRequestLabels is a list of labels to add to the created pull request.
 	pullRequestLabels []string
-	push              bool
-	languageRepo      gitrepo.Repository
-	sourceRepo        gitrepo.Repository
-	state             *config.LibrarianState
-	workRoot          string
+	// push declares whether or not to push the commits to GitHub.
+	push bool
+	// languageRepo is the git repository containing the language-specific libraries.
+	languageRepo gitrepo.Repository
+	// sourceRepo is the git repository containing the source protos.
+	sourceRepo gitrepo.Repository
+	// state is the librarian state.yaml contents.
+	state *config.LibrarianState
+	// workRoot is the directory that we stage code changes in.
+	workRoot string
+	// failedGenerations is the number of generations that failed.
 	failedGenerations int
 	// api is the api path of a library, only set this value during api onboarding.
 	api string
@@ -137,7 +151,9 @@ func newCommandRunner(cfg *config.Config) (*commandRunner, error) {
 		sourceRepo    gitrepo.Repository
 		sourceRepoDir string
 	)
-	if cfg.CommandName == generateCmdName {
+
+	// If APISource is set, checkout the protos repository.
+	if cfg.APISource != "" {
 		sourceRepo, err = cloneOrOpenRepo(cfg.WorkRoot, cfg.APISource, cfg.APISourceDepth, defaultAPISourceBranch, cfg.CI, cfg.GitHubToken)
 		if err != nil {
 			return nil, err
diff --git a/internal/librarian/command_test.go b/internal/librarian/command_test.go
@@ -1518,7 +1518,7 @@ func TestCommitAndPush(t *testing.T) {
 				return &mockGitHubClient{}
 			},
 			state:          &config.LibrarianState{},
-			prType:         4,
+			prType:         100,
 			push:           true,
 			wantErr:        true,
 			expectedErrMsg: "failed to create pull request body",
diff --git a/internal/librarian/generate_command_test.go b/internal/librarian/generate_command_test.go
@@ -34,6 +34,7 @@ func TestNewGenerateRunner(t *testing.T) {
 		cfg        *config.Config
 		wantErr    bool
 		wantErrMsg string
+		setupFunc  func(*config.Config) error
 	}{
 		{
 			name: "valid config",
@@ -102,15 +103,27 @@ func TestNewGenerateRunner(t *testing.T) {
 			name: "clone googleapis fails",
 			cfg: &config.Config{
 				API:            "some/api",
-				APISource:      "", // This will trigger the clone of googleapis
+				APISource:      "https://github.com/googleapis/googleapis", // This will trigger the clone of googleapis
 				APISourceDepth: 1,
 				Repo:           newTestGitRepo(t).GetDir(),
 				WorkRoot:       t.TempDir(),
 				Image:          "gcr.io/test/test-image",
 				CommandName:    generateCmdName,
 			},
 			wantErr:    true,
-			wantErrMsg: "repo must be specified",
+			wantErrMsg: "repository does not exist",
+			setupFunc: func(cfg *config.Config) error {
+				// The function will try to clone googleapis into the current work directory.
+				// To make it fail, create a non-empty, non-git directory.
+				googleapisDir := filepath.Join(cfg.WorkRoot, "googleapis")
+				if err := os.MkdirAll(googleapisDir, 0755); err != nil {
+					return err
+				}
+				if err := os.WriteFile(filepath.Join(googleapisDir, "some-file"), []byte("foo"), 0644); err != nil {
+					return err
+				}
+				return nil
+			},
 		},
 		{
 			name: "valid config with local repo",
@@ -127,32 +140,11 @@ func TestNewGenerateRunner(t *testing.T) {
 	} {
 		t.Run(test.name, func(t *testing.T) {
 			t.Parallel()
-			if test.cfg.APISource == "" && test.cfg.WorkRoot != "" {
-				if test.name == "clone googleapis fails" {
-					// The function will try to clone googleapis into the current work directory.
-					// To make it fail, create a non-empty, non-git directory.
-					googleapisDir := filepath.Join(test.cfg.WorkRoot, "googleapis")
-					if err := os.MkdirAll(googleapisDir, 0755); err != nil {
-						t.Fatalf("os.MkdirAll() = %v", err)
-					}
-					if err := os.WriteFile(filepath.Join(googleapisDir, "some-file"), []byte("foo"), 0644); err != nil {
-						t.Fatalf("os.WriteFile() = %v", err)
-					}
-				} else {
-					// The function will try to clone googleapis into the current work directory.
-					// To prevent a real clone, we can pre-create a fake googleapis repo.
-					googleapisDir := filepath.Join(test.cfg.WorkRoot, "googleapis")
-					if err := os.MkdirAll(googleapisDir, 0755); err != nil {
-						t.Fatalf("os.MkdirAll() = %v", err)
-					}
-					runGit(t, googleapisDir, "init")
-					runGit(t, googleapisDir, "config", "user.email", "test@example.com")
-					runGit(t, googleapisDir, "config", "user.name", "Test User")
-					if err := os.WriteFile(filepath.Join(googleapisDir, "README.md"), []byte("test"), 0644); err != nil {
-						t.Fatalf("os.WriteFile: %v", err)
-					}
-					runGit(t, googleapisDir, "add", "README.md")
-					runGit(t, googleapisDir, "commit", "-m", "initial commit")
+
+			// custom setup
+			if test.setupFunc != nil {
+				if err := test.setupFunc(test.cfg); err != nil {
+					t.Fatalf("error in setup %v", err)
 				}
 			}
 
diff --git a/internal/librarian/help.go b/internal/librarian/help.go
@@ -134,4 +134,20 @@ Examples:
 
   # Find and process all pending merged release PRs in a repository.
   librarian release tag-and-release --repo=https://github.com/googleapis/google-cloud-go`
+
+	updateImageLongHelp = `The 'update-image' command is used to update the 'image' SHA
+of the language container for a language repository.
+
+This command's primary responsibilities are to:
+
+- Update the 'image' field in '.librarian/state.yaml'
+- Regenerate each library with the new language container using googleapis'
+  proto definitions at the 'last_generated_commit'
+  
+Examples:
+  # Create a PR that updates the language container to latest image.
+  librarian update-image --commit --push
+
+  # Create a PR that updates the language container to the specified image.
+  librarian update-image --commit --push --image=<some-image-with-sha>`
 )
diff --git a/internal/librarian/librarian.go b/internal/librarian/librarian.go
@@ -69,6 +69,7 @@ func newLibrarianCommand() *cli.Command {
 		Commands: []*cli.Command{
 			newCmdGenerate(),
 			cmdRelease,
+			newCmdUpdateImage(),
 			cmdVersion,
 		},
 	}
@@ -177,3 +178,39 @@ func newCmdInit() *cli.Command {
 	addFlagVerbose(cmdInit.Flags, &verbose)
 	return cmdInit
 }
+
+func newCmdUpdateImage() *cli.Command {
+	var verbose bool
+	cmdUpdateImage := &cli.Command{
+		Short:     "update-image updates configured language image container",
+		UsageLine: "librarian update-image [flags]",
+		Long:      updateImageLongHelp,
+		Action: func(ctx context.Context, cmd *cli.Command) error {
+			setupLogger(verbose)
+			slog.Debug("update image command verbose logging")
+			if err := cmd.Config.SetDefaults(); err != nil {
+				return fmt.Errorf("failed to initialize config: %w", err)
+			}
+			if _, err := cmd.Config.IsValid(); err != nil {
+				return fmt.Errorf("failed to validate config: %s", err)
+			}
+			runner, err := newUpdateImageRunner(cmd.Config)
+			if err != nil {
+				return err
+			}
+			return runner.run(ctx)
+		},
+	}
+	cmdUpdateImage.Init()
+	addFlagAPISource(cmdUpdateImage.Flags, cmdUpdateImage.Config)
+	addFlagBuild(cmdUpdateImage.Flags, cmdUpdateImage.Config)
+	addFlagCommit(cmdUpdateImage.Flags, cmdUpdateImage.Config)
+	addFlagHostMount(cmdUpdateImage.Flags, cmdUpdateImage.Config)
+	addFlagImage(cmdUpdateImage.Flags, cmdUpdateImage.Config)
+	addFlagRepo(cmdUpdateImage.Flags, cmdUpdateImage.Config)
+	addFlagBranch(cmdUpdateImage.Flags, cmdUpdateImage.Config)
+	addFlagWorkRoot(cmdUpdateImage.Flags, cmdUpdateImage.Config)
+	addFlagPush(cmdUpdateImage.Flags, cmdUpdateImage.Config)
+	addFlagVerbose(cmdUpdateImage.Flags, &verbose)
+	return cmdUpdateImage
+}
diff --git a/internal/librarian/mocks_test.go b/internal/librarian/mocks_test.go
@@ -131,7 +131,13 @@ type mockContainerClient struct {
 	failGenerateForID string
 	// Set this value if you want an error when
 	// generate a library with a specific id.
-	generateErrForID    error
+	generateErrForID error
+	// Set this value if you want an error when
+	// build a library with a specific id.
+	failBuildForID string
+	// Set this value if you want an error when
+	// build a library with a specific id.
+	buildErrForID       error
 	requestLibraryID    string
 	noBuildResponse     bool
 	noConfigureResponse bool
@@ -164,6 +170,13 @@ func (m *mockContainerClient) Build(ctx context.Context, request *docker.BuildRe
 	if err := os.WriteFile(filepath.Join(request.RepoDir, ".librarian", config.BuildResponse), []byte(libraryStr), 0755); err != nil {
 		return err
 	}
+
+	if m.failBuildForID != "" {
+		if request.LibraryID == m.failBuildForID {
+			return m.buildErrForID
+		}
+	}
+
 	return m.buildErr
 }
 
@@ -308,6 +321,7 @@ type MockRepository struct {
 	RemotesValue                           []*gitrepo.Remote
 	RemotesError                           error
 	CommitCalls                            int
+	LastCommitMessage                      string
 	GetCommitError                         error
 	GetLatestCommitError                   error
 	GetCommitByHash                        map[string]*gitrepo.Commit
@@ -330,6 +344,8 @@ type MockRepository struct {
 	RestoreError                           error
 	HeadHashValue                          string
 	HeadHashError                          error
+	CheckoutCalls                          int
+	CheckoutError                          error
 }
 
 func (m *MockRepository) HeadHash() (string, error) {
@@ -355,6 +371,7 @@ func (m *MockRepository) AddAll() error {
 
 func (m *MockRepository) Commit(msg string) error {
 	m.CommitCalls++
+	m.LastCommitMessage = msg
 	return m.CommitError
 }
 
@@ -470,3 +487,27 @@ func (m *MockRepository) Push(name string) error {
 func (m *MockRepository) Restore(paths []string) error {
 	return m.RestoreError
 }
+
+func (m *MockRepository) CleanUntracked(paths []string) error {
+	return nil
+}
+
+func (m *MockRepository) Checkout(commitHash string) error {
+	m.CheckoutCalls++
+	if m.CheckoutError != nil {
+		return m.CheckoutError
+	}
+	return nil
+}
+
+// mockImagesClient is a mock implementation of the ImageRegistryClient interface for testing.
+type mockImagesClient struct {
+	latestImage     string
+	err             error
+	findLatestCalls int
+}
+
+func (m *mockImagesClient) FindLatest(ctx context.Context, imageName string) (string, error) {
+	m.findLatestCalls++
+	return m.latestImage, m.err
+}
diff --git a/internal/librarian/update_image.go b/internal/librarian/update_image.go
diff --git a/internal/librarian/update_image_test.go b/internal/librarian/update_image_test.go
