gr_hz/generator/layout.go

/*
 * Copyright 2022 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package generator

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"reflect"
	"strings"

	"gopkg.in/yaml.v2"
	"gr_hz/meta"
	"gr_hz/util"
)

// Layout contains the basic information of idl
type Layout struct {
	OutDir          string
	GoModule        string
	ServiceName     string
	UseApacheThrift bool
	HasIdl          bool
	NeedGoMod       bool
	ModelDir        string
	HandlerDir      string
	RouterDir       string
}

// LayoutGenerator contains the information generated by generating the layout template
type LayoutGenerator struct {
	ConfigPath string
	TemplateGenerator
}

var (
	layoutConfig  = defaultLayoutConfig
	packageConfig = defaultPkgConfig
)

func SetDefaultTemplateConfig() {
	layoutConfig = defaultLayoutConfig
	packageConfig = defaultPkgConfig
}

func (lg *LayoutGenerator) Init() error {
	config := layoutConfig
	// unmarshal from user-defined config file if it exists
	if lg.ConfigPath != "" {
		cdata, err := ioutil.ReadFile(lg.ConfigPath)
		if err != nil {
			return fmt.Errorf("read layout config from  %s failed, err: %v", lg.ConfigPath, err.Error())
		}
		config = TemplateConfig{}
		if err = yaml.Unmarshal(cdata, &config); err != nil {
			return fmt.Errorf("unmarshal layout config failed, err: %v", err.Error())
		}
	}

	if reflect.DeepEqual(config, TemplateConfig{}) {
		return errors.New("empty config")
	}
	lg.Config = &config

	return lg.TemplateGenerator.Init()
}

// checkInited initialize template definition
func (lg *LayoutGenerator) checkInited() error {
	if lg.tpls == nil || lg.dirs == nil {
		if err := lg.Init(); err != nil {
			return fmt.Errorf("init layout config failed, err: %v", err.Error())
		}
	}
	return nil
}

func (lg *LayoutGenerator) Generate(data map[string]interface{}) error {
	if err := lg.checkInited(); err != nil {
		return err
	}
	return lg.TemplateGenerator.Generate(data, "", "", false)
}

func (lg *LayoutGenerator) GenerateByService(service Layout) error {
	if err := lg.checkInited(); err != nil {
		return err
	}

	if len(service.HandlerDir) != 0 {
		// override the default "biz/handler/ping.go" to "HANDLER_DIR/ping.go"
		defaultPingDir := defaultHandlerDir + sp + "ping.go"
		if tpl, exist := lg.tpls[defaultPingDir]; exist {
			delete(lg.tpls, defaultPingDir)
			newPingDir := filepath.Clean(service.HandlerDir + sp + "ping.go")
			lg.tpls[newPingDir] = tpl
		}
	}

	if len(service.RouterDir) != 0 {
		defaultRegisterDir := defaultRouterDir + sp + registerTplName
		if tpl, exist := lg.tpls[defaultRegisterDir]; exist {
			delete(lg.tpls, defaultRegisterDir)
			newRegisterDir := filepath.Clean(service.RouterDir + sp + registerTplName)
			lg.tpls[newRegisterDir] = tpl
		}
	}

	if !service.NeedGoMod {
		gomodFile := "go.mod"
		if _, exist := lg.tpls[gomodFile]; exist {
			delete(lg.tpls, gomodFile)
		}
	}

	if util.IsWindows() {
		buildSh := "build.sh"
		bootstrapSh := defaultScriptDir + sp + "bootstrap.sh"
		if _, exist := lg.tpls[buildSh]; exist {
			delete(lg.tpls, buildSh)
		}
		if _, exist := lg.tpls[bootstrapSh]; exist {
			delete(lg.tpls, bootstrapSh)
		}
	}

	sd, err := serviceToLayoutData(service)
	if err != nil {
		return err
	}

	rd, err := serviceToRouterData(service)
	if err != nil {
		return err
	}
	if service.HasIdl {
		for k := range lg.tpls {
			if strings.Contains(k, registerTplName) {
				delete(lg.tpls, k)
				break
			}
		}
	}

	data := map[string]interface{}{
		"*":                                    sd,
		layoutConfig.Layouts[routerIndex].Path: rd, // router.go
		layoutConfig.Layouts[routerGenIndex].Path: rd, // router_gen.go
	}

	return lg.Generate(data)
}

// serviceToLayoutData stores go mod, serviceName, UseApacheThrift mapping
func serviceToLayoutData(service Layout) (map[string]interface{}, error) {
	goMod := service.GoModule
	if goMod == "" {
		return nil, errors.New("please specify go-module")
	}
	handlerPkg := filepath.Base(defaultHandlerDir)
	if len(service.HandlerDir) != 0 {
		handlerPkg = filepath.Base(service.HandlerDir)
	}
	routerPkg := filepath.Base(defaultRouterDir)
	if len(service.RouterDir) != 0 {
		routerPkg = filepath.Base(service.RouterDir)
	}
	serviceName := service.ServiceName
	if len(serviceName) == 0 {
		serviceName = meta.DefaultServiceName
	}

	return map[string]interface{}{
		"GoModule":        goMod,
		"ServiceName":     serviceName,
		"UseApacheThrift": service.UseApacheThrift,
		"HandlerPkg":      handlerPkg,
		"RouterPkg":       routerPkg,
	}, nil
}

// serviceToRouterData stores the registers function, router import path, handler import path
func serviceToRouterData(service Layout) (map[string]interface{}, error) {
	routerDir := sp + defaultRouterDir
	handlerDir := sp + defaultHandlerDir
	if len(service.RouterDir) != 0 {
		routerDir = sp + service.RouterDir
	}
	if len(service.HandlerDir) != 0 {
		handlerDir = sp + service.HandlerDir
	}
	return map[string]interface{}{
		"Registers":      []string{},
		"RouterPkgPath":  service.GoModule + util.PathToImport(routerDir, ""),
		"HandlerPkgPath": service.GoModule + util.PathToImport(handlerDir, ""),
	}, nil
}

func (lg *LayoutGenerator) GenerateByConfig(configPath string) error {
	if err := lg.checkInited(); err != nil {
		return err
	}
	buf, err := ioutil.ReadFile(configPath)
	if err != nil {
		return fmt.Errorf("read data file '%s' failed, err: %v", configPath, err.Error())
	}
	var data map[string]interface{}
	if err := json.Unmarshal(buf, &data); err != nil {
		return fmt.Errorf("unmarshal json data failed, err: %v", err.Error())
	}
	return lg.Generate(data)
}

func (lg *LayoutGenerator) Degenerate() error {
	return lg.TemplateGenerator.Degenerate()
}