package main import ( "fmt" "go/ast" "go/token" "golang.org/x/tools/go/packages" "log" "mongo.games.com/mongoctl/template" "os" "path/filepath" "reflect" "strings" ) const ( symbolBacktick = "`" ) const ( symbolBacktickKey = "SymbolBacktick" ) const ( varPackagesKey = "VarPackages" varModelClassNameKey = "VarModelClassName" varModelPackageNameKey = "VarModelPackageName" varModelPackagePathKey = "VarModelPackagePath" varModelVariableNameKey = "VarModelVariableName" varModelColumnsDefineKey = "VarModelColumnsDefine" varModelColumnsInstanceKey = "VarModelColumnsInstance" varDaoClassNameKey = "VarDaoClassName" varDaoVariableNameKey = "VarDaoVariableName" varDaoPackageNameKey = "VarDaoPackageName" varDaoPackagePathKey = "VarDaoPackagePath" varDaoPrefixNameKey = "VarDaoPrefixName" varCollectionNameKey = "VarCollectionName" varAutofillCodeKey = "VarAutofillCode" ) const defaultCounterName = "Counter" type options struct { modelDir string // modelPkgPath string // modelPkgAlias string // modelNames []string // daoDir string daoPkgPath string subPkgEnable bool subPkgStyle style counterName string fileNameStyle style } type generator struct { opts *options counter *counter modelNames map[string]struct{} } func newGenerator(opts *options) *generator { modelNames := make(map[string]struct{}, len(opts.modelNames)) for _, modelName := range opts.modelNames { if isExportable(modelName) { // 是否导出字段 modelNames[modelName] = struct{}{} } } if len(modelNames) == 0 { log.Fatalf("error: %d model type names found", len(modelNames)) } if opts.counterName == "" { opts.counterName = defaultCounterName } return &generator{ opts: opts, counter: newCounter(opts), modelNames: modelNames, } } func (g *generator) makeDao() { models := g.parseModels() for _, m := range models { g.makeModelInternalDao(m) g.makeModelExternalDao(m) fmt.Printf("%s's dao file generated successfully\n", m.modelName) if !m.isDependCounter { continue } g.makeCounterInternalDao() g.makeCounterExternalDao() } } // generate an internal dao file based on model func (g *generator) makeModelInternalDao(m *model) { replaces := make(map[string]string) replaces[varModelClassNameKey] = m.modelClassName replaces[varModelPackageNameKey] = m.modelPkgName replaces[varModelPackagePathKey] = m.modelPkgPath replaces[varModelVariableNameKey] = m.modelVariableName replaces[varDaoPrefixNameKey] = m.daoPrefixName replaces[varDaoClassNameKey] = m.daoClassName replaces[varDaoVariableNameKey] = m.daoVariableName replaces[varCollectionNameKey] = m.collectionName replaces[varModelColumnsDefineKey] = m.modelColumnsDefined() replaces[varModelColumnsInstanceKey] = m.modelColumnsInstance() replaces[varAutofillCodeKey] = m.autoFillCode() replaces[varPackagesKey] = m.packages() file := m.daoOutputDir + "/internal/" + m.daoOutputFile err := doWrite(file, template.InternalTemplate, replaces) if err != nil { log.Fatal(err) } } // generate an external dao file based on model func (g *generator) makeModelExternalDao(m *model) { file := m.daoOutputDir + "/" + m.daoOutputFile _, err := os.Stat(file) if err != nil { switch { case os.IsNotExist(err): // ignore case os.IsExist(err): return default: log.Fatal(err) } } else { return } replaces := make(map[string]string) replaces[varDaoClassNameKey] = m.daoClassName replaces[varDaoPrefixNameKey] = m.daoPrefixName replaces[varDaoPackageNameKey] = m.daoPkgName replaces[varDaoPackagePathKey] = m.daoPkgPath err = doWrite(file, template.ExternalTemplate, replaces) if err != nil { log.Fatal(err) } } // generate an internal dao file based on counter model func (g *generator) makeCounterInternalDao() { replaces := make(map[string]string) replaces[varDaoClassNameKey] = g.counter.daoClassName replaces[varDaoPrefixNameKey] = g.counter.daoPrefixName replaces[varDaoVariableNameKey] = g.counter.daoVariableName replaces[varCollectionNameKey] = g.counter.collectionName replaces[symbolBacktickKey] = symbolBacktick file := g.counter.daoOutputDir + "/internal/" + g.counter.daoOutputFile err := doWrite(file, template.CounterInternalTemplate, replaces) if err != nil { log.Fatal(err) } } // generate an external dao file based on counter model func (g *generator) makeCounterExternalDao() { file := g.counter.daoOutputDir + "/" + g.counter.daoOutputFile _, err := os.Stat(file) if err != nil { switch { case os.IsNotExist(err): // ignore case os.IsExist(err): return default: log.Fatal(err) } } else { return } replaces := make(map[string]string) replaces[varDaoClassNameKey] = g.counter.daoClassName replaces[varDaoPackageNameKey] = g.counter.daoPkgName replaces[varDaoPackagePathKey] = g.counter.daoPkgPath err = doWrite(file, template.CounterExternalTemplate, replaces) if err != nil { log.Fatal(err) } } // parse multiple models from the go file func (g *generator) parseModels() []*model { var ( pkg = g.loadPackage() models = make([]*model, 0, len(pkg.Syntax)) daoPkgPath = g.opts.daoPkgPath modelPkgPath = g.opts.modelPkgPath modelPkgName = g.opts.modelPkgAlias ) if g.opts.daoPkgPath == "" && pkg.Module != nil { outPath, err := filepath.Abs(g.opts.daoDir) if err != nil { log.Fatal(err) } daoPkgPath = pkg.Module.Path + outPath[len(pkg.Module.Dir):] } daoPkgPath = strings.ReplaceAll(daoPkgPath, `\`, `/`) g.counter.setDaoPkgPath(daoPkgPath) for _, file := range pkg.Syntax { if g.opts.modelPkgPath == "" && pkg.Module != nil && pkg.Fset != nil { filePath := filepath.Dir(pkg.Fset.Position(file.Package).Filename) modelPkgPath = pkg.Module.Path + filePath[len(pkg.Module.Dir):] } modelPkgPath = strings.ReplaceAll(modelPkgPath, `\`, `/`) modelPkgName = file.Name.Name ast.Inspect(file, func(node ast.Node) bool { decl, ok := node.(*ast.GenDecl) if !ok || decl.Tok != token.TYPE { return true } for _, s := range decl.Specs { spec, ok := s.(*ast.TypeSpec) if !ok { continue } _, ok = g.modelNames[spec.Name.Name] if !ok { continue } st, ok := spec.Type.(*ast.StructType) if !ok { continue } model := newModel(g.opts) model.setModelName(spec.Name.Name) model.setModelPkg(modelPkgName, modelPkgPath) model.setDaoPkgPath(daoPkgPath) for _, item := range st.Fields.List { name := item.Names[0].Name if !isExportable(name) { continue } field := &field{name: name, column: name} if item.Tag != nil && len(item.Tag.Value) > 2 { runes := []rune(item.Tag.Value) if runes[0] != '`' || runes[len(runes)-1] != '`' { continue } tag := reflect.StructTag(runes[1 : len(runes)-1]) if column := tag.Get("bson"); column != "" { field.column = column } val, ok := tag.Lookup("gen") if ok { parts := strings.Split(val, ";") for _, part := range parts { if part == "" { continue } switch eles := strings.SplitN(part, ":", 2); eles[0] { case "autoFill": expr, ok := item.Type.(*ast.SelectorExpr) if !ok { continue } switch fmt.Sprintf("%s.%s", expr.X.(*ast.Ident).Name, expr.Sel.Name) { case "primitive.ObjectID": field.autoFill = objectID model.addImport(pkg3) case "primitive.DateTime": field.autoFill = dateTime model.addImport(pkg1) model.addImport(pkg3) } case "autoIncr": if len(eles) != 2 || eles[1] == "" { continue } expr, ok := item.Type.(*ast.Ident) if !ok { continue } switch expr.Name { case "int", "int8", "int16", "int32", "int64", "uint", "uint8", "uint16", "uint32", "uint64": field.autoFill = autoIncr field.autoIncrFieldName = eles[1] if g.opts.subPkgEnable { model.addImport(g.counter.daoPkgPath) } switch expr.Name { case "int": field.autoIncrFieldKind = reflect.Int case "int8": field.autoIncrFieldKind = reflect.Int8 case "int16": field.autoIncrFieldKind = reflect.Int16 case "int32": field.autoIncrFieldKind = reflect.Int32 case "int64": field.autoIncrFieldKind = reflect.Int64 case "uint": field.autoIncrFieldKind = reflect.Uint case "uint8": field.autoIncrFieldKind = reflect.Uint8 case "uint16": field.autoIncrFieldKind = reflect.Uint16 case "uint32": field.autoIncrFieldKind = reflect.Uint32 case "uint64": field.autoIncrFieldKind = reflect.Uint64 } } } } } } if item.Doc != nil { field.documents = make([]string, 0, len(item.Doc.List)) for _, doc := range item.Doc.List { field.documents = append(field.documents, doc.Text) } } if item.Comment != nil { field.comment = item.Comment.List[0].Text } model.addFields(field) } models = append(models, model) } return true }) } return models } func (g *generator) loadPackage() *packages.Package { cfg := &packages.Config{ Mode: packages.NeedName | packages.NeedFiles | packages.NeedCompiledGoFiles | packages.NeedImports | packages.NeedTypes | packages.NeedTypesSizes | packages.NeedSyntax | packages.NeedTypesInfo | packages.NeedModule, Tests: false, } pkgs, err := packages.Load(cfg, g.opts.modelDir) if err != nil { log.Fatal(err) } if len(pkgs) != 1 { log.Fatalf("error: %d packages found", len(pkgs)) } return pkgs[0] }