kratos/cmd/protoc-gen-go-errors/errors.go

package main

import (
	"strings"

	"github.com/go-kratos/kratos/cmd/protoc-gen-go-errors/v2/errors"
	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/proto"
)

const (
	errorsPackage = protogen.GoImportPath("github.com/go-kratos/kratos/v2/errors")
	fmtPackage    = protogen.GoImportPath("fmt")
)

// generateFile generates a _http.pb.go file containing kratos errors definitions.
func generateFile(gen *protogen.Plugin, file *protogen.File) *protogen.GeneratedFile {
	var hasCode bool
	for _, enum := range file.Enums {
		if code := defaultErrorCode(enum); code > 0 {
			hasCode = true
		}
	}
	if len(file.Enums) == 0 || !hasCode {
		return nil
	}
	filename := file.GeneratedFilenamePrefix + "_errors.pb.go"
	g := gen.NewGeneratedFile(filename, file.GoImportPath)
	g.P("// Code generated by protoc-gen-go-errors. DO NOT EDIT.")
	g.P()
	g.P("package ", file.GoPackageName)
	g.P()
	g.QualifiedGoIdent(fmtPackage.Ident(""))
	generateFileContent(gen, file, g)
	return g
}

// generateFileContent generates the kratos errors definitions, excluding the package statement.
func generateFileContent(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile) {
	if len(file.Enums) == 0 {
		return
	}

	g.P("// This is a compile-time assertion to ensure that this generated file")
	g.P("// is compatible with the kratos package it is being compiled against.")
	g.P("const _ = ", errorsPackage.Ident("SupportPackageIsVersion1"))
	g.P()
	for _, enum := range file.Enums {
		genErrorsReason(gen, file, g, enum)
	}
}

func defaultErrorCode(enum *protogen.Enum) int {
	defaultCode := proto.GetExtension(enum.Desc.Options(), errors.E_DefaultCode)
	if code, ok := defaultCode.(int32); ok && code > 0 {
		return int(code)
	}
	return 0
}

func genErrorsReason(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile, enum *protogen.Enum) {
	var ew errorWrapper
	defCode := defaultErrorCode(enum)
	for _, v := range enum.Values {
		code := int(proto.GetExtension(v.Desc.Options(), errors.E_Code).(int32))
		if code == 0 {
			code = defCode
		}
		if code > 600 || code < 200 {
			panic("httpstatus code must be greater than or equal to 200 and less than 600")
		}
		err := &errorInfo{
			Name:       string(enum.Desc.Name()),
			Value:      string(v.Desc.Name()),
			CamelValue: case2Camel(string(v.Desc.Name())),
			HttpCode:   code,
		}
		ew.Errors = append(ew.Errors, err)
	}
	g.P(ew.execute())
}

func case2Camel(name string) string {
	if !strings.Contains(name, "_") {
		return name
	}
	name = strings.ToLower(name)
	name = strings.Replace(name, "_", " ", -1)
	name = strings.Title(name)
	return strings.Replace(name, " ", "", -1)
}