go-zero/tools/goctl/pkg/env/env.go

package env

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"

	"github.com/zeromicro/go-zero/tools/goctl/internal/version"
	sortedmap "github.com/zeromicro/go-zero/tools/goctl/pkg/collection"
	"github.com/zeromicro/go-zero/tools/goctl/pkg/protoc"
	"github.com/zeromicro/go-zero/tools/goctl/pkg/protocgengo"
	"github.com/zeromicro/go-zero/tools/goctl/pkg/protocgengogrpc"
	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
)

var goctlEnv *sortedmap.SortedMap

const (
	GoctlOS                = "GOCTL_OS"
	GoctlArch              = "GOCTL_ARCH"
	GoctlHome              = "GOCTL_HOME"
	GoctlDebug             = "GOCTL_DEBUG"
	GoctlCache             = "GOCTL_CACHE"
	GoctlVersion           = "GOCTL_VERSION"
	GoctlExperimental      = "GOCTL_EXPERIMENTAL"
	ProtocVersion          = "PROTOC_VERSION"
	ProtocGenGoVersion     = "PROTOC_GEN_GO_VERSION"
	ProtocGenGoGRPCVersion = "PROTO_GEN_GO_GRPC_VERSION"

	envFileDir      = "env"
	ExperimentalOn  = "on"
	ExperimentalOff = "off"
)

// init initializes the goctl environment variables, the environment variables of the function are set in order,
// please do not change the logic order of the code.
func init() {
	defaultGoctlHome, err := pathx.GetDefaultGoctlHome()
	if err != nil {
		log.Fatalln(err)
	}
	goctlEnv = sortedmap.New()
	goctlEnv.SetKV(GoctlOS, runtime.GOOS)
	goctlEnv.SetKV(GoctlArch, runtime.GOARCH)
	existsEnv := readEnv(defaultGoctlHome)
	if existsEnv != nil {
		goctlHome, ok := existsEnv.GetString(GoctlHome)
		if ok && len(goctlHome) > 0 {
			goctlEnv.SetKV(GoctlHome, goctlHome)
		}
		if debug := existsEnv.GetOr(GoctlDebug, "").(string); debug != "" {
			if strings.EqualFold(debug, "true") || strings.EqualFold(debug, "false") {
				goctlEnv.SetKV(GoctlDebug, debug)
			}
		}
		if value := existsEnv.GetStringOr(GoctlCache, ""); value != "" {
			goctlEnv.SetKV(GoctlCache, value)
		}
		experimental := existsEnv.GetOr(GoctlExperimental, ExperimentalOff)
		goctlEnv.SetKV(GoctlExperimental, experimental)
	}

	if !goctlEnv.HasKey(GoctlHome) {
		goctlEnv.SetKV(GoctlHome, defaultGoctlHome)
	}
	if !goctlEnv.HasKey(GoctlDebug) {
		goctlEnv.SetKV(GoctlDebug, "False")
	}

	if !goctlEnv.HasKey(GoctlCache) {
		cacheDir, _ := pathx.GetCacheDir()
		goctlEnv.SetKV(GoctlCache, cacheDir)
	}

	if !goctlEnv.HasKey(GoctlExperimental) {
		goctlEnv.SetKV(GoctlExperimental, ExperimentalOff)
	}

	goctlEnv.SetKV(GoctlVersion, version.BuildVersion)

	protocVer, _ := protoc.Version()
	goctlEnv.SetKV(ProtocVersion, protocVer)

	protocGenGoVer, _ := protocgengo.Version()
	goctlEnv.SetKV(ProtocGenGoVersion, protocGenGoVer)

	protocGenGoGrpcVer, _ := protocgengogrpc.Version()
	goctlEnv.SetKV(ProtocGenGoGRPCVersion, protocGenGoGrpcVer)
}

func Print(args ...string) string {
	if len(args) == 0 {
		return strings.Join(goctlEnv.Format(), "\n")
	}

	var values []string
	for _, key := range args {
		value, ok := goctlEnv.GetString(key)
		if !ok {
			value = fmt.Sprintf("%s=%%not found%%", key)
		}
		values = append(values, fmt.Sprintf("%s=%s", key, value))
	}
	return strings.Join(values, "\n")
}

func Get(key string) string {
	return GetOr(key, "")
}

// Set sets the environment variable for testing
func Set(t *testing.T, key, value string) {
	goctlEnv.SetKV(key, value)
	t.Cleanup(func() {
		goctlEnv.Remove(key)
	})
}

func GetOr(key, def string) string {
	return goctlEnv.GetStringOr(key, def)
}

func UseExperimental() bool {
	return GetOr(GoctlExperimental, ExperimentalOff) == ExperimentalOn
}

func readEnv(goctlHome string) *sortedmap.SortedMap {
	envFile := filepath.Join(goctlHome, envFileDir)
	data, err := os.ReadFile(envFile)
	if err != nil {
		return nil
	}
	dataStr := string(data)
	lines := strings.Split(dataStr, "\n")
	sm := sortedmap.New()
	for _, line := range lines {
		_, _, err = sm.SetExpression(line)
		if err != nil {
			continue
		}
	}
	return sm
}

func WriteEnv(kv []string) error {
	defaultGoctlHome, err := pathx.GetDefaultGoctlHome()
	if err != nil {
		log.Fatalln(err)
	}
	data := sortedmap.New()
	for _, e := range kv {
		_, _, err := data.SetExpression(e)
		if err != nil {
			return err
		}
	}
	data.RangeIf(func(key, value any) bool {
		switch key.(string) {
		case GoctlHome, GoctlCache:
			path := value.(string)
			if !pathx.FileExists(path) {
				err = fmt.Errorf("[writeEnv]: path %q is not exists", path)
				return false
			}
		}
		if goctlEnv.HasKey(key) {
			goctlEnv.SetKV(key, value)
			return true
		} else {
			err = fmt.Errorf("[writeEnv]: invalid key: %v", key)
			return false
		}
	})
	if err != nil {
		return err
	}
	envFile := filepath.Join(defaultGoctlHome, envFileDir)
	return os.WriteFile(envFile, []byte(strings.Join(goctlEnv.Format(), "\n")), 0o777)
}