Do not commit on each update when batching, and allow for setting the max batch size

README.md

@@ -54,6 +54,16 @@ func main() {
 }
 ```
 
+## Running the tests
+
+For running the tests you will need a Postgres database listening on localhost. The easiest way of getting one is by using Docker:
+
+``` sh
+docker run --rm -p 5432:5432 -e POSTGRES_PASSWORD=postgres postgres:11.14-alpine &
+```
+
+Running the full test suite will take a long time, so you will have to run `go test -timeout 60m`
+
 ## API
 
 [GoDoc Reference](https://godoc.org/github.com/alanshaw/ipfs-ds-postgres)

batching.go

@@ -15,29 +15,43 @@ type batch struct {
 
 // Batch creates a set of deferred updates to the database.
 func (d *Datastore) Batch() (ds.Batch, error) {
-	return &batch{ds: d, batch: &pgx.Batch{}}, nil
+	b := &batch{ds: d, batch: &pgx.Batch{}}
+	b.batch.Queue("BEGIN")
+	return b, nil
+}
+
+func (b *batch) commitIfBatchFull() error {
+	var err error
+
+	if b.ds.maxBatchSize != 0 && b.batch.Len() >= int(b.ds.maxBatchSize) {
+		err = b.CommitContext(context.Background())
+	}
+
+	if err != nil {
+		b.batch = &pgx.Batch{}
+		b.batch.Queue("BEGIN")
+	}
+
+	return err
 }
 
 func (b *batch) Put(key ds.Key, value []byte) error {
-	b.batch.Queue("BEGIN")
 	sql := fmt.Sprintf("INSERT INTO %s (key, data) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET data = $2", b.ds.table)
 	b.batch.Queue(sql, key.String(), value)
-	b.batch.Queue("COMMIT")
-	return nil
+	return b.commitIfBatchFull()
 }
 
 func (b *batch) Delete(key ds.Key) error {
-	b.batch.Queue("BEGIN")
 	b.batch.Queue(fmt.Sprintf("DELETE FROM %s WHERE key = $1", b.ds.table), key.String())
-	b.batch.Queue("COMMIT")
-	return nil
+	return b.commitIfBatchFull()
 }
 
 func (b *batch) Commit() error {
 	return b.CommitContext(context.Background())
 }
 
 func (b *batch) CommitContext(ctx context.Context) error {
+	b.batch.Queue("COMMIT")
 	res := b.ds.pool.SendBatch(ctx, b.batch)
 	defer res.Close()
 

datastore.go

@@ -12,8 +12,9 @@ import (
 
 // Datastore is a PostgreSQL backed datastore.
 type Datastore struct {
-	table string
-	pool  *pgxpool.Pool
+	table        string
+	pool         *pgxpool.Pool
+	maxBatchSize uint16
 }
 
 // NewDatastore creates a new PostgreSQL datastore
@@ -26,7 +27,7 @@ func NewDatastore(ctx context.Context, connString string, options ...Option) (*D
 		return nil, err
 	}
 
-	return &Datastore{table: cfg.Table, pool: pool}, nil
+	return &Datastore{table: cfg.Table, pool: pool, maxBatchSize: cfg.MaxBatchSize}, nil
 }
 
 // PgxPool exposes the underlying pool of connections to Postgres.

datastore_test.go

@@ -27,7 +27,7 @@ func initPG(t *testing.T) {
 		connConf, err := pgx.ParseConfig(fmt.Sprintf(
 			"postgres://%s:%s@%s/%s?sslmode=disable",
 			envString(t, "PG_USER", "postgres"),
-			envString(t, "PG_PASS", ""),
+			envString(t, "PG_PASS", "postgres"),
 			envString(t, "PG_HOST", "127.0.0.1"),
 			envString(t, "PG_DB", envString(t, "PG_USER", "postgres")),
 		))
@@ -62,7 +62,7 @@ func newDS(t *testing.T) (*Datastore, func()) {
 	connString := fmt.Sprintf(
 		"postgres://%s:%s@%s/%s?sslmode=disable",
 		envString(t, "PG_USER", "postgres"),
-		envString(t, "PG_PASS", ""),
+		envString(t, "PG_PASS", "postgres"),
 		envString(t, "PG_HOST", "127.0.0.1"),
 		"test_datastore",
 	)

options.go

@@ -6,7 +6,8 @@ import (
 
 // Options are Datastore options
 type Options struct {
-	Table string
+	Table        string
+	MaxBatchSize uint16
 }
 
 // Option is the Datastore option type.
@@ -26,6 +27,7 @@ func (o *Options) Apply(opts ...Option) error {
 // prepended to any options you pass to the Hydra Head constructor.
 var OptionDefaults = func(o *Options) error {
 	o.Table = "blocks"
+	o.MaxBatchSize = 0
 	return nil
 }
 
@@ -39,3 +41,12 @@ func Table(t string) Option {
 		return nil
 	}
 }
+
+// MaxBatchSize sets the maximum number of updates that will be batched before committing.
+// Default to 0, which means that the batch is only commited when Commit() is explicitly called.
+func MaxBatchSize(size uint16) Option {
+	return func(o *Options) error {
+		o.MaxBatchSize = size
+		return nil
+	}
+}