add support for bulk insert & connection pooling

Author: thatmattlove

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/gookit/gcli/v3 v3.0.1
 	github.com/hako/durafmt v0.0.0-20210608085754-5c1018a4e16b
 	github.com/jedib0t/go-pretty v4.3.0+incompatible
+	github.com/lib/pq v1.10.9
 	github.com/stretchr/testify v1.7.0
 	github.com/thatmattlove/go-macaddr v0.0.7
 	github.com/urfave/cli/v2 v2.3.0
@@ -25,12 +26,10 @@ require (
 	github.com/gookit/color v1.5.0 // indirect
 	github.com/gookit/goutil v0.4.3 // indirect
 	github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect
-	github.com/lib/pq v1.10.9 // indirect
 	github.com/mattn/go-isatty v0.0.14 // indirect
 	github.com/mattn/go-runewidth v0.0.13 // indirect
 	github.com/mitchellh/mapstructure v1.4.3 // indirect
 	github.com/oklog/ulid v1.3.1 // indirect
-	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 // indirect
 	github.com/rivo/uniseg v0.2.0 // indirect

go.sum

@@ -79,7 +79,6 @@ github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e h1:fD57ERR4JtEqsWb
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
 github.com/oklog/ulid v1.3.1 h1:EGfNDEx6MqHz8B3uNV6QAib1UR2Lm97sHi3ocA6ESJ4=
 github.com/oklog/ulid v1.3.1/go.mod h1:CirwcVhetQ6Lv90oh/F+FBtV6XMibvdAFo93nm5qn4U=
-github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=

internal/util/util.go

@@ -43,3 +43,19 @@ func PathExists(n string) bool {
 func TimeSince(t time.Time) string {
 	return durafmt.Parse(time.Since(t)).LimitFirstN(1).String()
 }
+
+func SplitSlice[T any](slice []T, max int) [][]T {
+	result := make([][]T, 0)
+
+	for i := 0; i < len(slice); i += max {
+		end := i + max
+		if end > len(slice) {
+			end = len(slice)
+		}
+		part := slice[i:end]
+		if part != nil {
+			result = append(result, part)
+		}
+	}
+	return result
+}

internal/util/util_test.go

@@ -2,6 +2,7 @@ package util_test
 
 import (
 	"os"
+	"reflect"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -67,3 +68,39 @@ func Test_pathExists(t *testing.T) {
 		assert.True(t, r)
 	})
 }
+
+func Test_SplitSlice(t *testing.T) {
+	t.Run("equal parts", func(t *testing.T) {
+		t.Parallel()
+		original := []int{1, 2, 3, 4, 5, 6}
+		max := 2
+		expected := [][]int{
+			{1, 2},
+			{3, 4},
+			{5, 6},
+		}
+		result := util.SplitSlice(original, max)
+		assert.True(t, reflect.DeepEqual(expected, result))
+	})
+	t.Run("unequal parts", func(t *testing.T) {
+		t.Parallel()
+		original := []int{1, 2, 3, 4, 5, 6}
+		max := 4
+		expected := [][]int{
+			{1, 2, 3, 4},
+			{5, 6},
+		}
+		result := util.SplitSlice(original, max)
+		assert.True(t, reflect.DeepEqual(expected, result))
+	})
+
+	t.Run("empty", func(t *testing.T) {
+		t.Parallel()
+		original := []string{}
+		max := 3
+		expected := [][]string{}
+		result := util.SplitSlice(original, max)
+		assert.Equal(t, expected, result)
+	})
+
+}

main.go

@@ -17,7 +17,7 @@ func getArgs() []string {
 
 func main() {
 	args := getArgs()
-	err := cmd.New(TABLE_VERSION).Run(args)
+	err := cmd.New(Version).Run(args)
 	if err != nil {
 		fmt.Println(err)
 		os.Exit(1)

oui/db.go

@@ -4,9 +4,11 @@ import (
 	"database/sql"
 	"fmt"
 	"strconv"
+	"strings"
 
 	"github.com/gookit/gcli/v3/progress"
 	"github.com/thatmattlove/go-macaddr"
+	"github.com/thatmattlove/oui/v2/internal/util"
 )
 
 type OUIDB struct {
@@ -81,6 +83,59 @@ func (ouidb *OUIDB) Insert(d *VendorDef) (res sql.Result, err error) {
 	return
 }
 
+func (ouidb *OUIDB) BulkInsert(defs []*VendorDef) (int64, error) {
+
+	var statement string
+	var tmpl string
+	var maxRecords int
+	switch ouidb.dialect {
+	case dialectSqlite:
+		tmpl = "(?,?,?,?)"
+		statement = "INSERT INTO %s(prefix, length, org, registry) values%s"
+		maxRecords = maxVarsSqlite
+	case dialectPsql:
+		tmpl = "($%d,$%d,$%d,$%d)"
+		statement = "INSERT INTO %s(prefix, length, org, registry) values%s ON CONFLICT (prefix, length, registry) DO UPDATE SET prefix = excluded.prefix, length = excluded.length, registry = excluded.registry"
+		maxRecords = maxVarsPsql
+	}
+
+	splitDefs := util.SplitSlice(defs, maxRecords/4)
+	inserted := int64(0)
+
+	for _, split := range splitDefs {
+		placeholders := make([]string, 0, len(split))
+		args := make([]interface{}, 0, len(split)*4)
+		i := 0
+		for _, def := range split {
+			def := def
+			var placeholder string
+			switch ouidb.dialect {
+			case dialectSqlite:
+				placeholder = tmpl
+			case dialectPsql:
+				placeholder = fmt.Sprintf(tmpl, i*4+1, i*4+2, i*4+3, i*4+4)
+			}
+			placeholders = append(placeholders, placeholder)
+			args = append(args, def.Prefix, def.Length, def.Org, def.Registry)
+			i++
+		}
+		q := fmt.Sprintf(statement, ouidb.Version, strings.Join(placeholders, ","))
+		res, err := ouidb.Connection.Exec(q, args...)
+		if err != nil {
+			return inserted, err
+		}
+		rows, err := res.RowsAffected()
+		if err != nil {
+			return inserted, err
+		}
+		inserted += rows
+		if err != nil {
+			return inserted, err
+		}
+	}
+	return inserted, nil
+}
+
 func (ouidb *OUIDB) Populate() (records int64, err error) {
 	err = ouidb.Clear()
 	if err != nil {
@@ -98,12 +153,9 @@ func (ouidb *OUIDB) Populate() (records int64, err error) {
 	if err != nil {
 		return
 	}
-	for _, def := range defs {
-		_, err = ouidb.Insert(def)
-		if err != nil {
-			return
-		}
-		records++
+	records, err = ouidb.BulkInsert(defs)
+	if err != nil {
+		return
 	}
 	return
 }
@@ -184,6 +236,7 @@ func New(opts ...Option) (*OUIDB, error) {
 		}
 	} else if options.dialect == dialectPsql {
 		q := fmt.Sprintf("CREATE TABLE IF NOT EXISTS %v ( id SERIAL PRIMARY KEY, prefix VARCHAR(32) NOT NULL, length INT NOT NULL, org VARCHAR(256) NOT NULL, registry VARCHAR(32) NOT NULL, UNIQUE(prefix, length, registry) )", options.Version)
+		options.Connection.SetMaxOpenConns(int(options.MaxConnections))
 		_, err := options.Connection.Exec(q)
 		if err != nil {
 			return nil, err

oui/db_test.go

@@ -20,7 +20,7 @@ func Test_New(t *testing.T) {
 	t.Run("postgres", func(t *testing.T) {
 		t.Parallel()
 		password := os.Getenv("POSTGRES_PASSWORD")
-		require.NotEqual(t, "", password)
+		require.NotEqual(t, "", password, "missing POSTGRES_PASSWORD environment variable")
 		cs := fmt.Sprintf("postgresql://oui:%s@localhost/oui?sslmode=disable", password)
 		psql, err := oui.CreatePostgresOption(cs)
 		require.NoError(t, err)

oui/options.go

@@ -12,11 +12,12 @@ import (
 )
 
 type Options struct {
-	Logger     *LoggerType
-	Progress   *progress.Progress
-	Version    string
-	Connection *sql.DB
-	dialect    int
+	Logger         *LoggerType
+	Progress       *progress.Progress
+	Version        string
+	Connection     *sql.DB
+	dialect        int
+	MaxConnections uint
 }
 
 type Option func(*Options)
@@ -45,12 +46,19 @@ func WithConnection(conn *sql.DB) Option {
 	}
 }
 
+func WithMaxConnections(max uint) Option {
+	return func(opts *Options) {
+		opts.MaxConnections = max
+	}
+}
+
 func getOptions(setters ...Option) *Options {
 	options := &Options{
-		Logger:     nil,
-		Progress:   nil,
-		Version:    "default",
-		Connection: nil,
+		Logger:         nil,
+		Progress:       nil,
+		Version:        "default",
+		Connection:     nil,
+		MaxConnections: 0,
 	}
 	for _, setter := range setters {
 		setter(options)

oui/types.go

@@ -26,3 +26,8 @@ const (
 	dialectSqlite int = iota
 	dialectPsql
 )
+
+const (
+	maxVarsSqlite int = 999
+	maxVarsPsql   int = 65535
+)