github.com/unionj-cloud/go-doudou@v1.3.8-0.20221011095552-0088008e5b31/cmd/internal/svc/codegen/db.go (about)

     1  package codegen
     2  
     3  import (
     4  	"bufio"
     5  	"github.com/sirupsen/logrus"
     6  	"github.com/unionj-cloud/go-doudou/version"
     7  	"os"
     8  	"path/filepath"
     9  	"strings"
    10  	"text/template"
    11  )
    12  
    13  var dbTmpl = `/**
    14  * Generated by go-doudou {{.Version}}.
    15  * You can edit it as your need.
    16  */
    17  package db
    18  
    19  import (
    20  	"{{.ConfigPackage}}"
    21  	"fmt"
    22  	_ "github.com/go-sql-driver/mysql"
    23  	"github.com/iancoleman/strcase"
    24  	"github.com/jmoiron/sqlx"
    25  	"github.com/pkg/errors"
    26  )
    27  
    28  func NewDb(conf config.DbConfig) (*sqlx.DB, error) {
    29  	conn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=%s",
    30  		conf.User,
    31  		conf.Passwd,
    32  		conf.Host,
    33  		conf.Port,
    34  		conf.Schema,
    35  		conf.Charset)
    36  	conn += "&loc=Asia%2FShanghai&parseTime=True"
    37  
    38  	db, err := sqlx.Connect(conf.Driver, conn)
    39  	if err != nil {
    40  		return nil, errors.Wrap(err, "database connection failed")
    41  	}
    42  	db.MapperFunc(strcase.ToSnake)
    43  	return db, nil
    44  }
    45  `
    46  
    47  var MkdirAll = os.MkdirAll
    48  var Open = os.Open
    49  var Create = os.Create
    50  var Stat = os.Stat
    51  
    52  // GenDb generates db connection code
    53  func GenDb(dir string) {
    54  	var (
    55  		err       error
    56  		dbfile    string
    57  		f         *os.File
    58  		tpl       *template.Template
    59  		dbDir     string
    60  		modfile   string
    61  		modName   string
    62  		firstLine string
    63  	)
    64  	dbDir = filepath.Join(dir, "db")
    65  	if err = MkdirAll(dbDir, os.ModePerm); err != nil {
    66  		panic(err)
    67  	}
    68  
    69  	dbfile = filepath.Join(dbDir, "db.go")
    70  	if _, err = Stat(dbfile); os.IsNotExist(err) {
    71  		modfile = filepath.Join(dir, "go.mod")
    72  		if f, err = Open(modfile); err != nil {
    73  			panic(err)
    74  		}
    75  		reader := bufio.NewReader(f)
    76  		firstLine, _ = reader.ReadString('\n')
    77  		modName = strings.TrimSpace(strings.TrimPrefix(firstLine, "module"))
    78  
    79  		if f, err = Create(dbfile); err != nil {
    80  			panic(err)
    81  		}
    82  		defer f.Close()
    83  
    84  		if tpl, err = template.New("db.go.tmpl").Parse(dbTmpl); err != nil {
    85  			panic(err)
    86  		}
    87  		_ = tpl.Execute(f, struct {
    88  			ConfigPackage string
    89  			Version       string
    90  		}{
    91  			ConfigPackage: modName + "/config",
    92  			Version:       version.Release,
    93  		})
    94  	} else {
    95  		logrus.Warnf("file %s already exists", dbfile)
    96  	}
    97  }