Add CommandTag to Rows interface

This allows handling queries where it is unknown if there will be a
result set or not. If it is not a result set returning query the
command tag will still be available.
pull/594/head
Jack Christensen 2019-09-02 09:53:26 -05:00
parent 0ef89ae0b7
commit b5ce0220f8
3 changed files with 43 additions and 11 deletions

View File

@ -1,6 +1,7 @@
package pgxpool package pgxpool
import ( import (
"github.com/jackc/pgconn"
"github.com/jackc/pgproto3/v2" "github.com/jackc/pgproto3/v2"
"github.com/jackc/pgx/v4" "github.com/jackc/pgx/v4"
) )
@ -11,6 +12,7 @@ type errRows struct {
func (errRows) Close() {} func (errRows) Close() {}
func (e errRows) Err() error { return e.err } func (e errRows) Err() error { return e.err }
func (errRows) CommandTag() pgconn.CommandTag { return nil }
func (errRows) FieldDescriptions() []pgproto3.FieldDescription { return nil } func (errRows) FieldDescriptions() []pgproto3.FieldDescription { return nil }
func (errRows) Next() bool { return false } func (errRows) Next() bool { return false }
func (e errRows) Scan(dest ...interface{}) error { return e.err } func (e errRows) Scan(dest ...interface{}) error { return e.err }
@ -43,6 +45,10 @@ func (rows *poolRows) Err() error {
return rows.r.Err() return rows.r.Err()
} }
func (rows *poolRows) CommandTag() pgconn.CommandTag {
return rows.CommandTag()
}
func (rows *poolRows) FieldDescriptions() []pgproto3.FieldDescription { func (rows *poolRows) FieldDescriptions() []pgproto3.FieldDescription {
return rows.r.FieldDescriptions() return rows.r.FieldDescriptions()
} }

View File

@ -48,6 +48,8 @@ func TestConnQueryScan(t *testing.T) {
t.Fatalf("conn.Query failed: %v", err) t.Fatalf("conn.Query failed: %v", err)
} }
assert.Equal(t, "SELECT 10", string(rows.CommandTag()))
if rowCount != 10 { if rowCount != 10 {
t.Error("Select called onDataRow wrong number of times") t.Error("Select called onDataRow wrong number of times")
} }
@ -56,6 +58,19 @@ func TestConnQueryScan(t *testing.T) {
} }
} }
func TestConnQueryWithoutResultSetCommandTag(t *testing.T) {
t.Parallel()
conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
defer closeConn(t, conn)
rows, err := conn.Query(context.Background(), "create temporary table t (id serial);")
assert.NoError(t, err)
rows.Close()
assert.NoError(t, rows.Err())
assert.Equal(t, "CREATE TABLE", string(rows.CommandTag()))
}
func TestConnQueryScanWithManyColumns(t *testing.T) { func TestConnQueryScanWithManyColumns(t *testing.T) {
t.Parallel() t.Parallel()

33
rows.go
View File

@ -21,7 +21,12 @@ type Rows interface {
// to call Close after rows is already closed. // to call Close after rows is already closed.
Close() Close()
// Err returns any error that occurred while reading.
Err() error Err() error
// CommandTag returns the command tag from this query. It is only available after Rows is closed.
CommandTag() pgconn.CommandTag
FieldDescriptions() []pgproto3.FieldDescription FieldDescriptions() []pgproto3.FieldDescription
// Next prepares the next row for reading. It returns true if there is another // Next prepares the next row for reading. It returns true if there is another
@ -76,16 +81,17 @@ type rowLog interface {
// connRows implements the Rows interface for Conn.Query. // connRows implements the Rows interface for Conn.Query.
type connRows struct { type connRows struct {
ctx context.Context ctx context.Context
logger rowLog logger rowLog
connInfo *pgtype.ConnInfo connInfo *pgtype.ConnInfo
values [][]byte values [][]byte
rowCount int rowCount int
err error err error
startTime time.Time commandTag pgconn.CommandTag
sql string startTime time.Time
args []interface{} sql string
closed bool args []interface{}
closed bool
resultReader *pgconn.ResultReader resultReader *pgconn.ResultReader
multiResultReader *pgconn.MultiResultReader multiResultReader *pgconn.MultiResultReader
@ -103,7 +109,8 @@ func (rows *connRows) Close() {
rows.closed = true rows.closed = true
if rows.resultReader != nil { if rows.resultReader != nil {
_, closeErr := rows.resultReader.Close() var closeErr error
rows.commandTag, closeErr = rows.resultReader.Close()
if rows.err == nil { if rows.err == nil {
rows.err = closeErr rows.err = closeErr
} }
@ -128,6 +135,10 @@ func (rows *connRows) Close() {
} }
} }
func (rows *connRows) CommandTag() pgconn.CommandTag {
return rows.commandTag
}
func (rows *connRows) Err() error { func (rows *connRows) Err() error {
return rows.err return rows.err
} }