Make BatchResults.Close safe to be called multiple times
https://github.com/jackc/pgx/issues/1138 https://github.com/jackc/pgx/issues/938
This commit is contained in:
@@ -3,6 +3,7 @@ package pgx
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/jackc/pgconn"
|
||||
)
|
||||
@@ -46,17 +47,18 @@ type BatchResults interface {
|
||||
|
||||
// Close closes the batch operation. This must be called before the underlying connection can be used again. Any error
|
||||
// that occurred during a batch operation may have made it impossible to resyncronize the connection with the server.
|
||||
// In this case the underlying connection will have been closed.
|
||||
// In this case the underlying connection will have been closed. Close is safe to call multiple times.
|
||||
Close() error
|
||||
}
|
||||
|
||||
type batchResults struct {
|
||||
ctx context.Context
|
||||
conn *Conn
|
||||
mrr *pgconn.MultiResultReader
|
||||
err error
|
||||
b *Batch
|
||||
ix int
|
||||
ctx context.Context
|
||||
conn *Conn
|
||||
mrr *pgconn.MultiResultReader
|
||||
err error
|
||||
b *Batch
|
||||
ix int
|
||||
closed bool
|
||||
}
|
||||
|
||||
// Exec reads the results from the next query in the batch as if the query has been sent with Exec.
|
||||
@@ -64,6 +66,9 @@ func (br *batchResults) Exec() (pgconn.CommandTag, error) {
|
||||
if br.err != nil {
|
||||
return nil, br.err
|
||||
}
|
||||
if br.closed {
|
||||
return nil, fmt.Errorf("batch already closed")
|
||||
}
|
||||
|
||||
query, arguments, _ := br.nextQueryAndArgs()
|
||||
|
||||
@@ -114,6 +119,11 @@ func (br *batchResults) Query() (Rows, error) {
|
||||
return &connRows{err: br.err, closed: true}, br.err
|
||||
}
|
||||
|
||||
if br.closed {
|
||||
alreadyClosedErr := fmt.Errorf("batch already closed")
|
||||
return &connRows{err: alreadyClosedErr, closed: true}, alreadyClosedErr
|
||||
}
|
||||
|
||||
rows := br.conn.getRows(br.ctx, query, arguments)
|
||||
|
||||
if !br.mrr.NextResult() {
|
||||
@@ -140,6 +150,10 @@ func (br *batchResults) Query() (Rows, error) {
|
||||
|
||||
// QueryFunc reads the results from the next query in the batch as if the query has been sent with Conn.QueryFunc.
|
||||
func (br *batchResults) QueryFunc(scans []interface{}, f func(QueryFuncRow) error) (pgconn.CommandTag, error) {
|
||||
if br.closed {
|
||||
return nil, fmt.Errorf("batch already closed")
|
||||
}
|
||||
|
||||
rows, err := br.Query()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -179,6 +193,11 @@ func (br *batchResults) Close() error {
|
||||
return br.err
|
||||
}
|
||||
|
||||
if br.closed {
|
||||
return nil
|
||||
}
|
||||
br.closed = true
|
||||
|
||||
// log any queries that haven't yet been logged by Exec or Query
|
||||
for {
|
||||
query, args, ok := br.nextQueryAndArgs()
|
||||
|
||||
Reference in New Issue
Block a user