go语言ent ORM框架增强-自定义分页查询

ent原始分页查询方法

func (s sysUserServiceImpl) ListSysUser(c *trace.Context, req *request.SysUserListReq) (*result.PageResult, error) {
	query := s.CoreDB().SysUser.Query()
	
	if req.UserName != "" {
		query.Where(sysuser.UserNameContains(req.UserName))
	}
	if req.UserStatus != "" {
		query.Where(sysuser.UserStatus(req.UserStatus))
	}
	
	total, err := query.Clone().Count(c)
	if err != nil {
		return nil, err
	}
	
	if req.Reverse {
		query.Order(entcore.Desc(sysuser.FieldID))
	} else {
		query.Order(entcore.Desc(sysuser.FieldID))
	}
	
	offset, limit := req.GetTerm()
	list, err := query.Offset(offset).Limit(limit).All(c)
	if err != nil {
		return nil, err
	}
	return result.NewPageResult(total, list), nil
}

每个实体的查询方法都类似，先计算总量，然后查询。能否封装成通用方法呢。

自定义模板

通过ent自定义模板功能，为每个模型都绑定Page方法，封装分页查询。

{{/*
// Code generated by ent with template, DO NOT EDIT.
*/}}

{{ define "pagination" }}
    {{- /* gotype: entgo.io/ent/entc/gen.Graph */ -}}

    {{ template "header" $ }}

    {{ $pkg := base $.Config.Package }}

    {{ template "import" $ }}

    func getTerm(pageNo int, pageSize int) (int, int) {
    	if pageSize <= 0 {
    		pageSize = 10
    	}
    	if pageNo <= 0 {
    		pageNo = 1
    	}

    	offset := (pageNo - 1) * pageSize
    	limit := pageSize
    	return offset, limit
    }

    {{ range $node := $.Nodes }}
        {{ $name := $node.Name }}
        {{ $r := $node.Receiver }}
        {{ $queryName := $node.QueryName }}

        type {{ $queryName }}Option func(*{{ $queryName }})

        func ({{ $r }} *{{ $queryName }}) Page(
            ctx context.Context,
            pageNo int,
            pageSize int,
            opts ...{{ $queryName }}Option,
        ) (int, []*{{ $name }}, error) {
            countQuery := {{ $r }}.Clone()

            for _, opt := range opts {
                opt({{ $r }})
                opt(countQuery)
            }

            countQuery.ctx.Fields = nil
            total, err := countQuery.Count(ctx)
            if err != nil {
                return 0, nil, err
            }

            offset, limit := getTerm(pageNo, pageSize)
            rows, err := {{ $r }}.
                Offset(offset).
                Limit(limit).
                All(ctx)
            if err != nil {
                return 0, nil, err
            }

            return total, rows, nil
        }
    {{ end }}
{{ end }}

生成代码

在生成代码的命令中加入--template选项指定模板：

go run -mod=mod entgo.io/ent/cmd/ent generate --feature intercept,schema/snapshot ./schema  --target ./schema/entcode --template glob="./template/*.tmpl"

生成的代码：

type SysUserQueryOption func(*SysUserQuery)

func (_m *SysUserQuery) Page(
	ctx context.Context,
	pageNo int,
	pageSize int,
	opts ...SysUserQueryOption,
) (int, []*SysUser, error) {
	countQuery := _m.Clone()

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

	countQuery.ctx.Fields = nil
	total, err := countQuery.Count(ctx)
	if err != nil {
		return 0, nil, err
	}

	offset, limit := getTerm(pageNo, pageSize)
	rows, err := _m.
		Offset(offset).
		Limit(limit).
		All(ctx)
	if err != nil {
		return 0, nil, err
	}

	return total, rows, nil
}

其中SysUserQueryOption为SysUserQuery的包装器，以支持指定其他查询参数，如Where条件或order by 排序规则等。不过我更喜欢拼接好条件后再执行Page方法。

使用Page

func (s sysUserServiceImpl) ListSysUser(c *trace.Context, req *request.SysUserListReq) (*result.PageResult, error) {
	query := s.CoreDB().SysUser.Query()

	if !req.DeptID.EqualsInt64(0) {
		query.Where(sysuser.DeptIDEQ(req.DeptID))
	}

	if req.UserName != "" {
		query.Where(sysuser.UserNameContains(req.UserName))
	}
	if req.UserStatus != "" {
		query.Where(sysuser.UserStatus(req.UserStatus))
	}

	query.OrderByIDDesc()

	total, rows, err := query.Page(c, req.PageNo, req.PageSize)
	if err != nil {
		return nil, err
	}

	return result.NewPageResult(total, rows), nil
}

注意

所有模型的Page查询方法都在文件pagination.go中，理想情况是追加到各自模型的xxx_query.go中。