batch_insert.go

package ddl

import (
	"bytes"
	"context"
	"database/sql"
	"fmt"
	"io"
	"math"
	"strconv"
)

// BatchInsert is used to insert data into a table in batches.
type BatchInsert struct {
	// BatchSize is the batch size used in INSERT statements e.g.
	// INSERT INTO tbl VALUES (1)           /* batchsize=1 */
	// INSERT INTO tbl VALUES (1), (2), (3) /* batchsize=3 */
	// If Batchsize <= 0, a default batch size of 500 will be used.
	Batchsize int

	// Dialect is the database dialect. Possible values: "sqlite", "postgres",
	// "mysql", "sqlserver".
	Dialect string

	// TableSchema is the name of the schema that the table belongs to.
	TableSchema string

	// TableName is the name of the table.
	TableName string

	// Columns is the list of columns participating in the INSERT.
	Columns []string

	// KeyColumns is the list of columns in the table that uniquely identify a
	// row e.g. PRIMARY KEY or UNIQUE columns. If provided, the batch insert
	// will become an upsert instead (e.g. ON CONFLICT DO UPDATE).
	KeyColumns []string

	// IdentityColumns is the list of columns in the table that are identity
	// columns e.g. GENERATED BY DEFAULT AS IDENTITY (Postgres) or IDENTITY
	// (SQL Server). Postgres or SQL Server only. If provided:
	//
	// (Postgres) The identity sequence for each identity column will be
	// updated accordingly at the end of the batch insert e.g. SELECT
	// setval(pg_get_serial_sequence('table', 'id'), max(id)) FROM table;
	//
	// (SQL Server) SET IDENTITY_INSERT will be enabled for the table at the
	// start of batch insert and disabled at the end. It is vital to pass in a
	// stateful database connection (an *sql.Conn or *sql.Tx) because we need
	// to SET IDENTITY_INSERT ON and SET IDENTITY_INSERT OFF on the same
	// connection.
	IdentityColumns []string
}

// ExecContext will execute the batch insert with the given database connection and the
// `next` iterator function.
//
// The `next` iterator function must populate the []any slice with each
// successive row's values when called, and must return io.EOF once there
// are no more results. The length of the []any slice is guaranteed to match
// the length of the Columns field of the BatchInsert.
//
// Here is a basic example:
//
//   bi := &BatchInsert{
//       Dialect:         "postgres",
//       TableName:       "actors",
//       Columns:         []string{"actor_id", "first_name", "last_name"},
//       KeyColumns:      []string{"actor_id"},
//       IdentityColumns: []string{"actor_id"},
//   }
//   actors := []Actor{
//       {ActorID: 1, FirstName: "PENELOPE", LastName: "GUINESS"},
//       {ActorID: 2, FirstName: "NICK", LastName: "WAHLBERG"},
//       {ActorID: 3, FirstName: "ED", LastName: "CHASE"},
//   }
//   i := 0
//   _, err := bi.ExecContext(ctx, db, func(row []any) error {
//       if i >= len(actors) {
//           return io.EOF
//       }
//       actor := actors[i]
//       row[0] = actor.ActorID
//       row[1] = actor.FirstName
//       row[2] = actor.LastName
//       i++
//       return nil
//   })
func (bi *BatchInsert) ExecContext(ctx context.Context, db DB, next func([]any) error) (rowsAffected int64, err error) {
	if db == nil {
		return 0, fmt.Errorf("db is nil")
	}
	if next == nil {
		return 0, fmt.Errorf("next is nil")
	}
	isColumn := make(map[string]struct{})
	for _, column := range bi.Columns {
		isColumn[column] = struct{}{}
	}
	keyColumns := bi.KeyColumns
	for _, column := range bi.KeyColumns {
		if _, ok := isColumn[column]; !ok {
			keyColumns = keyColumns[:0]
			break
		}
	}
	isKeyColumn := make(map[string]struct{})
	for _, column := range keyColumns {
		isKeyColumn[column] = struct{}{}
	}
	nonKeyColumns := make([]string, 0, len(bi.Columns)-len(keyColumns))
	for _, column := range bi.Columns {
		if _, ok := isKeyColumn[column]; ok {
			continue
		}
		nonKeyColumns = append(nonKeyColumns, column)
	}
	identityColumns := make([]string, 0, len(bi.IdentityColumns))
	for _, column := range bi.IdentityColumns {
		if _, ok := isColumn[column]; !ok {
			continue
		}
		identityColumns = append(identityColumns, column)
	}
	isIdentityColumn := make(map[string]struct{})
	for _, column := range identityColumns {
		isIdentityColumn[column] = struct{}{}
	}

	const tempTableName = "tmp_"
	batchsize := bi.Batchsize
	if batchsize <= 0 {
		batchsize = 500
	}
	if bi.Dialect == DialectSQLServer && batchsize > 25 {
		// SQLServer should be capped to a batchsize of 25, beyond that
		// performance will drop.
		// https://www.red-gate.com/simple-talk/databases/sql-server/performance-sql-server/comparing-multiple-rows-insert-vs-single-row-insert-with-three-data-load-methods/
		batchsize = 25
	}
	// adjust batchsize if it exceeds known upper bounds
	// https://stackoverflow.com/q/6581573
	// https://www.jooq.org/doc/3.12/manual/sql-building/dsl-context/custom-settings/settings-inline-threshold/
	var paramlimit int
	switch bi.Dialect {
	case DialectSQLite:
		paramlimit = 32766 // prior to v3.32.0 the limit was 999 https://www.sqlite.org/limits.html
	case DialectPostgres:
		paramlimit = 65535 // https://stackoverflow.com/a/49379324
	case DialectMySQL:
		paramlimit = 65535
	case DialectSQLServer:
		paramlimit = 2100
	}
	if paramlimit > 0 {
		maxBatchsize := int(math.Floor(float64(paramlimit)/float64(len(bi.Columns)))) - 1
		if batchsize > maxBatchsize {
			batchsize = maxBatchsize
		}
	}

	buf := bufpool.Get().(*bytes.Buffer)
	buf.Reset()
	defer bufpool.Put(buf)

	// tableName
	tableName := QuoteIdentifier(bi.Dialect, bi.TableName)
	if bi.TableSchema != "" {
		tableName = QuoteIdentifier(bi.Dialect, bi.TableSchema) + "." + tableName
	}

	// columnNames
	if len(bi.Columns) == 0 {
		return 0, fmt.Errorf("no columns provided")
	}
	buf.Reset()
	for i, column := range bi.Columns {
		if i > 0 {
			buf.WriteString(", ")
		}
		buf.WriteString(QuoteIdentifier(bi.Dialect, column))
	}
	columnNames := buf.String()

	// SET IDENTITY_INSERT <table> ON
	if len(identityColumns) > 0 && bi.Dialect == DialectSQLServer {
		_, err = db.ExecContext(ctx, "SET IDENTITY_INSERT "+tableName+" ON")
		if err != nil {
			return 0, err
		}
	}

	// updateQuery: updateStart + middle + updateEnd (SQLServer only)
	// insertQuery: insertStart + middle + insertEnd

	// insertStart
	buf.Reset()
	buf.WriteString("INSERT INTO " + tableName + " (" + columnNames + ")")
	if len(identityColumns) > 0 && bi.Dialect == DialectPostgres {
		buf.WriteString(" OVERRIDING SYSTEM VALUE VALUES ")
	} else if len(keyColumns) > 0 && bi.Dialect == DialectSQLServer {
		buf.WriteString(" SELECT " + columnNames + " FROM (VALUES ")
	} else {
		buf.WriteString(" VALUES ")
	}
	insertStart := buf.String()

	// middle
	buf.Reset()
	writeRowValues(ctx, bi.Dialect, buf, batchsize, len(bi.Columns))
	middle := buf.String()

	// insertEnd
	var insertEnd string
	if len(keyColumns) > 0 {
		buf.Reset()
		switch bi.Dialect {
		case DialectSQLite, DialectPostgres:
			buf.WriteString(" ON CONFLICT (")
			for i, column := range keyColumns {
				if i > 0 {
					buf.WriteString(", ")
				}
				buf.WriteString(QuoteIdentifier(bi.Dialect, column))
			}
			buf.WriteString(") ")
			if len(nonKeyColumns) > 0 {
				buf.WriteString("DO UPDATE SET ")
				for i, column := range nonKeyColumns {
					if i > 0 {
						buf.WriteString(", ")
					}
					columnName := QuoteIdentifier(bi.Dialect, column)
					buf.WriteString(columnName + " = EXCLUDED." + columnName)
				}
			} else {
				buf.WriteString("DO NOTHING")
			}
		case DialectMySQL:
			buf.WriteString(" ON DUPLICATE KEY UPDATE ")
			if len(nonKeyColumns) > 0 {
				for i, column := range nonKeyColumns {
					if i > 0 {
						buf.WriteString(", ")
					}
					columnName := QuoteIdentifier(bi.Dialect, column)
					buf.WriteString(columnName + " = VALUES(" + columnName + ")")
				}
			} else {
				columnName := QuoteIdentifier(bi.Dialect, keyColumns[0])
				buf.WriteString(columnName + " = " + columnName)
			}
		case DialectSQLServer:
			buf.WriteString(") AS " + tempTableName + " (" + columnNames + ") WHERE NOT EXISTS (SELECT 1 FROM " + tableName + " WHERE ")
			for i, column := range keyColumns {
				if i > 0 {
					buf.WriteString(" AND ")
				}
				columnName := QuoteIdentifier(bi.Dialect, column)
				buf.WriteString(tableName + "." + columnName + " = " + tempTableName + "." + columnName)
			}
			buf.WriteString(")")
		}
		insertEnd = buf.String()
	}

	// prepare SQL insertStmt
	insertStmt, err := db.PrepareContext(ctx, insertStart+middle+insertEnd)
	if err != nil {
		return 0, err
	}
	defer insertStmt.Close()

	// updateStart, updateEnd (SQLServer only)
	var updateStart, updateEnd string
	var updateStmt *sql.Stmt
	if len(keyColumns) > 0 && len(nonKeyColumns) > 0 && bi.Dialect == DialectSQLServer {
		buf.Reset()
		buf.WriteString("UPDATE " + tableName + " WITH (UPDLOCK, SERIALIZABLE) SET ")
		written := false
		for _, column := range bi.Columns {
			_, isKey := isKeyColumn[column]
			_, isIdentity := isIdentityColumn[column]
			if isKey || isIdentity {
				continue
			}
			if !written {
				written = true
			} else {
				buf.WriteString(", ")
			}
			columnName := QuoteIdentifier(bi.Dialect, column)
			buf.WriteString(columnName + " = " + tempTableName + "." + columnName)
		}
		buf.WriteString(" FROM " + tableName + " JOIN (VALUES ")
		updateStart = buf.String()

		buf.Reset()
		buf.WriteString(") AS " + tempTableName + " (" + columnNames + ") ON ")
		for i, column := range keyColumns {
			if i > 0 {
				buf.WriteString(" AND ")
			}
			columnName := QuoteIdentifier(bi.Dialect, column)
			buf.WriteString(tableName + "." + columnName + " = " + tempTableName + "." + columnName)
		}
		updateEnd = buf.String()

		// prepare SQL updateStmt
		updateStmt, err = db.PrepareContext(ctx, updateStart+middle+updateEnd)
		if err != nil {
			return 0, err
		}
		defer updateStmt.Close()
	}

	// insert values using insertStmt (and updateStmt), batch by batch
	var batchcounter int
	args := make([]any, 0, bi.Batchsize*len(bi.Columns))
	rowvalue := make([]any, len(bi.Columns))
	for {
		err = next(rowvalue)
		if err == io.EOF {
			break
		}
		if err != nil {
			return rowsAffected, err
		}
		batchcounter++
		args = append(args, rowvalue...)
		if batchcounter >= batchsize {
			if updateStmt != nil {
				result, updateErr := updateStmt.ExecContext(ctx, args...)
				if updateErr != nil {
					query, err := Sprintf(bi.Dialect, updateStart+middle+updateEnd, args)
					if err == nil {
						return rowsAffected, fmt.Errorf("%w\n"+query, updateErr)
					}
					return rowsAffected, updateErr
				}
				n, updateErr := result.RowsAffected()
				if updateErr != nil {
					if query, err2 := Sprintf(bi.Dialect, updateStart+middle+updateEnd, args); err2 != nil {
						return rowsAffected, fmt.Errorf("%w\n"+query, updateErr)
					}
					return rowsAffected, updateErr
				}
				rowsAffected += n
			}
			result, insertErr := insertStmt.ExecContext(ctx, args...)
			if insertErr != nil {
				query, err := Sprintf(bi.Dialect, insertStart+middle+insertEnd, args)
				if err == nil {
					return rowsAffected, fmt.Errorf("%w\n"+query, insertErr)
				}
				return rowsAffected, insertErr
			}
			n, insertErr := result.RowsAffected()
			if insertErr != nil {
				return rowsAffected, insertErr
			}
			rowsAffected += n
			batchcounter = 0
			args = args[:0]
		}
	}

	// insert leftover args
	if batchcounter > 0 {
		buf.Reset()
		writeRowValues(ctx, bi.Dialect, buf, batchcounter, len(bi.Columns))
		middle = buf.String()
		if updateStmt != nil {
			result, updateErr := db.ExecContext(ctx, updateStart+middle+updateEnd, args...)
			if updateErr != nil {
				query, err := Sprintf(bi.Dialect, updateStart+middle+updateEnd, args)
				if err == nil {
					return rowsAffected, fmt.Errorf("%w\n"+query, updateErr)
				}
				return rowsAffected, updateErr
			}
			n, updateErr := result.RowsAffected()
			if updateErr != nil {
				return rowsAffected, updateErr
			}
			rowsAffected += n
		}
		result, insertErr := db.ExecContext(ctx, insertStart+middle+insertEnd, args...)
		if insertErr != nil {
			query, err := Sprintf(bi.Dialect, insertStart+middle+insertEnd, args)
			if err == nil {
				return rowsAffected, fmt.Errorf("%w\n"+query, insertErr)
			}
			return rowsAffected, insertErr
		}
		n, insertErr := result.RowsAffected()
		if insertErr != nil {
			return rowsAffected, insertErr
		}
		rowsAffected += n
	}

	// SET IDENTITY_INSERT <table> OFF
	if len(identityColumns) > 0 && bi.Dialect == DialectSQLServer {
		_, err = db.ExecContext(ctx, "SET IDENTITY_INSERT "+tableName+" OFF")
		if err != nil {
			return rowsAffected, err
		}
	}

	// SELECT setval(pg_get_serial_sequence('table', 'id'), max(id)) FROM table;
	if len(identityColumns) > 0 && bi.Dialect == DialectPostgres {
		for _, column := range identityColumns {
			columnName := QuoteIdentifier(bi.Dialect, column)
			_, err = db.ExecContext(ctx,
				"SELECT setval(pg_get_serial_sequence($1, $2), MAX("+columnName+")) FROM "+tableName,
				tableName, columnName,
			)
			if err != nil {
				return rowsAffected, err
			}
		}
	}
	return rowsAffected, nil
}

func writeRowValues(ctx context.Context, dialect string, buf *bytes.Buffer, batchsize, numFields int) {
	var counter int
	for i := 0; i < batchsize; i++ {
		if i > 0 {
			buf.WriteString(", ")
		}
		buf.WriteString("(")
		for j := 0; j < numFields; j++ {
			if j > 0 {
				buf.WriteString(", ")
			}
			counter++
			switch dialect {
			case DialectPostgres:
				buf.WriteString("$" + strconv.Itoa(counter))
			case DialectSQLServer:
				buf.WriteString("@p" + strconv.Itoa(counter))
			default:
				buf.WriteString("?")
			}
		}
		buf.WriteString(")")
	}
}