github.com/cloudwego/kitex@v0.9.0/tool/internal_pkg/pluginmode/thriftgo/plugin.go (about)

     1  // Copyright 2021 CloudWeGo Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //   http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package thriftgo
    16  
    17  import (
    18  	"errors"
    19  	"fmt"
    20  	"os"
    21  
    22  	"github.com/cloudwego/thriftgo/plugin"
    23  
    24  	"github.com/cloudwego/kitex/tool/internal_pkg/generator"
    25  	"github.com/cloudwego/kitex/tool/internal_pkg/util"
    26  )
    27  
    28  // PluginName is the link name when the kitex binary is used as a plugin for thriftgo.
    29  const PluginName = "thrift-gen-kitex"
    30  
    31  // TheUseOptionMessage indicates that the generating of kitex_gen is aborted due to the -use option.
    32  const TheUseOptionMessage = "kitex_gen is not generated due to the -use option"
    33  
    34  // Run is an entry of the plugin mode of kitex for thriftgo.
    35  // It reads a plugin request from the standard input and writes out a response.
    36  func Run() int {
    37  	data, err := util.ReadInput()
    38  	if err != nil {
    39  		println("Failed to get input:", err.Error())
    40  		return 1
    41  	}
    42  
    43  	req, err := plugin.UnmarshalRequest(data)
    44  	if err != nil {
    45  		println("Failed to unmarshal request:", err.Error())
    46  		return 1
    47  	}
    48  
    49  	var conv converter
    50  	if err = conv.init(req); err != nil {
    51  		return conv.fail(err)
    52  	}
    53  
    54  	if err = conv.convertTypes(req); err != nil {
    55  		return conv.fail(err)
    56  	}
    57  
    58  	conv.fixImportConflicts()
    59  
    60  	var files []*generator.File
    61  	gen := generator.NewGenerator(&conv.Config, nil)
    62  
    63  	conv.Package.IDLName = util.IDLName(conv.Config.IDL)
    64  
    65  	if conv.Config.Use == "" {
    66  		for _, s := range conv.Services {
    67  			conv.Package.Dependencies = make(map[string]string)
    68  			conv.Package.ServiceInfo = s
    69  			conv.Package.Namespace = conv.svc2ast[s].GetNamespaceOrReferenceName("go")
    70  			fs, err := gen.GenerateService(&conv.Package)
    71  			if err != nil {
    72  				return conv.fail(err)
    73  			}
    74  			files = append(files, fs...)
    75  		}
    76  	}
    77  
    78  	if conv.Config.GenerateMain {
    79  		if len(conv.Services) == 0 {
    80  			return conv.fail(errors.New("no service defined in the IDL"))
    81  		}
    82  		conv.Package.ServiceInfo = conv.Services[len(conv.Services)-1]
    83  		fs, err := gen.GenerateMainPackage(&conv.Package)
    84  		if err != nil {
    85  			return conv.fail(err)
    86  		}
    87  		files = append(files, fs...)
    88  	}
    89  
    90  	if conv.Config.TemplateDir != "" {
    91  		if len(conv.Services) == 0 {
    92  			return conv.fail(errors.New("no service defined in the IDL"))
    93  		}
    94  		conv.Package.ServiceInfo = conv.Services[len(conv.Services)-1]
    95  		fs, err := gen.GenerateCustomPackage(&conv.Package)
    96  		if err != nil {
    97  			return conv.fail(err)
    98  		}
    99  		files = append(files, fs...)
   100  	}
   101  
   102  	res := &plugin.Response{
   103  		Warnings: conv.Warnings,
   104  	}
   105  	for _, f := range files {
   106  		res.Contents = append(res.Contents, &plugin.Generated{
   107  			Name:    &f.Name,
   108  			Content: f.Content,
   109  		})
   110  	}
   111  
   112  	if conv.Config.Use != "" {
   113  		err = conv.persist(res)
   114  		if err == nil {
   115  			err = errors.New(TheUseOptionMessage)
   116  		}
   117  		return conv.fail(err)
   118  	}
   119  	p := &patcher{
   120  		noFastAPI:             conv.Config.NoFastAPI,
   121  		utils:                 conv.Utils,
   122  		module:                conv.Config.ModuleName,
   123  		copyIDL:               conv.Config.CopyIDL,
   124  		version:               conv.Config.Version,
   125  		record:                conv.Config.Record,
   126  		recordCmd:             conv.Config.RecordCmd,
   127  		deepCopyAPI:           conv.Config.DeepCopyAPI,
   128  		protocol:              conv.Config.Protocol,
   129  		handlerReturnKeepResp: conv.Config.HandlerReturnKeepResp,
   130  	}
   131  	patches, err := p.patch(req)
   132  	if err != nil {
   133  		return conv.fail(fmt.Errorf("patch: %w", err))
   134  	}
   135  	res.Contents = append(res.Contents, patches...)
   136  
   137  	return exit(res)
   138  }
   139  
   140  func exit(res *plugin.Response) int {
   141  	data, err := plugin.MarshalResponse(res)
   142  	if err != nil {
   143  		println("Failed to marshal response:", err.Error())
   144  		return 1
   145  	}
   146  	_, err = os.Stdout.Write(data)
   147  	if err != nil {
   148  		println("Error at writing response out:", err.Error())
   149  		return 1
   150  	}
   151  	return 0
   152  }