github.com/mstephano/gqlgen-schemagen@v0.0.0-20230113041936-dd2cd4ea46aa/plugin/schemagen/schema.go

package schemagen

import (
	"bytes"
	// Embedded file
	_ "embed"
	"fmt"
	"go/constant"
	"go/types"
	"html/template"
	"log"
	"os"
	"reflect"
	"regexp"
	"sort"
	"strings"

	"github.com/mstephano/gqlgen-schemagen/codegen"
	"github.com/mstephano/gqlgen-schemagen/codegen/config"
	"github.com/mstephano/gqlgen-schemagen/internal/code"
	"github.com/mstephano/gqlgen-schemagen/plugin"
	"github.com/vektah/gqlparser/v2/ast"
	"github.com/vektah/gqlparser/v2/parser"
	"golang.org/x/tools/go/packages"
)

//go:embed schema.gotpl
var fileTemplate string

const basePath = "graph/schema/"

type addingModelType int

const (
	composition addingModelType = iota
	reference
)

type generator struct {
	config                *config.Config
	filePath              string
	mapDataTypes          map[string]string
	mapDataTypeExclusions map[string]string
	modelObjects          map[string]*modelObject
	modelEnums            map[string]*enumObject
	excludedModelObjects  map[string]*types.Object
	packages              []*packages.Package
}

type modelObject struct {
	Model        types.Object
	Fields       map[string]*fieldObject
	Compositions map[string]struct{}
	References   map[string]struct{}
}

type enumObject struct {
	Model  types.Object
	Values []string
}

type fieldObject struct {
	Field    *types.Var
	TagValue string
	TypeName string
	Required bool
}

type templateData struct {
	SortedEnumsSlice   []*enumObject
	SortedObjectsSlice []*modelObject
}

// New creates a new schemagen plugin
func New(cfg *config.Config, fileName string, mapDataTypes map[string]string, mapDataTypeExclusions map[string]string) plugin.Plugin {
	return &generator{
		config:                cfg,
		filePath:              getFilePath(fileName),
		mapDataTypes:          mapDataTypes,
		mapDataTypeExclusions: mapDataTypeExclusions,
		modelObjects:          make(map[string]*modelObject, 0),
		modelEnums:            make(map[string]*enumObject, 0),
		excludedModelObjects:  make(map[string]*types.Object, 0),
		packages:              make([]*packages.Package, 0),
	}
}

func (g *generator) Name() string {
	return "schemagen"
}

func (g *generator) InjectSourceEarly() *ast.Source {
	g.deleteSource()

	pkgs := &code.Packages{}
	g.packages = pkgs.LoadAll(g.config.AutoBind...)

	// First level models from schemaTypes
	astSchemaDoc, err := parser.ParseSchemas(g.config.Sources...)
	if err != nil {
		panic(err)
	}

	schemaTypeNames := make([]string, 0)
	for _, def := range astSchemaDoc.Definitions {
		for _, field := range def.Fields {
			if field.Type.Elem != nil {
				schemaTypeNames = append(schemaTypeNames, field.Type.Elem.NamedType)
			} else if field.Type.NamedType != "" {
				schemaTypeNames = append(schemaTypeNames, field.Type.NamedType)
			}
		}
	}

	for _, p := range g.packages {
		for _, schemaTypeName := range schemaTypeNames {
			if schemaType := p.Types.Scope().Lookup(schemaTypeName); schemaType != nil {
				if validateType(p, schemaType) {
					if g.addModel(schemaType, schemaType, reference) {
						fmt.Printf("adding schemaType: %+v\n", schemaTypeName)
					}
				} else {
					fmt.Printf("ignoring schemaType: %+v, %+v\n", p.PkgPath, schemaTypeName)
				}
			}
		}
	}

	// Add all models from all packages
	loop := 0
	for {
		loop++
		modelAdded := 0

		// Add references & composition
		for _, mo := range g.modelObjects {
			switch typ := mo.Model.Type().Underlying().(type) {
			case *types.Struct:
				for i := 0; i < typ.NumFields(); i++ {
					field := typ.Field(i)
					tagValue, _, hidden := getJSONTagValue(typ.Tag(i))
					if hidden {
						continue
					}

					if g.preAddModel(mo, field, tagValue) {
						modelAdded++
					}
				}
			case *types.Basic: // enum
				if g.addEnum(mo) {
					modelAdded++
				}
			default:
				fmt.Printf("Not supported type: %+v, %+v, %+v\n", mo.Model.Name(), typ.String(), reflect.TypeOf(mo.Model.Type().Underlying()))
			}
		}
		fmt.Printf("Loop # %+v - models added: %+v\n", loop, modelAdded)

		if modelAdded == 0 {
			break
		}
	}

	// Add all fields for every modelObject
	modelObjects := make(map[string]*modelObject, 0) // keep single model by name, models with same name will have their fields consolidated
	for _, mo := range g.modelObjects {
		moToAddFields := mo
		if obj, ok := modelObjects[mo.Model.Name()]; ok {
			moToAddFields = obj // when same name exist, consolidate all the fields to the same object

			// keep the one with references
			if len(obj.References) == 0 && len(mo.References) > 0 {
				mo.Fields = obj.Fields
				moToAddFields = mo
				modelObjects[mo.Model.Name()] = mo
			}
		} else {
			modelObjects[mo.Model.Name()] = mo
		}

		g.addFieldsToModelObject(mo.Model.Type().Underlying(), *moToAddFields, false)
	}

	// Complete composition with missing fields
	loop = 0
	for {
		loop++
		fieldAdded := 0
		for _, v := range modelObjects {
			for c := range v.Compositions {
				if obj, exist := modelObjects[c]; exist {
					for _, f := range v.Fields {
						if _, fieldExist := obj.Fields[f.TagValue]; !fieldExist {
							obj.Fields[f.TagValue] = &fieldObject{
								Field:    f.Field,
								TagValue: f.TagValue,
								TypeName: f.TypeName,
								Required: f.Required,
							}

							fieldAdded++
						}
					}
				}
			}
		}
		fmt.Printf("Loop # %+v - fields added: %+v\n", loop, fieldAdded)

		if fieldAdded == 0 {
			break
		}
	}

	// Sort modelObjects
	modelObjectsSlice := make([]*modelObject, 0, len(modelObjects))
	for _, v := range modelObjects {
		modelObjectsSlice = append(modelObjectsSlice, v)
	}
	sort.Slice(modelObjectsSlice, func(i, j int) bool {
		return modelObjectsSlice[i].Model.Name() < modelObjectsSlice[j].Model.Name()
	})
	// Sort modelObjects fields
	for _, mo := range modelObjectsSlice {
		fields := make(map[string]*fieldObject, len(mo.Fields))
		fieldKeys := make([]string, 0, len(mo.Fields))
		for k := range mo.Fields {
			fieldKeys = append(fieldKeys, k)
		}
		sort.Strings((fieldKeys))

		for _, fieldKey := range fieldKeys {
			field := mo.Fields[fieldKey]
			fields[fieldKey] = field
		}
		mo.Fields = fields
	}

	// Sort modelEnums
	modelEnumsSlice := make([]*enumObject, 0, len(g.modelEnums))
	for _, v := range g.modelEnums {
		modelEnumsSlice = append(modelEnumsSlice, v)
	}
	sort.Slice(modelEnumsSlice, func(i, j int) bool {
		return modelEnumsSlice[i].Model.Name() < modelEnumsSlice[j].Model.Name()
	})
	// Sort modelEnums values
	for _, v := range modelEnumsSlice {
		sort.Strings(v.Values)
	}

	// Create input
	input := g.renderTemplate(&templateData{
		SortedEnumsSlice:   modelEnumsSlice,
		SortedObjectsSlice: modelObjectsSlice,
	})

	// Create source
	fmt.Printf("Generating schema file ...\n")
	source := &ast.Source{
		Name:    g.filePath,
		BuiltIn: false,
		Input:   input,
	}
	// fmt.Printf("%s\n", source.Input)

	return source
}

func (g *generator) GenerateCode(data *codegen.Data) error {
	for _, s := range data.Config.Sources {
		if s.Name == g.filePath {
			f, err := os.OpenFile(g.filePath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0o755)
			if err != nil {
				log.Fatal(err)
			}
			defer f.Close()
			f.WriteString(s.Input)
			fmt.Printf("Generated schema file location: %s", g.filePath)
			break
		}
	}

	return nil
}

func (g *generator) preAddModel(mo *modelObject, field *types.Var, tagValue string) bool {
	addModelType := reference
	addModelTypeName := "reference"
	if tagValue == "" {
		addModelType = composition
		addModelTypeName = "composition"
	}

	pkgPath, fullTypeName, baseTypeName, typeName, isBasic := g.getTypeName(field.Name(), field.Type(), false)
	if addModelType == reference && isBasic {
		return false
	}

	added := false
	exist := false
	for _, p := range g.packages {
		if p.PkgPath != pkgPath {
			continue
		}

		if modelType := p.Types.Scope().Lookup(baseTypeName); modelType != nil {
			if !exist {
				if g.addModel(modelType, mo.Model, addModelType) {
					fmt.Printf("adding %s: %+v\n", addModelTypeName, modelType.Name())
					added = true
				}
				exist = true
			} else {
				fmt.Printf("%s found but already exist: %+v, %+v\n", addModelTypeName, fullTypeName, modelType.Name())
			}
		}
	}
	if !added && !exist && !strings.HasPrefix(typeName, "Map") {
		fmt.Printf("missing %s - model: %+v - field: %+v - pkgPath: %+v\n", addModelTypeName, mo.Model.Name(), field.Name(), pkgPath)
	}

	return added
}

func (g *generator) addEnum(mo *modelObject) bool {
	addModelTypeName := "enum"

	pkgPath, fullTypeName, baseTypeName, typeName, isBasic := g.getEnumName(mo.Model.Type(), false)
	if isBasic {
		return false
	}

	added := false
	exist := false
	for _, p := range g.packages {
		if p.PkgPath != pkgPath {
			continue
		}

		if modelType := p.Types.Scope().Lookup(baseTypeName); modelType != nil {
			if !exist {
				if _, ok := g.modelEnums[typeName]; !ok {
					g.modelEnums[typeName] = &enumObject{
						Model:  mo.Model,
						Values: make([]string, 0),
					}
					fmt.Printf("adding %s: %+v\n", addModelTypeName, modelType.Name())
					added = true

					for _, n := range p.Types.Scope().Names() {
						if o := p.Types.Scope().Lookup(n); modelType != nil {
							if typ, ok := o.(*types.Const); ok {
								if typ.Type().String() == fullTypeName && typ.Type().Underlying().String() == "string" {
									g.modelEnums[typeName].Values = append(g.modelEnums[typeName].Values, constant.StringVal(typ.Val()))
								}
							}
						}
					}
				}
				exist = true
			} else {
				fmt.Printf("%s found but already exist: %+v, %+v\n", addModelTypeName, fullTypeName, modelType.Name())
			}
		}
	}
	if !added && !exist {
		for k := range mo.References {
			fmt.Printf("missing %s - model: %+v - field: %+v - pkgPath: %+v\n", addModelTypeName, mo.Model.Name(), k, pkgPath)
			break
		}
	}

	return added
}

func (g *generator) addModel(model types.Object, fromModel types.Object, addModelType addingModelType) bool {
	if _, exist := g.mapDataTypeExclusions[model.Pkg().Name()+"."+model.Name()]; exist {
		if _, exist := g.excludedModelObjects[model.Name()]; !exist {
			g.excludedModelObjects[model.Name()] = &model
			fmt.Printf("excluded model: %+v\n", model.Name())
		}
		return false
	}

	added := false
	typeName := fmt.Sprintf("%s.%s", model.Pkg().Path(), model.Name())
	if _, exist := g.modelObjects[typeName]; !exist {
		g.modelObjects[typeName] = &modelObject{
			Model:        model,
			Fields:       make(map[string]*fieldObject, 0),
			Compositions: make(map[string]struct{}, 0),
			References:   make(map[string]struct{}, 0),
		}
		added = true
	}

	mo := g.modelObjects[typeName]
	switch addModelType {
	case composition:
		if _, exist := mo.Compositions[fromModel.Name()]; !exist {
			mo.Compositions[fromModel.Name()] = struct{}{}
		}
	case reference:
		if _, exist := mo.References[fromModel.Name()]; !exist {
			mo.References[fromModel.Name()] = struct{}{}
		}
	}

	return added
}

func (g *generator) getTypeName(fieldName string, obj types.Type, isSlice bool) (string, string, string, string, bool) {
	pkgPath := ""
	fullTypeName := ""
	baseTypeName := ""
	typeName := ""
	isBasic := false

	switch typ := obj.(type) {
	case *types.Basic:
		pkgPath = typ.Underlying().String()
		fullTypeName = typ.Underlying().String() + "." + typ.Name()
		baseTypeName = typ.Name()
		isBasic = true
	case *types.Named:
		if _, ok := typ.Obj().Type().Underlying().(*types.Map); ok {
			pkgPath = fmt.Sprintf("Map%s", typ.Obj().Name())
			fullTypeName = pkgPath
			baseTypeName = pkgPath
		} else {
			pkgPath = typ.Obj().Pkg().Path()
			fullTypeName = typ.Obj().Pkg().Path() + "." + typ.Obj().Name()
			baseTypeName = typ.Obj().Pkg().Name() + "." + typ.Obj().Name()
			if _, ok := g.mapDataTypes[baseTypeName]; !ok {
				baseTypeName = typ.Obj().Name()
				if _, ok := typ.Obj().Type().Underlying().(*types.Slice); ok {
					isSlice = true
				}
			}
		}
	case *types.Pointer:
		pkgPath, fullTypeName, baseTypeName, typeName, isBasic = g.getTypeName(fieldName, typ.Elem(), isSlice)
		return pkgPath, fullTypeName, baseTypeName, typeName, isBasic
	case *types.Slice:
		pkgPath, fullTypeName, baseTypeName, typeName, isBasic = g.getTypeName(fieldName, typ.Elem(), true)
		return pkgPath, fullTypeName, baseTypeName, typeName, isBasic
	case *types.Map:
		pkgPath = fmt.Sprintf("Map%s", typ.Elem().String())
		fullTypeName = pkgPath
		baseTypeName = pkgPath
	default:
		fmt.Printf("type not supported - field: %+v - type: %+v\n", fieldName, reflect.TypeOf(obj))
	}

	if mapName, ok := g.mapDataTypes[baseTypeName]; ok {
		baseTypeName = mapName
		isBasic = true
	}
	if isSlice {
		typeName = "[" + baseTypeName + "!]"
	} else {
		typeName = baseTypeName
	}
	return pkgPath, fullTypeName, baseTypeName, typeName, isBasic
}

func (g *generator) getEnumName(obj types.Type, isSlice bool) (string, string, string, string, bool) {
	pkgPath := ""
	fullTypeName := ""
	baseTypeName := ""
	typeName := ""
	isBasic := false

	switch typ := obj.(type) {
	case *types.Named:
		pkgPath = typ.Obj().Pkg().Path()
		fullTypeName = typ.Obj().Pkg().Path() + "." + typ.Obj().Name()
		baseTypeName = typ.Obj().Pkg().Name() + "." + typ.Obj().Name()
		if _, ok := g.mapDataTypes[baseTypeName]; !ok {
			baseTypeName = typ.Obj().Name()
			if _, ok := typ.Obj().Type().Underlying().(*types.Slice); ok {
				isSlice = true
			}
		}
	default:
		fmt.Printf("type not supported - field: %+v - type: %+v\n", obj.String(), reflect.TypeOf(obj))
	}

	if mapName, ok := g.mapDataTypes[baseTypeName]; ok {
		baseTypeName = mapName
		isBasic = true
	}
	if isSlice {
		typeName = "[" + baseTypeName + "!]"
	} else {
		typeName = baseTypeName
	}
	return pkgPath, fullTypeName, baseTypeName, typeName, isBasic
}

func (g *generator) addFieldsToModelObject(obj types.Type, modelObject modelObject, isSlice bool) {
	switch typ := obj.(type) {
	case *types.Struct:
		for i := 0; i < typ.NumFields(); i++ {
			field := typ.Field(i)
			tagValue, required, hidden := getJSONTagValue(typ.Tag(i))
			if hidden {
				continue
			}

			_, fullTypeName, baseTypeName, typeName, isBasic := g.getTypeName(field.Name(), field.Type(), isSlice)
			if tagValue != "" {
				if isBasic {
					if baseTypeName != "" {
						modelObject.addField(field, tagValue, typeName, required)
					}
				} else {
					_, ok := g.modelObjects[fullTypeName]
					if ok || strings.HasPrefix(typeName, "Map") {
						modelObject.addField(field, tagValue, typeName, required)
					}
				}
			} else {
				// first level composition
				if o, ok := g.modelObjects[fullTypeName]; ok {
					if u, ok := o.Model.Type().Underlying().(*types.Struct); ok {
						for i := 0; i < u.NumFields(); i++ {
							field := u.Field(i)
							tagValue, required, hidden := getJSONTagValue(u.Tag(i))
							_, _, baseTypeName, typeName, isBasic := g.getTypeName(field.Name(), field.Type(), isSlice)
							if tagValue != "" && !hidden && isBasic && baseTypeName != "" {
								modelObject.addField(field, tagValue, typeName, required)
							}
						}
					}
				}
			}
		}
	case *types.Slice:
		g.addFieldsToModelObject(typ.Elem().Underlying(), modelObject, true)
		return
	case *types.Pointer:
		g.addFieldsToModelObject(typ.Elem().Underlying(), modelObject, false)
		return
	case *types.Basic: // enum
		return // managed in addEnum()
	default:
		fmt.Printf("field not supported - model: %s - field: %s - type: %+v\n", modelObject.Model.Name(), obj.String(), reflect.TypeOf(obj))
	}
}

func (g *generator) deleteSource() {
	// Delete existing source to avoid type collisions
	for i, s := range g.config.Sources {
		if s.Name == g.filePath {
			g.config.Sources = append(g.config.Sources[:i], g.config.Sources[i+1:]...)
			break
		}
	}
}

func (g *generator) renderTemplate(templateData *templateData) string {
	var buf bytes.Buffer
	if err := template.Must(template.New("test.graphqls").Funcs(template.FuncMap{
		"GetSchemaField": GetSchemaFieldGotpl,
	}).Parse(fileTemplate)).Execute(&buf, templateData); err != nil {
		panic(err)
	}
	return buf.String()
}

func (o *modelObject) addField(field *types.Var, tagValue, typeName string, required bool) bool {
	if _, exist := o.Fields[tagValue]; !exist {
		o.Fields[tagValue] = &fieldObject{
			Field:    field,
			TagValue: tagValue,
			TypeName: typeName,
			Required: required,
		}
		return true
	}
	return false
}

func validateType(pkg *packages.Package, obj types.Object) bool {
	if named, ok := obj.Type().(*types.Named); ok {
		if s, ok := named.Underlying().(*types.Struct); ok {
			for i := 0; i < s.NumFields(); i++ {
				if s.Field(i).Pkg().Path() == pkg.PkgPath {
					return true
				}
			}
		}
	}
	return false
}

func getJSONTagValue(tag string) (string, bool, bool) {
	required := false
	hidden := false
	if tag != "" && strings.Contains(tag, "json:") {
		required := !strings.Contains(tag, "omitempty")
		re := regexp.MustCompile(`json:\"(.*?)\"`)
		res := re.FindAllStringSubmatch(tag, -1)
		propertyName := strings.ReplaceAll(res[0][1], ",omitempty", "")
		propertyName = strings.ReplaceAll(propertyName, ",", "")

		if propertyName == "-" {
			hidden = true
		}
		return propertyName, required, hidden
	}
	return "", required, hidden
}

// GetSchemaFieldGotpl is used in the go template
func GetSchemaFieldGotpl(tagValue, typeName string, required bool) string {
	field := fmt.Sprintf("%s: %s", tagValue, typeName)
	if required {
		field = fmt.Sprintf("%s: %s!", tagValue, typeName)
	}
	return field
}

func getFilePath(filename string) string {
	return fmt.Sprintf("%s%s", basePath, filename)
}