feat: clean up signatures

Signed-off-by: Andres Taylor <[email protected]>

Author: systay

go/testdata/small-slow-query-transactions.json

@@ -1,15 +1,15 @@
 [
   {
-    "Queries": [
-      "update tblA set apa = :v where foo = :0 and id = :1",
-      "update tblB set monkey = :v where bar = :0 and id = :2"
+    "query-signatures": [
+      "update tblA where foo = :0 and id = :1 set apa",
+      "update tblB where bar = :0 and id = :2 set monkey"
     ],
-    "Predicates": [
+    "predicates": [
       "tblA.foo = 0",
-      "tblA.id = 1",
+      "tblA.id = ?",
       "tblB.bar = 0",
-      "tblB.id = 2"
+      "tblB.id = ?"
     ],
-    "Count": 2
+    "count": 2
   }
 ]

go/transactions/transaction_signature.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"hash/fnv"
 	"sort"
+	"strconv"
 
 	"vitess.io/vitess/go/vt/sqlparser"
 )
@@ -46,7 +47,11 @@ type (
 )
 
 func (pi predicateInfo) String() string {
-	return fmt.Sprintf("%s.%s %s %d", pi.Table, pi.Col, pi.Op.ToString(), pi.Val)
+	val := strconv.Itoa(pi.Val)
+	if pi.Val == -1 {
+		val = "?"
+	}
+	return fmt.Sprintf("%s.%s %s %s", pi.Table, pi.Col, pi.Op.ToString(), val)
 }
 
 func (pi predicateInfo) compareTo(b predicateInfo) int {
@@ -69,9 +74,9 @@ func (tx *TxSignature) MarshalJSON() ([]byte, error) {
 	}
 
 	return json.Marshal(struct {
-		Queries    []string
-		Predicates []string
-		Count      int
+		Queries    []string `json:"query-signatures"`
+		Predicates []string `json:"predicates"`
+		Count      int      `json:"count"`
 	}{
 		Queries:    tx.Queries,
 		Predicates: predicateStrings,
@@ -164,13 +169,49 @@ func (tx *TxSignature) Equals(other *TxSignature) bool {
 	return true
 }
 
+// CleanUp removes values that are only used once and replaces them with -1
+func (tx *TxSignature) CleanUp() *TxSignature {
+	newPredicates := make([]predicateInfo, 0, len(tx.Predicates))
+	usedValues := make(map[int]int)
+
+	// First let's count how many times each value is used
+	for _, pred := range tx.Predicates {
+		usedValues[pred.Val]++
+	}
+
+	// Now we replace values only used once with -1
+	newCount := 0
+	newValues := make(map[int]int)
+	for _, pred := range tx.Predicates {
+		if usedValues[pred.Val] == 1 {
+			pred.Val = -1
+		} else {
+			newVal, found := newValues[pred.Val]
+			if !found {
+				// Assign a new value to this predicate
+				newVal = newCount
+				newCount++
+				newValues[pred.Val] = newVal
+			}
+			pred.Val = newVal
+		}
+		newPredicates = append(newPredicates, pred)
+	}
+
+	return &TxSignature{
+		Queries:    tx.Queries,
+		Predicates: newPredicates,
+		Count:      tx.Count,
+	}
+}
+
 func (m *txSignatureMap) MarshalJSON() ([]byte, error) {
 	// Collect all interesting TxSignatures into a slice
 	var signatures []*TxSignature
 	for _, bucket := range m.data {
 		for _, txSig := range bucket {
 			if txSig.Count > 1 {
-				signatures = append(signatures, txSig)
+				signatures = append(signatures, txSig.CleanUp())
 			}
 		}
 	}

go/transactions/transactions.go

@@ -223,21 +223,39 @@ func (s *state) consume(ch <-chan []sqlparser.Statement, wg *sync.WaitGroup) {
 	}
 }
 
-func (s *state) consumeUpdate(query *sqlparser.Update, st *semantics.SemTable, n *normalizer, tx *TxSignature) {
-	defer func() {
-		tx.Queries = append(tx.Queries, sqlparser.String(query))
-	}()
+func querySignatureForUpd(query *sqlparser.Update) string {
+	buffer := sqlparser.NewTrackedBuffer(nil)
+	builder := buffer.Builder
+	builder.WriteString("update ")
+
+	for i, tbl := range query.TableExprs {
+		tbl.Format(buffer)
+		if i < len(query.TableExprs)-1 {
+			buffer.WriteString(", ")
+		}
+	}
+
+	if query.Where != nil {
+		query.Where.Format(buffer)
+	}
+
+	builder.WriteString(" set ")
 
-	// Normalize the AST our own way:
-	// 	- Replace the value in SET by "v"
-	// 	- Replace the literals found in where clause comparisons by the corresponding ID we got earlier
 	for i, expr := range query.Exprs {
-		query.Exprs[i] = &sqlparser.UpdateExpr{
-			Name: expr.Name,
-			Expr: sqlparser.NewArgument("v"),
+		expr.Name.Format(buffer)
+		if i < len(query.Exprs)-1 {
+			buffer.WriteString(", ")
 		}
 	}
 
+	return builder.String()
+}
+
+func (s *state) consumeUpdate(query *sqlparser.Update, st *semantics.SemTable, n *normalizer, tx *TxSignature) {
+	defer func() {
+		tx.Queries = append(tx.Queries, querySignatureForUpd(query))
+	}()
+
 	if query.Where == nil {
 		return
 	}