From 97ddb3b286d9bbdc86056206199c6fa75a261c2c Mon Sep 17 00:00:00 2001
From: Alfonso Acosta <alfonso@stellar.org>
Date: Mon, 22 Nov 2021 20:33:17 +0100
Subject: [PATCH] Add benchmark for BatchInsertBuilder (#4086)

---
 support/db/batch_insert_builder_test.go | 51 ++++++++++++++++++++++---
 support/db/dbtest/db.go                 |  6 +--
 2 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/support/db/batch_insert_builder_test.go b/support/db/batch_insert_builder_test.go
index 9c2dfb3b81..330631ad4b 100644
--- a/support/db/batch_insert_builder_test.go
+++ b/support/db/batch_insert_builder_test.go
@@ -2,6 +2,7 @@ package db
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/stellar/go/support/db/dbtest"
@@ -20,6 +21,44 @@ type invalidHungerRow struct {
 	LastName    string `db:"last_name"`
 }
 
+func BenchmarkBatchInsertBuilder(b *testing.B) {
+	b.StopTimer()
+	// In order to show SQL queries
+	// log.SetLevel(logrus.DebugLevel)
+	db := dbtest.Postgres(b).Load(testSchema)
+	defer db.Close()
+	sess := &Session{DB: db.Open()}
+	defer sess.DB.Close()
+	ctx := context.Background()
+	maxBatchSize := 1000
+	insertBuilder := &BatchInsertBuilder{
+		Table:        sess.GetTable("people"),
+		MaxBatchSize: maxBatchSize,
+	}
+	b.StartTimer()
+	for i := 0; i < b.N; i++ {
+		for j := 0; j < maxBatchSize; j++ {
+			err := insertBuilder.RowStruct(ctx, hungerRow{
+				Name:        fmt.Sprintf("bubba%d", i*maxBatchSize+j),
+				HungerLevel: "1202",
+			})
+			require.NoError(b, err)
+		}
+	}
+	err := insertBuilder.Exec(ctx)
+
+	b.StopTimer()
+	assert.NoError(b, err)
+	var count []int
+	err = sess.SelectRaw(ctx,
+		&count,
+		"SELECT COUNT(*) FROM people",
+	)
+	assert.NoError(b, err)
+	preexistingCount := 3
+	assert.Equal(b, b.N*maxBatchSize+preexistingCount, count[0])
+}
+
 func TestBatchInsertBuilder(t *testing.T) {
 	db := dbtest.Postgres(t).Load(testSchema)
 	defer db.Close()
@@ -87,8 +126,8 @@ func TestBatchInsertBuilder(t *testing.T) {
 		t,
 		found,
 		[]person{
-			person{Name: "bubba", HungerLevel: "120"},
-			person{Name: "bubba2", HungerLevel: "1202"},
+			{Name: "bubba", HungerLevel: "120"},
+			{Name: "bubba2", HungerLevel: "1202"},
 		},
 	)
 
@@ -122,8 +161,8 @@ func TestBatchInsertBuilder(t *testing.T) {
 		t,
 		found,
 		[]person{
-			person{Name: "bubba", HungerLevel: "120"},
-			person{Name: "bubba2", HungerLevel: "1202"},
+			{Name: "bubba", HungerLevel: "120"},
+			{Name: "bubba2", HungerLevel: "1202"},
 		},
 	)
 
@@ -145,8 +184,8 @@ func TestBatchInsertBuilder(t *testing.T) {
 		t,
 		found,
 		[]person{
-			person{Name: "bubba2", HungerLevel: "1202"},
-			person{Name: "bubba", HungerLevel: "1"},
+			{Name: "bubba2", HungerLevel: "1202"},
+			{Name: "bubba", HungerLevel: "1"},
 		},
 	)
 }
diff --git a/support/db/dbtest/db.go b/support/db/dbtest/db.go
index 79007a286e..bda67c49e8 100644
--- a/support/db/dbtest/db.go
+++ b/support/db/dbtest/db.go
@@ -22,7 +22,7 @@ type DB struct {
 	Dialect string
 	DSN     string
 	dbName  string
-	t       *testing.T
+	t       testing.TB
 	closer  func()
 	closed  bool
 }
@@ -101,7 +101,7 @@ func (db *DB) Version() (major int) {
 	return major
 }
 
-func execStatement(t *testing.T, pguser, query string) {
+func execStatement(t testing.TB, pguser, query string) {
 	db, err := sqlx.Open("postgres", fmt.Sprintf("postgres://%s@localhost/?sslmode=disable", pguser))
 	require.NoError(t, err)
 	_, err = db.Exec(query)
@@ -113,7 +113,7 @@ func execStatement(t *testing.T, pguser, query string) {
 // of the running process.  It assumes that you have postgres running on the
 // default port, have the command line postgres tools installed, and that the
 // current user has access to the server.  It panics on the event of a failure.
-func Postgres(t *testing.T) *DB {
+func Postgres(t testing.TB) *DB {
 	var result DB
 	result.dbName = randomName()
 	result.Dialect = "postgres"