From a5fcf24c04b09f3d526f13f3838c66ecea9fbe9c Mon Sep 17 00:00:00 2001
From: MarkJoyMa <64180138+MarkJoyMa@users.noreply.github.com>
Date: Fri, 6 Jan 2023 23:30:50 +0800
Subject: [PATCH] feat: add batch inserter (#2755)

---
 core/stores/mon/bulkinserter.go      | 17 ++++++++++++++++-
 core/stores/mon/bulkinserter_test.go | 18 ++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/core/stores/mon/bulkinserter.go b/core/stores/mon/bulkinserter.go
index c3fdfa23..74cbb7d9 100644
--- a/core/stores/mon/bulkinserter.go
+++ b/core/stores/mon/bulkinserter.go
@@ -25,8 +25,23 @@ type (
 	}
 )
 
-// NewBulkInserter returns a BulkInserter.
+// Deprecated. Use NewBatchInserter instead.
 func NewBulkInserter(coll *mongo.Collection, interval ...time.Duration) *BulkInserter {
+	return newBulkInserter(coll, interval...)
+}
+
+// NewBatchInserter returns a BulkInserter.
+func NewBatchInserter(coll Collection, interval ...time.Duration) (*BulkInserter, error) {
+	cloneColl, err := coll.Clone()
+	if err != nil {
+		return nil, err
+	}
+
+	return newBulkInserter(cloneColl, interval...), nil
+}
+
+// newBulkInserter returns a BulkInserter.
+func newBulkInserter(coll *mongo.Collection, interval ...time.Duration) *BulkInserter {
 	inserter := &dbInserter{
 		collection: coll,
 	}
diff --git a/core/stores/mon/bulkinserter_test.go b/core/stores/mon/bulkinserter_test.go
index 0d6aa214..f2889d35 100644
--- a/core/stores/mon/bulkinserter_test.go
+++ b/core/stores/mon/bulkinserter_test.go
@@ -25,3 +25,21 @@ func TestBulkInserter(t *testing.T) {
 		bulk.Flush()
 	})
 }
+
+func TestBatchInserter(t *testing.T) {
+	mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
+	defer mt.Close()
+
+	mt.Run("test", func(mt *mtest.T) {
+		mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
+		bulk, err := NewBatchInserter(createModel(mt).Collection)
+		assert.Equal(t, err, nil)
+		bulk.SetResultHandler(func(result *mongo.InsertManyResult, err error) {
+			assert.Nil(t, err)
+			assert.Equal(t, 2, len(result.InsertedIDs))
+		})
+		bulk.Insert(bson.D{{Key: "foo", Value: "bar"}})
+		bulk.Insert(bson.D{{Key: "foo", Value: "baz"}})
+		bulk.Flush()
+	})
+}