增加获得总记录数支持

Author: xfali

builder.go

@@ -26,6 +26,14 @@ func (b *builder) Page(page, pageSize int) *builder {
     return b
 }
 
+//分页
+//page 页码
+//pageSize 分页大小
+func (b *builder) PageWithTotal(page, pageSize int) *builder {
+    b.ctx = StartPageWithTotal(page, pageSize, b.ctx)
+    return b
+}
+
 //手动指定字段和排序
 //field 字段
 //order 排序 [ASC | DESC]

pagehelper.go

@@ -11,6 +11,7 @@ package pagehelper
 import (
     "context"
     "fmt"
+    "github.com/xfali/gobatis"
     "github.com/xfali/gobatis/common"
     "github.com/xfali/gobatis/executor"
     "github.com/xfali/gobatis/factory"
@@ -24,11 +25,18 @@ import (
 const (
     pageHelperValue  = "_page_helper_value"
     orderHelperValue = "_order_helper_value"
+    totalHelperValue = "_total_helper_value"
 
     ASC  = "ASC"
     DESC = "DESC"
 )
 
+var (
+    OrderByModifier = modifyOrderSql
+    PageModifier    = modifyPageSql
+    CountModifier   = modifyCountSql
+)
+
 type OrderParam struct {
     Field string
     Order string
@@ -37,12 +45,27 @@ type OrderParam struct {
 type PageParam struct {
     Page     int
     PageSize int
+
+    total bool
 }
 
 func New(f factory.Factory) *Factory {
     return &Factory{f}
 }
 
+func GetTotal(ctx context.Context) int64 {
+    if ctx == nil {
+        return 0
+    }
+    p := ctx.Value(totalHelperValue)
+    if p != nil {
+        if t, ok := p.(int64); ok {
+            return t
+        }
+    }
+    return 0
+}
+
 type Factory struct {
     fac factory.Factory
 }
@@ -57,7 +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})
+    return context.WithValue(ctx, pageHelperValue, &PageParam{Page: page, PageSize: pageSize, total: false})
+}
+
+//分页(包含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})
 }
 
 //排序
@@ -89,17 +120,22 @@ func (exec *Executor) Rollback(require bool) error {
 }
 
 func (exec *Executor) Query(ctx context.Context, result reflection.Object, sql string, params ...interface{}) error {
+    originSql := sql
     o := ctx.Value(orderHelperValue)
     if o != nil {
         if param, ok := o.(*OrderParam); ok {
-            sql = modifySqlOrder(sql, param)
+            sql = OrderByModifier(sql, param)
         }
     }
 
     p := ctx.Value(pageHelperValue)
     if p != nil {
         if param, ok := p.(*PageParam); ok {
-            sql = modifySql(sql, param)
+            sql = PageModifier(sql, param)
+            if param.total == true {
+                total := exec.getTotal(ctx, originSql, params...)
+                ctx = context.WithValue(ctx, totalHelperValue, total)
+            }
         }
     }
     exec.log(logging.DEBUG, "PageHelper Query: [%s], params: %s\n", sql, fmt.Sprint(params))
@@ -130,7 +166,7 @@ func (f *Factory) CreateExecutor(transaction transaction.Transaction) executor.E
     }
 }
 
-func modifySqlOrder(sql string, p *OrderParam) string {
+func modifyOrderSql(sql string, p *OrderParam) string {
     if p.Field == "" {
         return sql
     }
@@ -140,9 +176,28 @@ func modifySqlOrder(sql string, p *OrderParam) string {
     return b.String()
 }
 
-func modifySql(sql string, p *PageParam) string {
+func modifyPageSql(sql string, p *PageParam) string {
     b := strings.Builder{}
     b.WriteString(strings.TrimSpace(sql))
     b.WriteString(fmt.Sprintf(" LIMIT %d, %d ", p.Page*p.PageSize, p.PageSize))
     return b.String()
 }
+
+func (exec *Executor) getTotal(ctx context.Context, sql string, params ...interface{}) int64 {
+    totalSql := CountModifier(sql)
+    var total int64
+    obj, err := gobatis.ParseObject(&total)
+    if err == nil {
+        exec.exec.Query(ctx, obj, totalSql, params...)
+        return total
+    }
+    return 0
+}
+
+func modifyCountSql(sql string) string {
+    b := strings.Builder{}
+    b.WriteString("SELECT COUNT(0) FROM (")
+    b.WriteString(sql)
+    b.WriteString(")")
+    return b.String()
+}

pagehelper_test.go

@@ -124,40 +124,75 @@ func TestPageHelper2(t *testing.T) {
 }
 
 func TestModifyPage(t *testing.T) {
-    sql := modifySql("select * from x", &PageParam{1, 2})
+    sql := PageModifier("select * from x", &PageParam{1, 2, false})
     t.Log(sql)
+    if strings.TrimSpace(sql) != `select * from x LIMIT 2, 2` {
+        t.Fail()
+    }
 }
 
 func order(sql string, params ...interface{}) (string, []interface{}) {
-    return modifySqlOrder(sql, &OrderParam{"test", ASC}), params
+    return OrderByModifier(sql, &OrderParam{"test", ASC}), params
 }
 
 func TestModifyOrder(t *testing.T) {
     sql, p := order("select ? from x", "field1")
     t.Log(sql)
-    if len(p) != 2 {
-        t.Fatal()
-    }
     for _, v := range p {
         t.Log(v)
     }
+
+    if strings.TrimSpace(sql) != "select ? from x ORDER BY `test` ASC" {
+        t.Fail()
+    }
+}
+
+func TestModifyCount(t *testing.T) {
+    sql := CountModifier("select ? from x")
+    t.Log(sql)
+
+    if strings.TrimSpace(sql) != "SELECT COUNT(0) FROM (select ? from x)" {
+        t.Fail()
+    }
+}
+
+
+func TestChangeModifyCount(t *testing.T) {
+    CountModifier = func(sql string) string {
+        return "test " + sql
+    }
+    sql := CountModifier("select ? from x")
+    t.Log(sql)
+
+    if strings.TrimSpace(sql) != "test select ? from x" {
+        t.Fail()
+    }
+}
+
+func TestGetTotal(t *testing.T) {
+    ctx, _ := context.WithTimeout(context.Background(), 2*time.Second)
+    ctx = OrderBy("test", ASC, ctx)
+    ctx = StartPage(1, 2, ctx)
+
+    total := GetTotal(ctx)
+    t.Log(total)
+    if total != 0 {
+        t.Fail()
+    }
 }
 
 func TestModifyOrderAndPage(t *testing.T) {
     sql, p := order("select ? from x", "field1")
     t.Log(sql)
-    if len(p) != 2 {
-        t.Fatal()
-    }
 
-    sql = modifySql(sql, &PageParam{1, 2})
+    sql = PageModifier(sql, &PageParam{1, 2, false})
 
     t.Log(sql)
     for _, v := range p {
         t.Log(v)
     }
 
-    if strings.TrimSpace(sql) != "select ? from x ORDER BY ? ASC LIMIT 2, 2" {
+    if strings.TrimSpace(sql) != "select ? from x ORDER BY `test` ASC LIMIT 2, 2" {
         t.Fail()
     }
 }