bilibili-go-common/app/tool/bmproto/bmgen/main.go

package main

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"os/exec"
	"os/user"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/pkg/errors"
	"github.com/urfave/cli"
)

var tplFlag = false
var grpcFlag = false

// 使用app/tool/warden/protoc.sh 生成grpc .pb.go
var protocShRunned = false

var syncLiveDoc = false

func main() {
	app := cli.NewApp()
	app.Name = "bmgen"
	app.Usage = "根据proto文件生成bm框架或者grpc代码: \n" +
		"用法1，在项目的任何一个位置运行bmgen，会自动找到proto文件\n" +
		"用法2，bmgen proto文件(文件必须在一个项目的api中（项目包含cmd和api目录）)\n"
	app.Version = "1.0.0"
	app.Commands = []cli.Command{
		{
			Name:   "update",
			Usage:  "更新工具本身",
			Action: actionUpdate,
		},
		{
			Name:   "clean-live-doc",
			Usage:  "清除live-doc已经失效的分支的文档",
			Action: actionCleanLiveDoc,
		},
	}
	app.Action = actionGenerate
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "interface, i",
			Value: "",
			Usage: "generate for the interface project `RELATIVE-PATH`, eg. live/live-demo",
		},
		cli.BoolFlag{
			Name:        "grpc, g",
			Destination: &grpcFlag,
			Usage:       "废弃，会自动检测是否需要grpc",
		},
		cli.BoolFlag{
			Name:        "tpl, t",
			Destination: &tplFlag,
			Usage:       "是否生成service模板代码",
		},
		cli.BoolFlag{
			Name:        "live-doc, l",
			Destination: &syncLiveDoc,
			Usage:       "同步该项目的文档到live-doc https://git.bilibili.co/live-dev/live-doc/tree/doc/proto/$branch/$package/$filename.$Service.md",
		},
	}
	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}

func autoUpdate(ctx *cli.Context) (err error) {
	tfile, e := ioutil.ReadFile("/tmp/bmgentime")
	if e != nil {
		tfile = []byte{'0'}
	}
	ts, _ := strconv.ParseInt(string(tfile), 0, 64)
	current := time.Now().Unix()
	if (current - int64(ts)) > 3600*12 { // 12 hours old
		ioutil.WriteFile("/tmp/bmgentime", []byte(strconv.FormatInt(current, 10)), 0777)
		err = actionUpdate(ctx)
		if err != nil {
			return
		}
	}
	return
}

func actionCleanLiveDoc(ctx *cli.Context) (err error) {
	fmt.Println("暂未实现，敬请期待")
	return
}

func installDependencies() (err error) {
	// 检查protoc
	//
	e := runCmd("which protoc")
	if e != nil {
		var uname string
		uname, err = runCmdRet("uname")
		if err != nil {
			return
		}
		if uname == "Darwin" {
			err = runCmd("brew install protobuf")
			if err != nil {
				return
			}
		} else {
			fmt.Println("找不到protoc，请于 https://github.com/protocolbuffers/protobuf/releases 下载里面的protoc-$VERSION-$OS.zip 安装, 注意把文件拷贝到正确的未知")
			return errors.New("找不到protoc")
		}
	}
	err = runCmd("which protoc-gen-gogofast || go get github.com/gogo/protobuf/protoc-gen-gogofast")
	return
}

// actionGenerate invokes protoc to generate files
func actionGenerate(ctx *cli.Context) (err error) {
	if err = autoUpdate(ctx); err != nil {
		return
	}
	err = installDependencies()
	if err != nil {
		return
	}
	f := ctx.Args().Get(0)

	goPath := initGopath()
	if !fileExist(goPath) {
		return cli.NewExitError(fmt.Sprintf("GOPATH not exist: "+goPath), 1)
	}

	filesToGenerate := []string{f}
	iPath := ctx.String("i")
	if iPath != "" {
		iPath = goPath + "/src/go-common/app/interface/" + iPath
		if !fileExist(iPath) {
			return cli.NewExitError(fmt.Sprintf("interface project not found: "+iPath), 1)
		}
		pbs := filesWithSuffix(iPath+"/api", ".pb", ".proto")
		if len(pbs) == 0 {
			return cli.NewExitError(fmt.Sprintf("no pbs found in path: "+iPath+"/api"), 1)
		}
		filesToGenerate = pbs
		fmt.Printf(".pb files found %v\n", pbs)
	} else {
		if f == "" {
			// if is is empty, look up project that contains current dir
			abs, _ := filepath.Abs(".")
			proj := lookupProjPath(abs)
			if proj == "" {
				return cli.NewExitError("current dir is not in any project : "+abs, 1)
			}
			if proj != "" {
				pbs := filesWithSuffix(proj+"/api", ".pb", ".proto")
				if len(pbs) == 0 {
					return cli.NewExitError(fmt.Sprintf("no pbs found in path: "+proj+"/api"), 1)
				}
				filesToGenerate = pbs
				fmt.Printf(".pb files found %v\n", pbs)
			}
		}
	}

	for _, p := range filesToGenerate {
		if !fileExist(p) {
			return cli.NewExitError(fmt.Sprintf("file not exist: "+p), 1)
		}
		generateForFile(p, goPath)
	}
	if syncLiveDoc {
		err = actionSyncLiveDoc(ctx)
	}
	return
}

func generateForFile(f string, goPath string) {
	absPath, _ := filepath.Abs(f)
	base := filepath.Base(f)
	projPath := lookupProjPath(absPath)
	fileFolder := filepath.Dir(absPath)
	var relativePath string
	if projPath != "" {
		relativePath = absPath[len(projPath)+1:]
	}

	var cmd string
	if strings.Index(relativePath, "api/liverpc") == 0 {
		// need not generate for liverpc
		fmt.Printf("skip for liverpc \n")
		return
	}
	genTpl := 0
	if tplFlag {
		genTpl = 1
	}
	if !strings.Contains(relativePath, "api/http") {
		//非http 生成grpc和http的代码
		isInsideGoCommon := strings.Contains(fileFolder, "go-common")
		if !protocShRunned && isInsideGoCommon {
			//protoc.sh 只能在大仓库中使用
			protocShRunned = true
			cmd = fmt.Sprintf(`cd "%s" && "%s/src/go-common/app/tool/warden/protoc.sh"`, fileFolder, goPath)
			runCmd(cmd)
		}
		genGrpcArg := "--gogofast_out=plugins=grpc:."
		if isInsideGoCommon {
			// go-common中已经用上述命令生成过了
			genGrpcArg = ""
		}
		cmd = fmt.Sprintf(`cd "%s" && protoc --bm_out=tpl=%d:. `+
			`%s -I. -I%s/src -I"%s/src/go-common" -I"%s/src/go-common/vendor" "%s"`,
			fileFolder, genTpl, genGrpcArg, goPath, goPath, goPath, base)
	} else {
		// 只生成http的代码
		var pbOutArg string
		if strings.LastIndex(f, ".pb") == len(f)-3 {
			// ends with .pb
			log.Printf("\n\033[0;33m======！！！！WARNING！！！！========\n" +
				".pb文件生成代码的功能已经不再维护，请尽快迁移到.proto, 详情：\nhttp://info.bilibili.co/pages/viewpage.action?pageId=11864735#proto文件格式-.pb迁移到.proto\n" +
				"======！！！！WARNING！！！！========\033[0m\n")
			pbOutArg = "--gogofasterg_out=json_emitdefault=1,suffix=.pbg.go:."
		} else {
			pbOutArg = "--gogofast_out=."
		}
		cmd = fmt.Sprintf(`cd "%s" && protoc --bm_out=tpl=%d:. `+
			pbOutArg+` -I. -I"%s/src" -I"%s/src/go-common" -I"%s/src/go-common/vendor" "%s"`,
			fileFolder, genTpl, goPath, goPath, goPath, base)
	}

	err := runCmd(cmd)
	if err == nil {
		fmt.Println("ok")
	}
}

// files all files with suffix
func filesWithSuffix(dir string, suffixes ...string) (result []string) {
	filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			for _, suffix := range suffixes {
				if strings.HasSuffix(info.Name(), suffix) {
					result = append(result, path)
					break
				}
			}
		}
		return nil
	})
	return
}

// 扫描md文档并且同步到live-doc
func actionSyncLiveDoc(ctx *cli.Context) (err error) {
	//mdFiles := filesWithSuffix(
	// if is is empty, look up project that contains current dir
	abs, _ := filepath.Abs(".")
	proj := lookupProjPath(abs)
	if proj == "" {
		err = cli.NewExitError("current dir is not in any project : "+abs, 1)
		return
	}

	var branch string
	branch, err = runCmdRet("git rev-parse --abbrev-ref HEAD")
	if err != nil {
		return
	}
	branch = url.QueryEscape(branch)

	err = runCmd("[ -d ~/.brpc ] || mkdir ~/.brpc")
	if err != nil {
		return
	}
	var liveDocUrl = "[email protected]:live-dev/live-doc.git"
	err = runCmd("cd ~/.brpc && [ -d ~/.brpc/live-doc ] || git clone " + liveDocUrl)
	if err != nil {
		return
	}
	err = runCmd("cd ~/.brpc/live-doc && git checkout doc && git pull origin doc")
	if err != nil {
		return
	}

	mdFiles := filesWithSuffix(proj+"/api", ".md")

	var u *user.User
	u, err = user.Current()
	if err != nil {
		return err
	}
	var liveDocDir = u.HomeDir + "/.brpc/live-doc"

	for _, file := range mdFiles {
		fileHandle, _ := os.Open(file)
		fileScanner := bufio.NewScanner(fileHandle)
		for fileScanner.Scan() {
			var text = fileScanner.Text()
			var reg = regexp.MustCompile(`package=([\w\.]+)`)
			matches := reg.FindStringSubmatch(text)
			if len(matches) >= 2 {
				pkg := matches[1]
				relativeDir := strings.Replace(pkg, ".", "/", -1)
				var destDir = liveDocDir + "/protodoc/" + branch + "/" + relativeDir
				runCmd("mkdir -p " + destDir)
				err = runCmd(fmt.Sprintf("cp %s %s", file, destDir))
				if err != nil {
					return
				}
			} else {
				fmt.Println("package not found", file)
			}
			break
		}
		fileHandle.Close()
	}
	err = runCmd(`cd "` + liveDocDir + `" && git add -A`)
	if err != nil {
		return
	}
	var gitStatus string
	gitStatus, _ = runCmdRet(`cd "` + liveDocDir + `" && git status --porcelain`)
	if gitStatus != "" {
		err = runCmd(`cd "` + liveDocDir + `" && git commit -m 'update doc from proto' && git push origin doc`)
		if err != nil {
			return
		}
	}
	fmt.Printf("文档已经生成至 %s%s\n", "https://git.bilibili.co/live-dev/live-doc/tree/doc/protodoc/", url.QueryEscape(branch))
	return
}

// actionUpdate update the tools its self
func actionUpdate(ctx *cli.Context) (err error) {
	log.Print("Updating bmgen.....")
	goPath := initGopath()
	goCommonPath := goPath + "/src/go-common"
	if !fileExist(goCommonPath) {
		return cli.NewExitError("go-common not exist : "+goCommonPath, 1)
	}
	cmd := fmt.Sprintf(`go install "go-common/app/tool/bmproto/..."`)
	if err = runCmd(cmd); err != nil {
		err = cli.NewExitError(err.Error(), 1)
		return
	}
	log.Print("Updated!")
	return
}

// runCmd runs the cmd & print output (both stdout & stderr)
func runCmd(cmd string) (err error) {
	fmt.Printf("CMD: %s \n", cmd)
	out, err := exec.Command("/bin/bash", "-c", cmd).CombinedOutput()
	fmt.Print(string(out))
	return
}

func runCmdRet(cmd string) (out string, err error) {
	fmt.Printf("CMD: %s \n", cmd)
	outBytes, err := exec.Command("/bin/bash", "-c", cmd).CombinedOutput()
	out = strings.Trim(string(outBytes), "\n\r\t ")
	return
}

// lookupProjPath get project path by proto absolute path
// assume that proto is in the project's model directory
func lookupProjPath(protoAbs string) (result string) {
	lastIndex := len(protoAbs)
	curPath := protoAbs

	for lastIndex > 0 {
		if fileExist(curPath+"/cmd") && fileExist(curPath+"/api") {
			result = curPath
			return
		}
		lastIndex = strings.LastIndex(curPath, string(os.PathSeparator))
		curPath = protoAbs[:lastIndex]
	}
	result = ""
	return
}

func fileExist(file string) bool {
	_, err := os.Stat(file)
	return err == nil
}

func initGopath() string {
	root, err := goPath()
	if err != nil || root == "" {
		log.Printf("can not read GOPATH, use ~/go as default GOPATH")
		root = path.Join(os.Getenv("HOME"), "go")
	}
	return root
}

func goPath() (string, error) {
	gopaths := strings.Split(os.Getenv("GOPATH"), ":")
	if len(gopaths) == 1 {
		return gopaths[0], nil
	}
	for _, gp := range gopaths {
		absgp, err := filepath.Abs(gp)
		if err != nil {
			return "", err
		}
		if fileExist(absgp + "/src/go-common") {
			return absgp, nil
		}
	}
	return "", fmt.Errorf("can't found current gopath")
}