fix ut

Author: volgariver6

pkg/vm/engine/disttae/txn_test.go

@@ -16,11 +16,15 @@ package disttae
 
 import (
 	"bytes"
+	"context"
 	"sync"
 	"testing"
 
 	"github.com/matrixorigin/matrixone/pkg/container/types"
 	"github.com/matrixorigin/matrixone/pkg/objectio"
+	"github.com/matrixorigin/matrixone/pkg/sql/parsers"
+	"github.com/matrixorigin/matrixone/pkg/sql/parsers/dialect"
+	"github.com/matrixorigin/matrixone/pkg/txn/client"
 	"github.com/stretchr/testify/require"
 )
 
@@ -62,3 +66,16 @@ func Test_GetUncommittedS3Tombstone(t *testing.T) {
 		return true
 	})
 }
+
+func TestTransaction_StartStatement(t *testing.T) {
+	var txn Transaction
+	var fn func()
+	ctx := context.Background()
+	txn.op, fn = client.NewTestTxnOperator(ctx)
+	defer fn()
+	txn.StartStatement(nil)
+	stmts, err := parsers.Parse(ctx, dialect.MYSQL, "show databases", 0)
+	require.NoError(t, err)
+	require.Equal(t, len(stmts), 1)
+	txn.StartStatement(stmts[0])
+}

pkg/vm/engine/disttae/types.go

@@ -18,8 +18,8 @@ import (
 	"bytes"
 	"context"
 	"encoding/hex"
+	"flag"
 	"fmt"
-	"github.com/matrixorigin/matrixone/pkg/sql/parsers/tree"
 	"math"
 	"strconv"
 	"sync"
@@ -45,6 +45,8 @@ import (
 	"github.com/matrixorigin/matrixone/pkg/pb/timestamp"
 	qclient "github.com/matrixorigin/matrixone/pkg/queryservice/client"
 	"github.com/matrixorigin/matrixone/pkg/sql/colexec"
+	"github.com/matrixorigin/matrixone/pkg/sql/parsers/dialect"
+	"github.com/matrixorigin/matrixone/pkg/sql/parsers/tree"
 	"github.com/matrixorigin/matrixone/pkg/txn/client"
 	"github.com/matrixorigin/matrixone/pkg/txn/trace"
 	"github.com/matrixorigin/matrixone/pkg/udf"
@@ -530,9 +532,15 @@ func (txn *Transaction) StartStatement(stmt tree.Statement) {
 	if txn.startStatementCalled {
 		var sql string
 		if stmt != nil {
-			sql = stmt.String()
+			fmtCtx := tree.NewFmtCtx(dialect.MYSQL, tree.WithQuoteString(true))
+			stmt.Format(fmtCtx)
+			sql = fmtCtx.String()
+		}
+		log := logutil.Fatal
+		if flag.Lookup("test.v") != nil {
+			log = logutil.Error
 		}
-		logutil.Fatal("BUG: StartStatement called twice",
+		log("BUG: StartStatement called twice",
 			zap.String("txn", hex.EncodeToString(txn.op.Txn().ID)),
 			zap.String("SQL", sql),
 		)