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

@@ -9,27 +9,45 @@ import (
 )
 
 type batch struct {
-	ds    *Datastore
-	batch *pgx.Batch
+	ds           *Datastore
+	batch        *pgx.Batch
+	maxBatchSize uint16
 }
 
 // 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{}, maxBatchSize: 0}
+	b.batch.Queue("BEGIN")
+	return b, nil
+}
+
+// Set the max batch size (0 or default means unlimited - the batch is only cleared when calling Commit)
+func (b *batch) SetMaxBatchSize(size uint16) {
+	b.maxBatchSize = size
+}
+
+func (b *batch) checkMaxBatchSize() error {
+	var err error
+
+	if b.maxBatchSize != 0 && b.batch.Len() >= int(b.maxBatchSize) {
+		err = b.CommitContext(context.Background())
+	}
+
+	if err != nil {
+		b.batch = &pgx.Batch{}
+	}
+
+	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.checkMaxBatchSize()
 }
 
 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
 }
 
@@ -38,6 +56,7 @@ func (b *batch) Commit() error {
 }
 
 func (b *batch) CommitContext(ctx context.Context) error {
+	b.batch.Queue("COMMIT")
 	res := b.ds.pool.SendBatch(ctx, b.batch)
 	defer res.Close()
 

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",
 	)