Import path re-writer

The path re-writer is the go script to re-write imports
in the go source code files, proto files and go.mod file.
The script accepts path to the project dir where go.mod file
locates, current module version and desired module version.

Upgrading a module requires re-generating the gRPC stubs
from proto file that is why the code of the path re-writer
script is imported in a new 'upgrade-module' task which covers
that need.

Part of: https://gitlab.com/gitlab-org/gitaly/-/issues/3177

3 files changed, 344 insertions, 0 deletions

diff --git a/Makefile b/Makefile
index 62350f73b..f993e7d6d 100644
--- a/Makefile
+++ b/Makefile
@@ -341,6 +341,11 @@ smoke-test: TEST_PACKAGES := ${SOURCE_DIR}/internal/gitaly/rubyserver
 smoke-test: all rspec
 	$(call run_go_tests)
 
+.PHONY: upgrade-module
+upgrade-module:
+	${Q}go run ${SOURCE_DIR}/_support/module-updater/main.go -dir . -from=${FROM_MODULE} -to=${TO_MODULE}
+	${Q}${MAKE} proto
+
 .PHONY: git
 git: ${GIT_INSTALL_DIR}/bin/git
 
diff --git a/_support/module-updater/main.go b/_support/module-updater/main.go
new file mode 100644
index 000000000..c58aad535
--- /dev/null
+++ b/_support/module-updater/main.go
@@ -0,0 +1,324 @@
+package main
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"flag"
+	"fmt"
+	"go/format"
+	"go/parser"
+	"go/token"
+	"io/fs"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strconv"
+	"strings"
+)
+
+var skipDirs = map[string]bool{
+	".git":                              true,
+	".gitlab":                           true,
+	"_build":                            true,
+	"_support":                          true,
+	"changelogs":                        true,
+	"danger":                            true,
+	"doc":                               true,
+	"proto/go/gitalypb":                 true,
+	"proto/go/internal/linter/testdata": true,
+	"ruby":                              true,
+	"scripts":                           true,
+	"unreleased":                        true,
+}
+
+func main() {
+	if err := changeModuleVersion(); err != nil {
+		fmt.Println(err.Error())
+		os.Exit(1)
+	}
+}
+
+func changeModuleVersion() error {
+	dirLocation, fromVersion, toVersion, err := getInput()
+	if err != nil {
+		return err
+	}
+
+	moduleAbsRootPath, err := filepath.Abs(dirLocation)
+	if err != nil {
+		return fmt.Errorf("define absolute module path: %w", err)
+	}
+
+	if err := verifyModulePath(moduleAbsRootPath); err != nil {
+		return err
+	}
+
+	module, err := getModule(moduleAbsRootPath)
+	if err != nil {
+		return fmt.Errorf("define module path: %w", err)
+	}
+
+	prev, next, err := normaliseDesiredImportReplacement(module, fromVersion, toVersion)
+	if err != nil {
+		return err
+	}
+
+	if err := rewriteImports(moduleAbsRootPath, prev, next); err != nil {
+		return fmt.Errorf("re-write go imports: %s", err)
+	}
+
+	if err := rewriteProto(moduleAbsRootPath, prev, next); err != nil {
+		return fmt.Errorf("re-write .proto files: %s", err)
+	}
+
+	if err := rewriteGoMod(moduleAbsRootPath, next); err != nil {
+		return fmt.Errorf("re-write go.mod file: %s", err)
+	}
+
+	return nil
+}
+
+func getInput() (dirLocation string, fromVersion string, toVersion string, err error) {
+	flag.StringVar(&dirLocation, "dir", ".", "directory of the module with the go.mod file")
+	flag.StringVar(&fromVersion, "from", "", "module version to upgrade from")
+	flag.StringVar(&toVersion, "to", "", "module version to upgrade to")
+	flag.Parse()
+	return
+}
+
+func isModuleVersion(input string) error {
+	if input == "" {
+		return errors.New("empty module version")
+	}
+
+	if !strings.HasPrefix(input, "v") {
+		return fmt.Errorf("module version should start with 'v': %s", input)
+	}
+
+	rawVersion, err := strconv.ParseInt(input[1:], 10, 64)
+	if err != nil {
+		return err
+	}
+	if rawVersion < 1 {
+		return fmt.Errorf("version number should be positive: %d", rawVersion)
+	}
+
+	return nil
+}
+
+func getModule(modDir string) (string, error) {
+	cmd := exec.Command("go", "mod", "edit", "-json")
+	cmd.Dir = modDir
+	data, err := cmd.Output()
+	if err != nil {
+		var exitErr *exec.ExitError
+		if errors.As(err, &exitErr) {
+			return "", fmt.Errorf("command %q: %v", strings.Join(cmd.Args, " "), exitErr.Stderr)
+		}
+		return "", fmt.Errorf("command %q: %w", strings.Join(cmd.Args, " "), err)
+	}
+
+	var modInfo = struct{ Module struct{ Path string } }{}
+	if err := json.Unmarshal(data, &modInfo); err != nil {
+		return "", err
+	}
+
+	return modInfo.Module.Path, nil
+}
+
+// rewriteImports rewrites go source files by replacing old import path for the module with the new one.
+func rewriteImports(moduleAbsRootPath, prev, next string) error {
+	fileSet := token.NewFileSet()
+	if err := filepath.Walk(moduleAbsRootPath, func(path string, info fs.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+
+		if info.IsDir() {
+			relPath := strings.TrimPrefix(path, moduleAbsRootPath)
+			relPath = strings.TrimPrefix(relPath, "/")
+			if skipDirs[relPath] {
+				return fs.SkipDir
+			}
+			return nil
+		}
+
+		if filepath.Ext(info.Name()) == ".go" {
+			fileSet.AddFile(path, fileSet.Base(), 0)
+		}
+
+		return nil
+	}); err != nil {
+		return fmt.Errorf("scan directory for go source files: %w", err)
+	}
+
+	seen := map[string]struct{}{}
+	var rerr error
+	fileSet.Iterate(func(file *token.File) bool {
+		if _, found := seen[file.Name()]; found {
+			return false
+		}
+		seen[file.Name()] = struct{}{}
+
+		astFile, err := parser.ParseFile(fileSet, file.Name(), nil, parser.ParseComments)
+		if err != nil {
+			rerr = err
+			return false
+		}
+
+		for _, imprt := range astFile.Imports {
+			oldImport, err := strconv.Unquote(imprt.Path.Value)
+			if err != nil {
+				rerr = fmt.Errorf("unquote import: %s :%w", imprt.Path.Value, err)
+				return false
+			}
+
+			if newImport := strings.Replace(oldImport, prev, next, 1); newImport != oldImport {
+				imprt.EndPos = imprt.End()
+				imprt.Path.Value = strconv.Quote(newImport)
+			}
+		}
+
+		f, err := os.Create(file.Name())
+		if err != nil {
+			rerr = fmt.Errorf("open file %q: %w", file.Name(), err)
+			return false
+		}
+
+		ferr := format.Node(f, fileSet, astFile)
+		cerr := f.Close()
+
+		if ferr != nil {
+			rerr = fmt.Errorf("rewrite file %q: %w", file.Name(), err)
+			return false
+		}
+
+		if cerr != nil {
+			rerr = fmt.Errorf("close file %q: %w", file.Name(), err)
+			return false
+		}
+
+		return true
+	})
+
+	return rerr
+}
+
+func normaliseDesiredImportReplacement(module, from, to string) (string, string, error) {
+	if err := isModuleVersion(from); err != nil {
+		return "", "", fmt.Errorf("invalid 'from' version: %w", err)
+	}
+
+	if err := isModuleVersion(to); err != nil {
+		return "", "", fmt.Errorf("invalid 'to' version: %w", err)
+	}
+
+	prev := module
+	next := filepath.Join(filepath.Dir(module), to)
+	current := filepath.Base(module)
+	if err := isModuleVersion(current); err == nil {
+		if current != from {
+			return "", "", fmt.Errorf("existing module version is %q, but 'from' specified as %q", current, from)
+		}
+	} else {
+		next = filepath.Join(module, to)
+	}
+	return prev, next, nil
+}
+
+func verifyModulePath(moduleRootPath string) error {
+	st, err := os.Stat(moduleRootPath)
+	if err != nil {
+		return fmt.Errorf("inspect module root path: %w", err)
+	}
+
+	if !st.IsDir() {
+		return fmt.Errorf("provided module root path is not a directory: %s", moduleRootPath)
+	}
+
+	entries, err := os.ReadDir(moduleRootPath)
+	if err != nil {
+		return fmt.Errorf("inspect module root path: %w", err)
+	}
+
+	var modExists bool
+	for _, entry := range entries {
+		if entry.Name() == "go.mod" {
+			modExists = true
+			break
+		}
+	}
+
+	if !modExists {
+		return fmt.Errorf("provided module root path doesn't contain go.mod file: %s", moduleRootPath)
+	}
+
+	return nil
+}
+
+// rewriteProto re-write proto files by changing the go_package option declaration:
+// 1. option go_package = "gitlab.com/gitlab-org/gitaly/proto/go/gitalypb";
+// 2. option go_package = "gitlab.com/gitlab-org/gitaly/v14/proto/go/gitalypb";
+// 4. option go_package = "gitlab.com/gitlab-org/gitaly/v15/proto/go/gitalypb";
+func rewriteProto(moduleAbsRootPath, prev, next string) error {
+	protoDirPath := filepath.Join(moduleAbsRootPath, "proto")
+	if err := filepath.Walk(protoDirPath, func(path string, info fs.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+
+		if info.IsDir() {
+			if info.Name() == "go" {
+				return fs.SkipDir
+			}
+			return nil
+		}
+
+		if filepath.Ext(info.Name()) != ".proto" {
+			return nil
+		}
+
+		data, err := os.ReadFile(path)
+		if err != nil {
+			return err
+		}
+
+		var modified [][]byte
+		lines := bytes.Split(data, []byte{'\n'})
+		for _, line := range lines {
+			tokens := bytes.Fields(line)
+			pckg := bytes.Join(tokens, []byte{'~'})
+			if !bytes.HasPrefix(pckg, []byte(`option~go_package~=~"`+prev+`/proto/go/gitalypb";`)) {
+				modified = append(modified, line)
+				continue
+			}
+
+			modified = append(modified, bytes.ReplaceAll(line, []byte(prev), []byte(next)))
+		}
+		if len(modified) == 0 {
+			return nil
+		}
+		modifiedData := bytes.Join(modified, []byte{'\n'})
+
+		if err := os.WriteFile(path, modifiedData, info.Mode()); err != nil {
+			return fmt.Errorf("write modified file content: %s, %w", info.Name(), err)
+		}
+
+		return nil
+	}); err != nil {
+		return err
+	}
+	return nil
+}
+
+// rewriteGoMod modifies name of the module in the go.mod file.
+func rewriteGoMod(moduleAbsRootPath, next string) error {
+	cmd := exec.Command("go", "mod", "edit", "-module", next)
+	cmd.Dir = moduleAbsRootPath
+	if err := cmd.Run(); err != nil {
+		return fmt.Errorf("command %q: %w", strings.Join(cmd.Args, " "), err)
+	}
+
+	return nil
+}
diff --git a/doc/PROCESS.md b/doc/PROCESS.md
index 59ccf628c..92fa86b82 100644
--- a/doc/PROCESS.md
+++ b/doc/PROCESS.md
@@ -403,6 +403,21 @@ If the changes needed are not yet released, [create a release candidate](#creati
 - Checkout the tag to publish (vX.Y.Z)
 - run `_support/publish-gem X.Y.Z`
 
+### Publishing the go module
+
+If an [updated version](https://golang.org/doc/modules/release-workflow) of the go module is needed, it can be [published](https://golang.org/doc/modules/publishing)
+by tag creation.
+
+If a new [major module version update](https://golang.org/doc/modules/major-version) is needed,
+it can be changed by running `upgrade-module` `make` task with desired parameters:
+
+```bash
+make upgrade-module FROM_MODULE=v14 TO_MODULE=v15
+```
+
+It replaces old imports with the new version in the go source files,
+updates `*.proto` files and modifies `go.mod` file to use a new target version of the module.
+
 ##### Security release
 
 Security releases involve additional processes to ensure that recent releases