Add statement type convenience methods to CommandTag and optimize

Added convenient way to check whether a statement was a select, insert,
update, or delete. These methods do not allocate.

RowsAffected now does not allocate even when a large number of rows are
affected. It also is multiple times faster, though the absolute change
is inconsequential.

benchmark_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"os"
+	"strings"
 	"testing"
 
 	"github.com/jackc/pgconn"
@@ -252,3 +253,70 @@ func BenchmarkExecPreparedPossibleToCancel(b *testing.B) {
 // 		conn.ChanToSetDeadline().Ignore()
 // 	}
 // }
+
+func BenchmarkCommandTagRowsAffected(b *testing.B) {
+	benchmarks := []struct {
+		commandTag   string
+		rowsAffected int64
+	}{
+		{"UPDATE 1", 1},
+		{"UPDATE 123456789", 123456789},
+		{"INSERT 0 1", 1},
+		{"INSERT 0 123456789", 123456789},
+	}
+
+	for _, bm := range benchmarks {
+		ct := pgconn.CommandTag(bm.commandTag)
+		b.Run(bm.commandTag, func(b *testing.B) {
+			var n int64
+			for i := 0; i < b.N; i++ {
+				n = ct.RowsAffected()
+			}
+			if n != bm.rowsAffected {
+				b.Errorf("expected %d got %d", bm.rowsAffected, n)
+			}
+		})
+	}
+}
+
+func BenchmarkCommandTagTypeFromString(b *testing.B) {
+	ct := pgconn.CommandTag("UPDATE 1")
+
+	var update bool
+	for i := 0; i < b.N; i++ {
+		update = strings.HasPrefix(ct.String(), "UPDATE")
+	}
+	if !update {
+		b.Error("expected update")
+	}
+}
+
+func BenchmarkCommandTagInsert(b *testing.B) {
+	benchmarks := []struct {
+		commandTag string
+		is         bool
+	}{
+		{"INSERT 1", true},
+		{"INSERT 1234567890", true},
+		{"UPDATE 1", false},
+		{"UPDATE 1234567890", false},
+		{"DELETE 1", false},
+		{"DELETE 1234567890", false},
+		{"SELECT 1", false},
+		{"SELECT 1234567890", false},
+		{"UNKNOWN 1234567890", false},
+	}
+
+	for _, bm := range benchmarks {
+		ct := pgconn.CommandTag(bm.commandTag)
+		b.Run(bm.commandTag, func(b *testing.B) {
+			var is bool
+			for i := 0; i < b.N; i++ {
+				is = ct.Insert()
+			}
+			if is != bm.is {
+				b.Errorf("expected %v got %v", bm.is, is)
+			}
+		})
+	}
+}

pgconn.go

@@ -1,7 +1,6 @@
 package pgconn
 
 import (
-	"bytes"
 	"context"
 	"crypto/md5"
 	"crypto/tls"
@@ -10,7 +9,6 @@ import (
 	"io"
 	"math"
 	"net"
-	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -579,11 +577,25 @@ type CommandTag []byte
 // RowsAffected returns the number of rows affected. If the CommandTag was not
 // for a row affecting command (e.g. "CREATE TABLE") then it returns 0.
 func (ct CommandTag) RowsAffected() int64 {
-	idx := bytes.LastIndexByte([]byte(ct), ' ')
+	// Find last non-digit
+	idx := -1
+	for i := len(ct) - 1; i >= 0; i-- {
+		if ct[i] >= '0' && ct[i] <= '9' {
+			idx = i
+		} else {
+			break
+		}
+	}
+
 	if idx == -1 {
 		return 0
 	}
-	n, _ := strconv.ParseInt(string([]byte(ct)[idx+1:]), 10, 64)
+
+	var n int64
+	for _, b := range ct[idx:] {
+		n = n*10 + int64(b-'0')
+	}
+
 	return n
 }
 
@@ -591,6 +603,50 @@ func (ct CommandTag) String() string {
 	return string(ct)
 }
 
+// Insert is true if the command tag starts with "INSERT".
+func (ct CommandTag) Insert() bool {
+	return len(ct) >= 6 &&
+		ct[0] == 'I' &&
+		ct[1] == 'N' &&
+		ct[2] == 'S' &&
+		ct[3] == 'E' &&
+		ct[4] == 'R' &&
+		ct[5] == 'T'
+}
+
+// Update is true if the command tag starts with "UPDATE".
+func (ct CommandTag) Update() bool {
+	return len(ct) >= 6 &&
+		ct[0] == 'U' &&
+		ct[1] == 'P' &&
+		ct[2] == 'D' &&
+		ct[3] == 'A' &&
+		ct[4] == 'T' &&
+		ct[5] == 'E'
+}
+
+// Delete is true if the command tag starts with "DELETE".
+func (ct CommandTag) Delete() bool {
+	return len(ct) >= 6 &&
+		ct[0] == 'D' &&
+		ct[1] == 'E' &&
+		ct[2] == 'L' &&
+		ct[3] == 'E' &&
+		ct[4] == 'T' &&
+		ct[5] == 'E'
+}
+
+// Select is true if the command tag starts with "SELECT".
+func (ct CommandTag) Select() bool {
+	return len(ct) >= 6 &&
+		ct[0] == 'S' &&
+		ct[1] == 'E' &&
+		ct[2] == 'L' &&
+		ct[3] == 'E' &&
+		ct[4] == 'C' &&
+		ct[5] == 'T'
+}
+
 type StatementDescription struct {
 	Name      string
 	SQL       string

pgconn_test.go

@@ -973,20 +973,31 @@ func TestCommandTag(t *testing.T) {
 	var tests = []struct {
 		commandTag   pgconn.CommandTag
 		rowsAffected int64
+		isInsert     bool
+		isUpdate     bool
+		isDelete     bool
+		isSelect     bool
 	}{
-		{commandTag: pgconn.CommandTag("INSERT 0 5"), rowsAffected: 5},
+		{commandTag: pgconn.CommandTag("INSERT 0 5"), rowsAffected: 5, isInsert: true},
-		{commandTag: pgconn.CommandTag("UPDATE 0"), rowsAffected: 0},
+		{commandTag: pgconn.CommandTag("UPDATE 0"), rowsAffected: 0, isUpdate: true},
-		{commandTag: pgconn.CommandTag("UPDATE 1"), rowsAffected: 1},
+		{commandTag: pgconn.CommandTag("UPDATE 1"), rowsAffected: 1, isUpdate: true},
-		{commandTag: pgconn.CommandTag("DELETE 0"), rowsAffected: 0},
+		{commandTag: pgconn.CommandTag("DELETE 0"), rowsAffected: 0, isDelete: true},
-		{commandTag: pgconn.CommandTag("DELETE 1"), rowsAffected: 1},
+		{commandTag: pgconn.CommandTag("DELETE 1"), rowsAffected: 1, isDelete: true},
+		{commandTag: pgconn.CommandTag("DELETE 1234567890"), rowsAffected: 1234567890, isDelete: true},
+		{commandTag: pgconn.CommandTag("SELECT 1"), rowsAffected: 1, isSelect: true},
+		{commandTag: pgconn.CommandTag("SELECT 99999999999"), rowsAffected: 99999999999, isSelect: true},
 		{commandTag: pgconn.CommandTag("CREATE TABLE"), rowsAffected: 0},
 		{commandTag: pgconn.CommandTag("ALTER TABLE"), rowsAffected: 0},
 		{commandTag: pgconn.CommandTag("DROP TABLE"), rowsAffected: 0},
 	}
 
 	for i, tt := range tests {
-		actual := tt.commandTag.RowsAffected()
+		ct := tt.commandTag
-		assert.Equalf(t, tt.rowsAffected, actual, "%d. %v", i, tt.commandTag)
+		assert.Equalf(t, tt.rowsAffected, ct.RowsAffected(), "%d. %v", i, tt.commandTag)
+		assert.Equalf(t, tt.isInsert, ct.Insert(), "%d. %v", i, tt.commandTag)
+		assert.Equalf(t, tt.isUpdate, ct.Update(), "%d. %v", i, tt.commandTag)
+		assert.Equalf(t, tt.isDelete, ct.Delete(), "%d. %v", i, tt.commandTag)
+		assert.Equalf(t, tt.isSelect, ct.Select(), "%d. %v", i, tt.commandTag)
 	}
 }