支持指定统计总记录数的列

Author: xfali

builder.go

@@ -29,8 +29,8 @@ func (b *builder) Page(page, pageSize int) *builder {
 //分页
 //page 页码
 //pageSize 分页大小
-func (b *builder) PageWithTotal(page, pageSize int) *builder {
-    b.ctx = StartPageWithTotal(page, pageSize, b.ctx)
+func (b *builder) PageWithTotal(page, pageSize int, countColumn string) *builder {
+    b.ctx = StartPageWithTotal(page, pageSize, countColumn, b.ctx)
     return b
 }
 

pagehelper.go

@@ -25,7 +25,6 @@ import (
 const (
     pageHelperValue  = "_page_helper_value"
     orderHelperValue = "_order_helper_value"
-    totalHelperValue = "_total_helper_value"
 
     ASC  = "ASC"
     DESC = "DESC"
@@ -46,7 +45,8 @@ type PageParam struct {
     Page     int
     PageSize int
 
-    total bool
+    countColumn string
+    total       int64
 }
 
 func New(f factory.Factory) *Factory {
@@ -57,10 +57,10 @@ func GetTotal(ctx context.Context) int64 {
     if ctx == nil {
         return 0
     }
-    p := ctx.Value(totalHelperValue)
+    p := ctx.Value(pageHelperValue)
     if p != nil {
-        if t, ok := p.(int64); ok {
-            return t
+        if param, ok := p.(*PageParam); ok {
+            return param.total
         }
     }
     return 0
@@ -80,15 +80,15 @@ type Executor struct {
 //pageSize 分页大小
 //ctx 初始context
 func StartPage(page, pageSize int, ctx context.Context) context.Context {
-    return context.WithValue(ctx, pageHelperValue, &PageParam{Page: page, PageSize: pageSize, total: false})
+    return context.WithValue(ctx, pageHelperValue, &PageParam{Page: page, PageSize: pageSize, total: 0})
 }
 
 //分页(包含total信息)
 //page 页码
 //pageSize 分页大小
 //ctx 初始context
-func StartPageWithTotal(page, pageSize int, ctx context.Context) context.Context {
-    return context.WithValue(ctx, pageHelperValue, &PageParam{Page: page, PageSize: pageSize, total: true})
+func StartPageWithTotal(page, pageSize int, countColumn string, ctx context.Context) context.Context {
+    return context.WithValue(ctx, pageHelperValue, &PageParam{Page: page, PageSize: pageSize, total: -1, countColumn: countColumn})
 }
 
 //排序
@@ -132,9 +132,8 @@ func (exec *Executor) Query(ctx context.Context, result reflection.Object, sql s
     if p != nil {
         if param, ok := p.(*PageParam); ok {
             sql = PageModifier(sql, param)
-            if param.total == true {
-                total := exec.getTotal(ctx, originSql, params...)
-                ctx = context.WithValue(ctx, totalHelperValue, total)
+            if param.total == -1 {
+                param.total = exec.getTotal(ctx, originSql, param.countColumn, params...)
             }
         }
     }
@@ -183,8 +182,8 @@ func modifyPageSql(sql string, p *PageParam) string {
     return b.String()
 }
 
-func (exec *Executor) getTotal(ctx context.Context, sql string, params ...interface{}) int64 {
-    totalSql := CountModifier(sql)
+func (exec *Executor) getTotal(ctx context.Context, sql, countColumn string, params ...interface{}) int64 {
+    totalSql := CountModifier(sql, countColumn)
     var total int64
     obj, err := gobatis.ParseObject(&total)
     if err == nil {
@@ -194,10 +193,15 @@ func (exec *Executor) getTotal(ctx context.Context, sql string, params ...interf
     return 0
 }
 
-func modifyCountSql(sql string) string {
+func modifyCountSql(sql, countColumn string) string {
+    if countColumn == "" {
+        countColumn = "0"
+    }
     b := strings.Builder{}
-    b.WriteString("SELECT COUNT(0) FROM (")
-    b.WriteString(sql)
-    b.WriteString(")")
+    b.WriteString("SELECT COUNT(`")
+    b.WriteString(countColumn)
+    b.WriteString("`) FROM (")
+    b.WriteString(strings.TrimSpace(sql))
+    b.WriteString(") AS __hp_tempCountTl")
     return b.String()
 }

pagehelper_test.go

@@ -124,7 +124,7 @@ func TestPageHelper2(t *testing.T) {
 }
 
 func TestModifyPage(t *testing.T) {
-    sql := PageModifier("select * from x", &PageParam{1, 2, false})
+    sql := PageModifier("select * from x", &PageParam{Page: 1, PageSize: 2,})
     t.Log(sql)
     if strings.TrimSpace(sql) != `select * from x LIMIT 2, 2` {
         t.Fail()
@@ -148,20 +148,30 @@ func TestModifyOrder(t *testing.T) {
 }
 
 func TestModifyCount(t *testing.T) {
-    sql := CountModifier("select ? from x")
-    t.Log(sql)
+    t.Run("empty", func(t *testing.T) {
+        sql := CountModifier("select ? from x", "")
+        t.Log(sql)
 
-    if strings.TrimSpace(sql) != "SELECT COUNT(0) FROM (select ? from x)" {
-        t.Fail()
-    }
-}
+        if strings.TrimSpace(sql) != "SELECT COUNT(0) FROM (select ? from x) AS __hp_tempCountTl" {
+            t.Fail()
+        }
+    })
+
+    t.Run("test", func(t *testing.T) {
+        sql := CountModifier("select ? from x", "test")
+        t.Log(sql)
 
+        if strings.TrimSpace(sql) != "SELECT COUNT(`test`) FROM (select ? from x) AS __hp_tempCountTl" {
+            t.Fail()
+        }
+    })
+}
 
 func TestChangeModifyCount(t *testing.T) {
-    CountModifier = func(sql string) string {
+    CountModifier = func(sql, c string) string {
         return "test " + sql
     }
-    sql := CountModifier("select ? from x")
+    sql := CountModifier("select ? from x", "")
     t.Log(sql)
 
     if strings.TrimSpace(sql) != "test select ? from x" {
@@ -185,7 +195,7 @@ func TestModifyOrderAndPage(t *testing.T) {
     sql, p := order("select ? from x", "field1")
     t.Log(sql)
 
-    sql = PageModifier(sql, &PageParam{1, 2, false})
+    sql = PageModifier(sql, &PageParam{Page: 1, PageSize: 2,})
 
     t.Log(sql)
     for _, v := range p {