joybuywork/internal/utils/logic/context.go

package logic

import (
	"context"
	"joybuywork/internal/app"
	"joybuywork/internal/consts"
	"joybuywork/internal/repository/dao"
	"joybuywork/internal/service/provider/traces"
	"joybuywork/internal/utils/merr"
	"fmt"

	"go-micro.dev/v4/metadata"
)

type after struct {
	done    []func()
	errDone []func()
	okDone  []func()
}

func (s *after) AddDone(dc ...func()) {
	s.done = append(s.done, dc...)
}

func (s *after) AddErrDone(dc ...func()) {
	s.errDone = append(s.errDone, dc...)
}

func (s *after) AddOkDone(dc ...func()) {
	s.okDone = append(s.okDone, dc...)
}

func (s *after) AfterDo(err error) {
	for _, fc := range s.done {
		fc()
	}

	if err != nil {
		for _, fc := range s.errDone {
			fc()
		}
	} else {
		for _, fc := range s.okDone {
			fc()
		}
	}

}

type Context struct {
	Ctx      context.Context
	be       *merr.BusinessError
	dao      *dao.Query
	daoTx    *dao.QueryTx
	withTran bool
	rollback string

	fnTimer *traces.FnTimers

	after
}

// Init 初始化一个逻辑类
func (c *Context) Init(ctx context.Context, dao *dao.Query, be *merr.BusinessError) {
	c.BaseInit(ctx, be, WithDao(dao))
}

// BaseInit 基础初始化一个逻辑类
func (c *Context) BaseInit(ctx context.Context, be *merr.BusinessError, opts ...CtxOption) {
	c.Ctx = ctx
	c.be = be
	c.dao = dao.Q

	for _, opt := range opts {
		opt(c)
	}
}

// GetDao 获取dao
func (c *Context) GetDao() *dao.Query {
	if c.withTran {
		return c.daoTx.Query
	}

	return c.dao
}

// Tran 事务开启
func (c *Context) Tran() {
	c.daoTx = c.dao.Begin()
	c.withTran = true

	return
}

// TranCommit 事务提交
func (c *Context) TranCommit(err error) {
	if !c.withTran {
		return
	}

	c.withTran = false

	if err != nil {
		_ = c.daoTx.Rollback()
		app.Logger.Warn(c.Ctx, fmt.Sprintf("事务回滚了，原因：%v", err))
		return
	}

	if c.rollback != "" {
		_ = c.daoTx.Rollback()
		app.Logger.Warn(c.Ctx, fmt.Sprintf("事务回滚了，原因:%s", c.rollback))
		return
	}

	_ = c.daoTx.Commit()

	return
}

func (c *Context) BE() *merr.BusinessError {
	return c.be
}

func (c *Context) FnTm() *traces.FnTimers {
	return c.fnTimer
}

func (c *Context) Debug() {
	val, ok := metadata.Get(c.Ctx, consts.DebugKey)
	if !ok || val != "open" {
		return
	}
	c.fnTimer = traces.NewFnTimers(c.Ctx)
}

func (c *Context) Rollback(reason string) {
	c.rollback = reason
}