sqlreplay, lex: fix replaying with readonly causes inconsistent prepared statement IDs (#757)

Author: djshow832

lib/util/waitgroup/waitgroup.go

@@ -47,7 +47,7 @@ func recoverFromErr(wg *sync.WaitGroup, recoverFn func(r interface{}), logger *z
 	}()
 	if r != nil && logger != nil {
 		logger.Error("panic in the recoverable goroutine",
-			zap.Reflect("r", r),
+			zap.Any("err", r),
 			zap.Stack("stack trace"))
 	}
 	// Call Done() before recoverFn because recoverFn normally calls `Close()`, which may call `wg.Wait()`.

pkg/sqlreplay/cmd/cmd.go

@@ -14,7 +14,6 @@ import (
 	"github.com/pingcap/tidb/pkg/parser"
 	"github.com/pingcap/tiproxy/lib/util/errors"
 	pnet "github.com/pingcap/tiproxy/pkg/proxy/net"
-	"github.com/pingcap/tiproxy/pkg/util/lex"
 	"github.com/siddontang/go/hack"
 )
 
@@ -222,18 +221,6 @@ func (c *Command) QueryText() string {
 	return ""
 }
 
-func (c *Command) ReadOnly() bool {
-	switch c.Type {
-	case pnet.ComQuery, pnet.ComStmtPrepare:
-		return lex.IsReadOnly(c.QueryText())
-	case pnet.ComStmtExecute, pnet.ComStmtClose, pnet.ComStmtSendLongData, pnet.ComStmtReset, pnet.ComStmtFetch:
-		return lex.IsReadOnly(c.PreparedStmt)
-	case pnet.ComCreateDB, pnet.ComDropDB, pnet.ComDelayedInsert:
-		return false
-	}
-	return true
-}
-
 func writeString(key, value string, writer *bytes.Buffer) error {
 	var err error
 	if _, err = writer.WriteString(key); err != nil {

pkg/sqlreplay/cmd/cmd_test.go

@@ -152,58 +152,3 @@ func TestDigest(t *testing.T) {
 	require.Equal(t, cmd1.Digest(), cmd6.Digest())
 	require.Equal(t, "select ?", cmd6.QueryText())
 }
-
-func TestReadOnly(t *testing.T) {
-	tests := []struct {
-		cmd         pnet.Command
-		payload     []byte
-		prepareStmt string
-		readOnly    bool
-	}{
-		{
-			cmd:      pnet.ComQuery,
-			payload:  []byte("select 1"),
-			readOnly: true,
-		},
-		{
-			cmd:      pnet.ComQuery,
-			payload:  []byte("insert into t value(1)"),
-			readOnly: false,
-		},
-		{
-			cmd:      pnet.ComStmtPrepare,
-			payload:  []byte("select ?"),
-			readOnly: true,
-		},
-		{
-			cmd:         pnet.ComStmtExecute,
-			prepareStmt: "select ?",
-			readOnly:    true,
-		},
-		{
-			cmd:         pnet.ComStmtExecute,
-			prepareStmt: "insert into t value(?)",
-			readOnly:    false,
-		},
-		{
-			cmd:      pnet.ComStmtExecute,
-			readOnly: false,
-		},
-		{
-			cmd:         pnet.ComStmtClose,
-			prepareStmt: "select ?",
-			readOnly:    true,
-		},
-		{
-			cmd:      pnet.ComQuit,
-			readOnly: true,
-		},
-	}
-
-	for i, test := range tests {
-		packet := append([]byte{byte(test.cmd)}, test.payload...)
-		cmd := NewCommand(packet, time.Time{}, 100)
-		cmd.PreparedStmt = test.prepareStmt
-		require.Equal(t, test.readOnly, cmd.ReadOnly(), "case %d", i)
-	}
-}

pkg/sqlreplay/conn/conn.go

@@ -15,6 +15,8 @@ import (
 	"github.com/pingcap/tiproxy/pkg/proxy/backend"
 	pnet "github.com/pingcap/tiproxy/pkg/proxy/net"
 	"github.com/pingcap/tiproxy/pkg/sqlreplay/cmd"
+	"github.com/pingcap/tiproxy/pkg/util/lex"
+	"github.com/siddontang/go/hack"
 	"go.uber.org/zap"
 )
 
@@ -107,8 +109,7 @@ func (c *conn) Run(ctx context.Context) {
 				break
 			}
 			if c.readonly {
-				c.updateCmdForExecuteStmt(command.Value)
-				if !command.Value.ReadOnly() {
+				if !c.isReadOnly(command.Value) {
 					c.replayStats.FilteredCmds.Add(1)
 					continue
 				}
@@ -131,8 +132,30 @@ func (c *conn) Run(ctx context.Context) {
 	}
 }
 
-// updateCmdForExecuteStmt may be called multiple times, avoid duplicated works.
+func (c *conn) isReadOnly(command *cmd.Command) bool {
+	switch command.Type {
+	case pnet.ComQuery:
+		return lex.IsReadOnly(hack.String(command.Payload[1:]))
+	case pnet.ComStmtExecute, pnet.ComStmtSendLongData, pnet.ComStmtReset, pnet.ComStmtFetch:
+		stmtID := binary.LittleEndian.Uint32(command.Payload[1:5])
+		text, _, _ := c.backendConn.GetPreparedStmt(stmtID)
+		if len(text) == 0 {
+			c.lg.Error("prepared stmt not found", zap.Uint32("stmt_id", stmtID), zap.Stringer("cmd_type", command.Type))
+			return false
+		}
+		return lex.IsReadOnly(text)
+	case pnet.ComCreateDB, pnet.ComDropDB, pnet.ComDelayedInsert:
+		return false
+	}
+	// Treat ComStmtPrepare and ComStmtClose as read-only to make prepared stmt IDs in capture and replay phases the same.
+	// The problem is that it still requires write privilege. Better solutions are much more complex:
+	// - Replace all prepared DML statements with `SELECT 1`, including ComStmtPrepare and `SET SESSION_STATES`.
+	// - Remove all prepared DML statements and map catpure prepared stmt ID to replay prepared stmt ID, including ComStmtPrepare and `SET SESSION_STATES`.
+	return true
+}
+
 func (c *conn) updateCmdForExecuteStmt(command *cmd.Command) bool {
+	// updated before
 	if command.PreparedStmt != "" {
 		return true
 	}
@@ -147,7 +170,9 @@ func (c *conn) updateCmdForExecuteStmt(command *cmd.Command) bool {
 		if command.Type == pnet.ComStmtExecute {
 			_, args, _, err := pnet.ParseExecuteStmtRequest(command.Payload, paramNum, paramTypes)
 			if err != nil {
-				c.lg.Error("parsing ComExecuteStmt request failed", zap.Uint32("stmt_id", stmtID), zap.Error(err))
+				// Failing to parse the request is not critical, so don't return false.
+				c.lg.Error("parsing ComExecuteStmt request failed", zap.Uint32("stmt_id", stmtID), zap.String("sql", text),
+					zap.Int("param_num", paramNum), zap.ByteString("param_types", paramTypes), zap.Error(err))
 			}
 			command.Params = args
 		}

pkg/sqlreplay/conn/conn_test.go

@@ -178,7 +178,7 @@ func TestSkipReadOnly(t *testing.T) {
 		},
 		{
 			cmd:      &cmd.Command{Type: pnet.ComStmtPrepare, Payload: append([]byte{pnet.ComStmtPrepare.Byte()}, []byte("insert into t value(?)")...)},
-			readonly: false,
+			readonly: true,
 		},
 		{
 			cmd:      &cmd.Command{Type: pnet.ComStmtExecute, Payload: []byte{pnet.ComStmtExecute.Byte(), 2, 0, 0, 0}},
@@ -221,3 +221,71 @@ func TestSkipReadOnly(t *testing.T) {
 	cancel()
 	wg.Wait()
 }
+
+func TestReadOnly(t *testing.T) {
+	tests := []struct {
+		cmd      pnet.Command
+		stmt     string
+		readOnly bool
+	}{
+		{
+			cmd:      pnet.ComQuery,
+			stmt:     "select 1",
+			readOnly: true,
+		},
+		{
+			cmd:      pnet.ComQuery,
+			stmt:     "insert into t value(1)",
+			readOnly: false,
+		},
+		{
+			cmd:      pnet.ComStmtPrepare,
+			stmt:     "select ?",
+			readOnly: true,
+		},
+		{
+			cmd:      pnet.ComStmtPrepare,
+			stmt:     "insert into t value(?)",
+			readOnly: true,
+		},
+		{
+			cmd:      pnet.ComStmtExecute,
+			stmt:     "select ?",
+			readOnly: true,
+		},
+		{
+			cmd:      pnet.ComStmtExecute,
+			stmt:     "insert into t value(?)",
+			readOnly: false,
+		},
+		{
+			cmd:      pnet.ComStmtClose,
+			stmt:     "insert into t value(?)",
+			readOnly: true,
+		},
+		{
+			cmd:      pnet.ComQuit,
+			readOnly: true,
+		},
+		{
+			cmd:      pnet.ComCreateDB,
+			readOnly: false,
+		},
+	}
+
+	conn := &conn{}
+	backendConn := newMockBackendConn()
+	conn.backendConn = backendConn
+	for i, test := range tests {
+		var payload []byte
+		switch test.cmd {
+		case pnet.ComQuery:
+			payload = append([]byte{test.cmd.Byte()}, []byte(test.stmt)...)
+		default:
+			backendConn.prepared[1] = &preparedStmt{text: test.stmt}
+			payload = []byte{test.cmd.Byte(), 1, 0, 0, 0}
+		}
+		command := cmd.NewCommand(payload, time.Time{}, 100)
+		require.Equal(t, test.readOnly, conn.isReadOnly(command), "case %d", i)
+	}
+}

pkg/util/lex/filter.go

@@ -3,75 +3,66 @@
 
 package lex
 
+func startsWithKeyword(sql string, keywords [][]string) bool {
+	lexer := NewLexer(sql)
+	tokens := make([]string, 0, 2)
+	for _, kw := range keywords {
+		match := true
+		for i, t := range kw {
+			if len(tokens) <= i {
+				tokens = append(tokens, lexer.NextToken())
+			}
+			if tokens[i] != t {
+				match = false
+				break
+			}
+		}
+		if match {
+			return true
+		}
+	}
+	return false
+}
+
 var sensitiveKeywords = [][]string{
 	// contain passwords
-	{
-		"CREATE", "USER",
-	},
-	{
-		"ALTER", "USER",
-	},
-	{
-		"SET", "PASSWORD",
-	},
-	{
-		"GRANT",
-	},
+	{"CREATE", "USER"},
+	{"ALTER", "USER"},
+	{"SET", "PASSWORD"},
+	{"GRANT"},
 	// contain cloud storage url
-	{
-		"BACKUP",
-	},
-	{
-		"RESTORE",
-	},
-	{
-		"IMPORT",
-	},
+	{"BACKUP"},
+	{"RESTORE"},
+	{"IMPORT"},
 	// not supported yet
-	{
-		"LOAD", "DATA",
-	},
+	{"LOAD", "DATA"},
 }
 
-// ignore prepared statements in text-protocol
-// ignore EXPLAIN (including EXPLAIN ANALYZE) and TRACE
-// include SELECT FOR UPDATE because it releases locks immediately in auto-commit transactions
-// include SET because SET SESSION_STATES and SET session variables should be executed
-var readOnlyKeywords = []string{
-	"SELECT", "SHOW", "WITH", "SET", "USE", "DESC", "DESCRIBE", "TABLE", "DO",
+func IsSensitiveSQL(sql string) bool {
+	return startsWithKeyword(sql, sensitiveKeywords)
 }
 
-func IsSensitiveSQL(sql string) bool {
-	lexer := NewLexer(sql)
-	keyword := lexer.NextToken()
-	if len(keyword) == 0 {
-		return false
-	}
-	for _, kw := range sensitiveKeywords {
-		if keyword != kw[0] {
-			continue
-		}
-		if len(kw) <= 1 {
-			return true
-		}
-		keyword = lexer.NextToken()
-		if keyword == kw[1] {
-			return true
-		}
-	}
-	return false
+// ignore prepared statements in text-protocol
+// ignore EXPLAIN, EXPLAIN ANALYZE, and TRACE
+// include SELECT FOR UPDATE because it doesn't require write privilege
+// include SET because SET SESSION_STATES and SET session variables should be executed
+// include BEGIN / COMMIT in case the user sets autocommit to false, either in SET SESSION_STATES or SET @@autocommit
+var readOnlyKeywords = [][]string{
+	{"SELECT"},
+	{"SHOW"},
+	{"WITH"},
+	{"SET"},
+	{"USE"},
+	{"DESC"},
+	{"DESCRIBE"},
+	{"TABLE"},
+	{"DO"},
+	{"BEGIN"},
+	{"COMMIT"},
+	{"ROLLBACK"},
+	{"START", "TRANSACTION"},
 }
 
 func IsReadOnly(sql string) bool {
-	lexer := NewLexer(sql)
-	keyword := lexer.NextToken()
-	if len(keyword) == 0 {
-		return false
-	}
-	for _, kw := range readOnlyKeywords {
-		if keyword == kw {
-			return true
-		}
-	}
-	return false
+	return startsWithKeyword(sql, readOnlyKeywords)
 }

pkg/util/lex/filter_test.go

@@ -24,8 +24,8 @@ func TestSenstiveSQL(t *testing.T) {
 		{`set`, false},
 	}
 
-	for i, test := range tests {
-		require.Equal(t, test.sensitive, IsSensitiveSQL(test.sql), "case %d", i)
+	for _, test := range tests {
+		require.Equal(t, test.sensitive, IsSensitiveSQL(test.sql), test.sql)
 	}
 }
 
@@ -52,11 +52,11 @@ func TestReadOnlySQL(t *testing.T) {
 		{`do 1`, true},
 		{`/*hello */select 1`, true},
 		{`    select 1`, true},
-		{`/**/ start transaction`, false},
-		{`  COMMIT`, false},
+		{`/**/ start transaction`, true},
+		{`  COMMIT`, true},
 	}
 
-	for i, test := range tests {
-		require.Equal(t, test.readOnly, IsReadOnly(test.sql), "case %d", i)
+	for _, test := range tests {
+		require.Equal(t, test.readOnly, IsReadOnly(test.sql), test.sql)
 	}
 }