backend, lex: use a simpler lexer to parse begin statements (#758)

Author: djshow832

pkg/proxy/backend/cmd_processor_exec.go

@@ -5,12 +5,11 @@ package backend
 
 import (
 	"encoding/binary"
-	"strings"
 
 	"github.com/go-mysql-org/go-mysql/mysql"
-	"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"
 	"go.uber.org/zap"
 )
@@ -361,10 +360,5 @@ func (cp *CmdProcessor) needHoldRequest(request []byte) bool {
 		data = data[:len(data)-1]
 	}
 	query := hack.String(data)
-	return isBeginStmt(query)
-}
-
-func isBeginStmt(query string) bool {
-	normalized := parser.Normalize(query, "ON")
-	return strings.HasPrefix(normalized, "begin") || strings.HasPrefix(normalized, "start transaction")
+	return lex.IsStartTxn(query)
 }

pkg/proxy/backend/cmd_processor_test.go

@@ -918,61 +918,6 @@ func TestHoldRequest(t *testing.T) {
 	}
 }
 
-func TestBeginStmt(t *testing.T) {
-	tests := []struct {
-		stmt    string
-		isBegin bool
-	}{
-		{
-			stmt:    "begin",
-			isBegin: true,
-		},
-		{
-			stmt:    "BEGIN",
-			isBegin: true,
-		},
-		{
-			stmt:    "begin optimistic as of timestamp now()",
-			isBegin: true,
-		},
-		{
-			stmt:    "    begin",
-			isBegin: true,
-		},
-		{
-			stmt:    "start transaction",
-			isBegin: true,
-		},
-		{
-			stmt:    "START transaction",
-			isBegin: true,
-		},
-		{
-			stmt:    "start transaction with consistent snapshot",
-			isBegin: true,
-		},
-		{
-			stmt:    "begin; select 1",
-			isBegin: true,
-		},
-		{
-			stmt:    "/*+ some_hint */begin",
-			isBegin: true,
-		},
-		{
-			stmt:    "commit",
-			isBegin: false,
-		},
-		{
-			stmt:    "select 1; begin",
-			isBegin: false,
-		},
-	}
-	for _, test := range tests {
-		require.Equal(t, test.isBegin, isBeginStmt(test.stmt), test.stmt)
-	}
-}
-
 // Test forwarding multi-statements works well.
 func TestMultiStmt(t *testing.T) {
 	tc := newTCPConnSuite(t)

pkg/util/lex/filter.go

@@ -66,3 +66,12 @@ var readOnlyKeywords = [][]string{
 func IsReadOnly(sql string) bool {
 	return startsWithKeyword(sql, readOnlyKeywords)
 }
+
+var startTxnKeywords = [][]string{
+	{"START", "TRANSACTION"},
+	{"BEGIN"},
+}
+
+func IsStartTxn(sql string) bool {
+	return startsWithKeyword(sql, startTxnKeywords)
+}

pkg/util/lex/filter_test.go

@@ -60,3 +60,58 @@ func TestReadOnlySQL(t *testing.T) {
 		require.Equal(t, test.readOnly, IsReadOnly(test.sql), test.sql)
 	}
 }
+
+func TestStartTxn(t *testing.T) {
+	tests := []struct {
+		stmt    string
+		isBegin bool
+	}{
+		{
+			stmt:    "begin",
+			isBegin: true,
+		},
+		{
+			stmt:    "BEGIN",
+			isBegin: true,
+		},
+		{
+			stmt:    "begin optimistic as of timestamp now()",
+			isBegin: true,
+		},
+		{
+			stmt:    "    begin",
+			isBegin: true,
+		},
+		{
+			stmt:    "start transaction",
+			isBegin: true,
+		},
+		{
+			stmt:    "START transaction",
+			isBegin: true,
+		},
+		{
+			stmt:    "start transaction with consistent snapshot",
+			isBegin: true,
+		},
+		{
+			stmt:    "begin; select 1",
+			isBegin: true,
+		},
+		{
+			stmt:    "/*+ some_hint */begin",
+			isBegin: true,
+		},
+		{
+			stmt:    "commit",
+			isBegin: false,
+		},
+		{
+			stmt:    "select 1; begin",
+			isBegin: false,
+		},
+	}
+	for _, test := range tests {
+		require.Equal(t, test.isBegin, IsStartTxn(test.stmt), test.stmt)
+	}
+}