From fd2093cef8e97839e11bb13bb4a2c1b805ae62f5 Mon Sep 17 00:00:00 2001 From: Jack Christensen Date: Sat, 11 Jan 2020 18:42:31 -0600 Subject: [PATCH] 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 | 68 +++++++++++++++++++++++++++++++++++++++++++++++ pgconn.go | 64 +++++++++++++++++++++++++++++++++++++++++--- pgconn_test.go | 25 ++++++++++++----- 3 files changed, 146 insertions(+), 11 deletions(-) diff --git a/benchmark_test.go b/benchmark_test.go index 3295a90f..ced785b6 100644 --- a/benchmark_test.go +++ b/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) + } + }) + } +} diff --git a/pgconn.go b/pgconn.go index c46dc6a6..dce4bfb5 100644 --- a/pgconn.go +++ b/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 diff --git a/pgconn_test.go b/pgconn_test.go index 7ae6fdc5..2c303d81 100644 --- a/pgconn_test.go +++ b/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("UPDATE 0"), rowsAffected: 0}, - {commandTag: pgconn.CommandTag("UPDATE 1"), rowsAffected: 1}, - {commandTag: pgconn.CommandTag("DELETE 0"), rowsAffected: 0}, - {commandTag: pgconn.CommandTag("DELETE 1"), rowsAffected: 1}, + {commandTag: pgconn.CommandTag("INSERT 0 5"), rowsAffected: 5, isInsert: true}, + {commandTag: pgconn.CommandTag("UPDATE 0"), rowsAffected: 0, isUpdate: true}, + {commandTag: pgconn.CommandTag("UPDATE 1"), rowsAffected: 1, isUpdate: true}, + {commandTag: pgconn.CommandTag("DELETE 0"), rowsAffected: 0, isDelete: true}, + {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() - assert.Equalf(t, tt.rowsAffected, actual, "%d. %v", i, tt.commandTag) + ct := 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) } }