safe_pgx_batch.md

pgx.Batch is not thread-safe

Problem

Batch filling often happens in goroutines, but pgx.Batch stores queries in a simple slice and Queue() method just appends to it. So some queries will be missing.

See https://github.com/jackc/pgx/blob/dc0ad04ff58f72f4819289f54745a36124cdbec3/batch.go#L16:

// Batch queries are a way of bundling multiple queries together to avoid
// unnecessary network round trips.
type Batch struct {
	items []*batchItem
}

// Queue queues a query to batch b. query can be an SQL query or the name of a prepared statement.
func (b *Batch) Queue(query string, arguments ...interface{}) {
	b.items = append(b.items, &batchItem{
		query:     query,
		arguments: arguments,
	})
}

Solution

Just wrap it with your struct and add sync.Mutex:

type safeBatch struct {
	mu sync.Mutex
	pgx.Batch
}

And use mutex anywhere around Queue():

sb := &safeBatch{}
var wg sync.WaitGroup
for i := 1; i <= 5; i++ {
	wg.Add(1)
	go func(i int) {
		defer wg.Done()
		sb.mu.Lock()
		sb.Batch.Queue(`select $1`, i)
		sb.mu.Unlock()
	}(i)
}
wg.Wait()
print(sb.Len()) // always 5