From c9690500fa5e68a5e451dcadfa66e3acd34a4c0e Mon Sep 17 00:00:00 2001
From: Lunny Xiao <xiaolunwen@gmail.com>
Date: Tue, 19 Sep 2017 20:59:41 +0800
Subject: [PATCH] add buffer Iterate support (#729)

---
 engine.go               |  7 ++++++
 session_iterate.go      | 50 ++++++++++++++++++++++++++++++++++++++
 session_iterate_test.go | 54 +++++++++++++++++++++++++++++++++++++++++
 statement.go            |  2 ++
 4 files changed, 113 insertions(+)

diff --git a/engine.go b/engine.go
index 0d34bf96..17d16063 100644
--- a/engine.go
+++ b/engine.go
@@ -1574,3 +1574,10 @@ func (engine *Engine) CondDeleted(colName string) builder.Cond {
 	}
 	return builder.IsNull{colName}.Or(builder.Eq{colName: zeroTime1})
 }
+
+// BufferSize sets buffer size for iterate
+func (engine *Engine) BufferSize(size int) *Session {
+	session := engine.NewSession()
+	session.isAutoClose = true
+	return session.BufferSize(size)
+}
diff --git a/session_iterate.go b/session_iterate.go
index a2f957a2..071fce49 100644
--- a/session_iterate.go
+++ b/session_iterate.go
@@ -23,6 +23,10 @@ func (session *Session) Iterate(bean interface{}, fun IterFunc) error {
 		defer session.Close()
 	}
 
+	if session.statement.bufferSize > 0 {
+		return session.bufferIterate(bean, fun)
+	}
+
 	rows, err := session.Rows(bean)
 	if err != nil {
 		return err
@@ -44,3 +48,49 @@ func (session *Session) Iterate(bean interface{}, fun IterFunc) error {
 	}
 	return err
 }
+
+// BufferSize sets the buffersize for iterate
+func (session *Session) BufferSize(size int) *Session {
+	session.statement.bufferSize = size
+	return session
+}
+
+func (session *Session) bufferIterate(bean interface{}, fun IterFunc) error {
+	if session.isAutoClose {
+		defer session.Close()
+	}
+
+	var bufferSize = session.statement.bufferSize
+	var limit = session.statement.LimitN
+	if limit > 0 && bufferSize > limit {
+		bufferSize = limit
+	}
+	var start = session.statement.Start
+	v := rValue(bean)
+	sliceType := reflect.SliceOf(v.Type())
+	var idx = 0
+	for {
+		slice := reflect.New(sliceType)
+		if err := session.Limit(bufferSize, start).find(slice.Interface(), bean); err != nil {
+			return err
+		}
+
+		for i := 0; i < slice.Elem().Len(); i++ {
+			if err := fun(idx, slice.Elem().Index(i).Addr().Interface()); err != nil {
+				return err
+			}
+			idx++
+		}
+
+		start = start + slice.Elem().Len()
+		if limit > 0 && idx+bufferSize > limit {
+			bufferSize = limit - idx
+		}
+
+		if bufferSize <= 0 || slice.Elem().Len() < bufferSize || idx == limit {
+			break
+		}
+	}
+
+	return nil
+}
diff --git a/session_iterate_test.go b/session_iterate_test.go
index 2c58c3ab..9a7ec25f 100644
--- a/session_iterate_test.go
+++ b/session_iterate_test.go
@@ -34,5 +34,59 @@ func TestIterate(t *testing.T) {
 		cnt++
 		return nil
 	})
+	assert.NoError(t, err)
 	assert.EqualValues(t, 1, cnt)
 }
+
+func TestBufferIterate(t *testing.T) {
+	assert.NoError(t, prepareEngine())
+
+	type UserBufferIterate struct {
+		Id    int64
+		IsMan bool
+	}
+
+	assert.NoError(t, testEngine.Sync2(new(UserBufferIterate)))
+
+	var size = 20
+	for i := 0; i < size; i++ {
+		cnt, err := testEngine.Insert(&UserBufferIterate{
+			IsMan: true,
+		})
+		assert.NoError(t, err)
+		assert.EqualValues(t, 1, cnt)
+	}
+
+	var cnt = 0
+	err := testEngine.BufferSize(9).Iterate(new(UserBufferIterate), func(i int, bean interface{}) error {
+		user := bean.(*UserBufferIterate)
+		assert.EqualValues(t, cnt+1, user.Id)
+		assert.EqualValues(t, true, user.IsMan)
+		cnt++
+		return nil
+	})
+	assert.NoError(t, err)
+	assert.EqualValues(t, size, cnt)
+
+	cnt = 0
+	err = testEngine.Limit(20).BufferSize(9).Iterate(new(UserBufferIterate), func(i int, bean interface{}) error {
+		user := bean.(*UserBufferIterate)
+		assert.EqualValues(t, cnt+1, user.Id)
+		assert.EqualValues(t, true, user.IsMan)
+		cnt++
+		return nil
+	})
+	assert.NoError(t, err)
+	assert.EqualValues(t, size, cnt)
+
+	cnt = 0
+	err = testEngine.Limit(7).BufferSize(9).Iterate(new(UserBufferIterate), func(i int, bean interface{}) error {
+		user := bean.(*UserBufferIterate)
+		assert.EqualValues(t, cnt+1, user.Id)
+		assert.EqualValues(t, true, user.IsMan)
+		cnt++
+		return nil
+	})
+	assert.NoError(t, err)
+	assert.EqualValues(t, 7, cnt)
+}
diff --git a/statement.go b/statement.go
index dc8a0c9b..23346c71 100644
--- a/statement.go
+++ b/statement.go
@@ -73,6 +73,7 @@ type Statement struct {
 	decrColumns     map[string]decrParam
 	exprColumns     map[string]exprParam
 	cond            builder.Cond
+	bufferSize      int
 }
 
 // Init reset all the statement's fields
@@ -111,6 +112,7 @@ func (statement *Statement) Init() {
 	statement.decrColumns = make(map[string]decrParam)
 	statement.exprColumns = make(map[string]exprParam)
 	statement.cond = builder.NewCond()
+	statement.bufferSize = 0
 }
 
 // NoAutoCondition if you do not want convert bean's field as query condition, then use this function