Merge pull request #24 from blastrain/feature/fix-is-already-commit-query-log

Fix IsAlreadyCommittedQueryLog

Author: kanataxa

database/sql/query_log.go

@@ -5,6 +5,7 @@ import (
 
 	vtparser "github.com/blastrain/vitess-sqlparser/sqlparser"
 	"github.com/pkg/errors"
+	"go.knocknote.io/octillery/debug"
 	"go.knocknote.io/octillery/exec"
 	"go.knocknote.io/octillery/sqlparser"
 )
@@ -248,7 +249,30 @@ func (t *Tx) mergeComparisonExprs(comparisonExprs []*vtparser.ComparisonExpr) vt
 
 func (t *Tx) createEqualComparisonExprWithArgs(left vtparser.Expr, right vtparser.Expr, args []interface{}) *vtparser.ComparisonExpr {
 	parser, _ := sqlparser.New()
-	val := right.(*vtparser.SQLVal)
+	switch val := right.(type) {
+	case *vtparser.SQLVal:
+		return &vtparser.ComparisonExpr{
+			Operator: vtparser.EqualStr,
+			Left:     left,
+			Right:    t.covertValArgToRealValue(parser, val, args),
+		}
+	case vtparser.ValTuple:
+		values := make(vtparser.ValTuple, len(val))
+		for idx, value := range val {
+			values[idx] = t.covertValArgToRealValue(parser, value.(*vtparser.SQLVal), args)
+		}
+		return &vtparser.ComparisonExpr{
+			Operator: vtparser.InStr,
+			Left:     left,
+			Right:    values,
+		}
+	default:
+		debug.Printf("[WARN] unexpected expr value %T", val)
+		return nil
+	}
+}
+
+func (t *Tx) covertValArgToRealValue(parser *sqlparser.Parser, val *vtparser.SQLVal, args []interface{}) *vtparser.SQLVal {
 	if val.Type == vtparser.ValArg {
 		arg := args[parser.ValueIndexByValArg(val)-1]
 		switch arg.(type) {
@@ -258,11 +282,7 @@ func (t *Tx) createEqualComparisonExprWithArgs(left vtparser.Expr, right vtparse
 			val = vtparser.NewStrVal([]byte(arg.(string)))
 		}
 	}
-	return &vtparser.ComparisonExpr{
-		Operator: vtparser.EqualStr,
-		Left:     left,
-		Right:    val,
-	}
+	return val
 }
 
 func (t *Tx) exprToComparisonExprs(expr vtparser.Expr, args []interface{}) []*vtparser.ComparisonExpr {

database/sql/sql_test.go

@@ -340,7 +340,9 @@ func testPrepareContextWithNotShardingTable(ctx context.Context, t *testing.T, d
 				power     int32
 				createdAt time.Time
 			)
-			stmt.QueryRow(1).Scan(&name, &age, &isGod, &point, &power, &createdAt)
+			if err := stmt.QueryRow(1).Scan(&name, &age, &isGod, &point, &power, &createdAt); err != nil {
+				t.Fatal(err)
+			}
 			if name != "alice" {
 				t.Fatal("cannot scan")
 			}
@@ -354,7 +356,9 @@ func testPrepareContextWithNotShardingTable(ctx context.Context, t *testing.T, d
 				power     int32
 				createdAt time.Time
 			)
-			stmt.QueryRowContext(ctx, 1).Scan(&name, &age, &isGod, &point, &power, &createdAt)
+			if err := stmt.QueryRowContext(ctx, 1).Scan(&name, &age, &isGod, &point, &power, &createdAt); err != nil {
+				t.Fatal(err)
+			}
 			if name != "alice" {
 				t.Fatal("cannot scan")
 			}

transaction_test.go

@@ -402,9 +402,22 @@ func TestIsAlreadyCommittedQueryLogErrorCase(t *testing.T) {
 }
 
 func TestIsAlreadyCommittedDeleteQueryLog(t *testing.T) {
-	testIsAlreadyCommittedQueryLog(t, &sql.QueryLog{
-		Query: "DELETE from user_stages WHERE id = ? AND user_id = ?",
-		Args:  []interface{}{1, 10},
+	t.Run("simple", func(t *testing.T) {
+		testIsAlreadyCommittedQueryLog(t, &sql.QueryLog{
+			Query: "DELETE from user_stages WHERE id = ? AND user_id = ?",
+			Args:  []interface{}{1, 10},
+		})
+	})
+
+	t.Run("included IN", func(t *testing.T) {
+		testIsAlreadyCommittedQueryLog(t, &sql.QueryLog{
+			Query: "DELETE from user_stages WHERE id IN (1,?,?) AND user_id = ?",
+			Args:  []interface{}{2, 3, 10},
+		})
+		testIsAlreadyCommittedQueryLog(t, &sql.QueryLog{
+			Query: "DELETE from user_items WHERE id IN (?,?) AND user_id = 10",
+			Args:  []interface{}{1, 2},
+		})
 	})
 }
 
@@ -421,4 +434,8 @@ func TestIsAlreadyCommittedUpdateQueryLog(t *testing.T) {
 		Query: "UPDATE user_stages set name = ?, age = 5 where user_id = ?",
 		Args:  []interface{}{"alice", 10},
 	})
+	testIsAlreadyCommittedQueryLog(t, &sql.QueryLog{
+		Query: "UPDATE user_stages set name = ?, age = 5 where user_id IN (?,?)",
+		Args:  []interface{}{"alice", 5, 10},
+	})
 }