github.com/lastbackend/toolkit@v0.0.0-20241020043710-cafa37b95aad/protoc-gen-toolkit/descriptor/descriptor.go (about)

     1  /*
     2  Copyright [2014] - [2023] The Last.Backend authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package descriptor
    18  
    19  import (
    20  	"fmt"
    21  
    22  	"google.golang.org/protobuf/compiler/protogen"
    23  	"google.golang.org/protobuf/types/descriptorpb"
    24  )
    25  
    26  type Descriptor struct {
    27  	packageAliasesMap map[string]string
    28  	fileMap           map[string]*File
    29  	messageMap        map[string]*Message
    30  	enumMap           map[string]*Enum
    31  	packageMap        map[string]string
    32  }
    33  
    34  func NewDescriptor() *Descriptor {
    35  	return &Descriptor{
    36  		packageAliasesMap: make(map[string]string),
    37  		fileMap:           make(map[string]*File),
    38  		messageMap:        make(map[string]*Message),
    39  		enumMap:           make(map[string]*Enum),
    40  		packageMap:        make(map[string]string),
    41  	}
    42  }
    43  
    44  func (d *Descriptor) LoadFromPlugin(gen *protogen.Plugin) error {
    45  	for filePath, f := range gen.FilesByPath {
    46  		d.loadFile(filePath, f)
    47  	}
    48  
    49  	for filePath, f := range gen.FilesByPath {
    50  		if !f.Generate {
    51  			continue
    52  		}
    53  
    54  		file := d.fileMap[filePath]
    55  
    56  		if err := d.loadServices(file); err != nil {
    57  			return err
    58  		}
    59  	}
    60  
    61  	return nil
    62  }
    63  
    64  func (d *Descriptor) loadFile(filePath string, file *protogen.File) {
    65  	pkg := GoPackage{
    66  		Path: string(file.GoImportPath),
    67  		Name: string(file.GoPackageName),
    68  	}
    69  
    70  	if err := d.ReserveGoPackageAlias(pkg.Name, pkg.Path); err != nil {
    71  		for i := 0; ; i++ {
    72  			alias := fmt.Sprintf("%s_%d", pkg.Name, i)
    73  			if err := d.ReserveGoPackageAlias(alias, pkg.Path); err == nil {
    74  				pkg.Alias = alias
    75  				break
    76  			}
    77  		}
    78  	}
    79  
    80  	f := &File{
    81  		FileDescriptorProto:     file.Proto,
    82  		GoPkg:                   pkg,
    83  		GeneratedFilenamePrefix: file.GeneratedFilenamePrefix,
    84  	}
    85  
    86  	d.fileMap[filePath] = f
    87  
    88  	d.registerMsg(f, nil, file.Proto.MessageType)
    89  	d.registerEnum(f, nil, file.Proto.EnumType)
    90  }
    91  
    92  func (d *Descriptor) ReserveGoPackageAlias(alias, pkgpath string) error {
    93  	if taken, ok := d.packageAliasesMap[alias]; ok {
    94  		if taken == pkgpath {
    95  			return nil
    96  		}
    97  		return fmt.Errorf("package name %s is already taken. Use another alias", alias)
    98  	}
    99  	d.packageAliasesMap[alias] = pkgpath
   100  	return nil
   101  }
   102  
   103  func (d *Descriptor) registerMsg(file *File, outerPath []string, messages []*descriptorpb.DescriptorProto) {
   104  	for index, message := range messages {
   105  		m := &Message{
   106  			File:            file,
   107  			Outers:          outerPath,
   108  			DescriptorProto: message,
   109  			Index:           index,
   110  		}
   111  
   112  		for _, fd := range message.GetField() {
   113  			m.Fields = append(m.Fields, &Field{
   114  				Message:              m,
   115  				FieldDescriptorProto: fd,
   116  			})
   117  		}
   118  
   119  		file.Messages = append(file.Messages, m)
   120  		d.messageMap[m.FullyName()] = m
   121  
   122  		var outers []string
   123  		outers = append(outers, outerPath...)
   124  		outers = append(outers, m.GetName())
   125  
   126  		d.registerMsg(file, outers, m.GetNestedType())
   127  		d.registerEnum(file, outers, m.GetEnumType())
   128  	}
   129  }
   130  
   131  func (d *Descriptor) registerEnum(file *File, outerPath []string, enums []*descriptorpb.EnumDescriptorProto) {
   132  	for index, enum := range enums {
   133  		e := &Enum{
   134  			File:                file,
   135  			Outers:              outerPath,
   136  			EnumDescriptorProto: enum,
   137  			Index:               index,
   138  		}
   139  
   140  		file.Enums = append(file.Enums, e)
   141  		d.enumMap[e.FullyName()] = e
   142  	}
   143  }
   144  
   145  func (d *Descriptor) FindFile(name string) (*File, error) {
   146  	file, ok := d.fileMap[name]
   147  	if !ok {
   148  		return nil, fmt.Errorf("no such file: %s", name)
   149  	}
   150  
   151  	return file, nil
   152  }