Skip to content

Commit

Permalink
add support for bulk insert & connection pooling
Browse files Browse the repository at this point in the history
thatmattlove committed Sep 24, 2023

Verified

This commit was signed with the committer’s verified signature.
joshcooper Josh Cooper
1 parent 1b51297 commit 5361e34
Showing 9 changed files with 137 additions and 20 deletions.
3 changes: 1 addition & 2 deletions go.mod
Original file line number Diff line number Diff line change
@@ -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
1 change: 0 additions & 1 deletion go.sum
Original file line number Diff line number Diff line change
@@ -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=
16 changes: 16 additions & 0 deletions internal/util/util.go
Original file line number Diff line number Diff line change
@@ -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
}
37 changes: 37 additions & 0 deletions internal/util/util_test.go
Original file line number Diff line number Diff line change
@@ -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)
})

}
2 changes: 1 addition & 1 deletion main.go
Original file line number Diff line number Diff line change
@@ -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)
65 changes: 59 additions & 6 deletions oui/db.go
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion oui/db_test.go
Original file line number Diff line number Diff line change
@@ -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)
26 changes: 17 additions & 9 deletions oui/options.go
Original file line number Diff line number Diff line change
@@ -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)
5 changes: 5 additions & 0 deletions oui/types.go
Original file line number Diff line number Diff line change
@@ -26,3 +26,8 @@ const (
dialectSqlite int = iota
dialectPsql
)

const (
maxVarsSqlite int = 999
maxVarsPsql int = 65535
)

0 comments on commit 5361e34

Please sign in to comment.