github.com/yoogoc/kratos-scaffold@v0.0.0-20240402032722-a538b3c18955/project_generator/resources/data.ent.data.go.tmpl (about)

     1  {{- /*gotype: github.com/yoogoc/kratos-scaffold/project_generator.DataTmpl*/ -}}
     2  package data
     3  
     4  import (
     5  	"context"
     6  	"{{.AppPkgPath}}/internal/biz"
     7  	"{{.AppPkgPath}}/internal/conf"
     8  	"{{.AppPkgPath}}/internal/data/ent"
     9  
    10  	"github.com/go-kratos/kratos/v2/log"
    11  	"github.com/google/wire"
    12  
    13  	{{- if eq .DBDriverType "mysql"}}
    14  	_ "github.com/go-sql-driver/mysql"
    15  	{{- end}}
    16  	"github.com/pkg/errors"
    17  )
    18  
    19  // ProviderSet is data providers.
    20  var ProviderSet = wire.NewSet(NewEntClient, NewData, NewTxRepo)
    21  
    22  type Data struct {
    23  	log *log.Helper
    24  	db  *ent.Client
    25  }
    26  
    27  func (d *Data) WithEntTx(ctx context.Context, fn func(ctx context.Context) error) error {
    28  	// 如果事务内嵌套事务,则交由上层处理
    29  	if tx := ent.TxFromContext(ctx); tx != nil {
    30  		return fn(ctx)
    31  	}
    32  	tx, err := d.db.Tx(ctx)
    33  	txContext := ent.NewTxContext(ctx, tx)
    34  	if err != nil {
    35  		return err
    36  	}
    37  	defer func() {
    38  		if v := recover(); v != nil {
    39  			_ = tx.Rollback()
    40  			panic(v)
    41  		}
    42  	}()
    43  	if err := fn(txContext); err != nil {
    44  		if rerr := tx.Rollback(); rerr != nil {
    45  			err = errors.Wrapf(err, "rolling back transaction: %v", rerr)
    46  		}
    47  		return err
    48  	}
    49  	if err := tx.Commit(); err != nil {
    50  		return errors.Wrapf(err, "committing transaction: %v", err)
    51  	}
    52  	return nil
    53  }
    54  
    55  func (d *Data) EntClient(ctx context.Context) *ent.Client {
    56  	tx := ent.TxFromContext(ctx)
    57  	if tx != nil {
    58  		return tx.Client()
    59  	}
    60  	return d.db
    61  }
    62  
    63  func NewTxRepo(d *Data) biz.TxRepo {
    64  	return d
    65  }
    66  
    67  func NewEntClient(conf *conf.Data, logger log.Logger) *ent.Client {
    68  	l := log.NewHelper(log.With(logger, "module", "{{.LoggerName}}/ent"))
    69  
    70  	var ops []ent.Option
    71  	if conf.Database.EntDebug {
    72  		ops = append(ops, ent.Debug())
    73  	} else {
    74  		//ops = append(ops, ent.Log(log.Info))
    75  	}
    76  
    77  	client, err := ent.Open(
    78  		conf.Database.Driver,
    79  		conf.Database.Source,
    80  		ops...,
    81  	)
    82  	if err != nil {
    83  		l.Fatalf("failed opening connection to db: %v", err)
    84  	}
    85  	return client
    86  }
    87  
    88  func NewData(conf *conf.Data, logger log.Logger, db *ent.Client) (*Data, error) {
    89  	l := log.NewHelper(log.With(logger, "module", "{{.LoggerName}}"))
    90  	return &Data{
    91  		log: l,
    92  		db:  db,
    93  	}, nil
    94  }