sqlsmith, tests, tree: add prepared queries to tlp

Previously, we did not have statements with placeholders in TLP queries.
This could have caught a correctness bug (#71002).
In this PR, support for prepared queries is added.

Fixes #71216.

Release note: None

pkg/cmd/roachtest/tests/tlp.go

@@ -155,12 +155,11 @@ func runMutationStatement(conn *gosql.DB, smither *sqlsmith.Smither, logStmt fun
 	})
 }
 
-// runTLPQuery runs two queries to perform TLP. If the results of the query are
-// not equal, an error is returned. Currently GenerateTLP always returns
-// unpartitioned and partitioned queries of the form "SELECT count(*) ...". The
-// resulting counts of the queries are compared in order to verify logical
-// correctness. See GenerateTLP for more information on TLP and the generated
-// queries.
+// runTLPQuery runs two queries to perform TLP. One is partitioned and one is
+// unpartitioned. If the results of the queries are not equal, an error is
+// returned. GenerateTLP also returns any placeholder arguments needed for the
+// partitioned query. See GenerateTLP for more information on TLP and the
+// generated queries.
 func runTLPQuery(conn *gosql.DB, smither *sqlsmith.Smither, logStmt func(string)) error {
 	// Ignore panics from GenerateTLP.
 	defer func() {
@@ -169,7 +168,7 @@ func runTLPQuery(conn *gosql.DB, smither *sqlsmith.Smither, logStmt func(string)
 		}
 	}()
 
-	unpartitioned, partitioned := smither.GenerateTLP()
+	unpartitioned, partitioned, args := smither.GenerateTLP()
 
 	return runWithTimeout(func() error {
 		rows1, err := conn.Query(unpartitioned)
@@ -185,7 +184,7 @@ func runTLPQuery(conn *gosql.DB, smither *sqlsmith.Smither, logStmt func(string)
 			//nolint:returnerrcheck
 			return nil
 		}
-		rows2, err := conn.Query(partitioned)
+		rows2, err := conn.Query(partitioned, args...)
 		if err != nil {
 			// Ignore errors.
 			//nolint:returnerrcheck
@@ -203,8 +202,8 @@ func runTLPQuery(conn *gosql.DB, smither *sqlsmith.Smither, logStmt func(string)
 			logStmt(unpartitioned)
 			logStmt(partitioned)
 			return errors.Newf(
-				"expected unpartitioned results to equal partitioned results\n%s\nsql: %s\n%s",
-				diff, unpartitioned, partitioned)
+				"expected unpartitioned and partitioned results to be equal\n%s\nsql: %s\n%s\nwith args: %s",
+				diff, unpartitioned, partitioned, args)
 		}
 		return nil
 	})

pkg/internal/sqlsmith/scalar.go

@@ -11,6 +11,8 @@
 package sqlsmith
 
 import (
+	"strconv"
+
 	"github.com/cockroachdb/cockroach/pkg/sql/randgen"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
@@ -73,6 +75,15 @@ func makeBoolExpr(s *Smither, refs colRefs) tree.TypedExpr {
 	return makeBoolExprContext(s, emptyCtx, refs)
 }
 
+func makeBoolExprWithPlaceholders(s *Smither, refs colRefs) (tree.Expr, []interface{}) {
+	expr := makeBoolExprContext(s, emptyCtx, refs)
+
+	// Replace constants with placeholders if the type is numeric or bool.
+	visitor := replaceDatumPlaceholderVisitor{}
+	exprFmt := expr.Walk(&visitor)
+	return exprFmt, visitor.Args
+}
+
 func makeBoolExprContext(s *Smither, ctx Context, refs colRefs) tree.TypedExpr {
 	return makeScalarSample(s.boolExprSampler, s, ctx, types.Bool, refs)
 }
@@ -693,3 +704,31 @@ func makeScalarSubquery(s *Smither, typ *types.T, refs colRefs) (tree.TypedExpr,
 
 	return subq, true
 }
+
+// replaceDatumPlaceholderVisitor replaces occurrences of numeric and bool Datum
+// expressions with placeholders, and updates Args with the corresponding Datum
+// values. This is used to prepare and execute a statement with placeholders.
+type replaceDatumPlaceholderVisitor struct {
+	Args []interface{}
+}
+
+var _ tree.Visitor = &replaceDatumPlaceholderVisitor{}
+
+// VisitPre satisfies the tree.Visitor interface.
+func (v *replaceDatumPlaceholderVisitor) VisitPre(
+	expr tree.Expr,
+) (recurse bool, newExpr tree.Expr) {
+	switch t := expr.(type) {
+	case tree.Datum:
+		if t.ResolvedType().IsNumeric() || t.ResolvedType() == types.Bool {
+			v.Args = append(v.Args, expr)
+			placeholder, _ := tree.NewPlaceholder(strconv.Itoa(len(v.Args)))
+			return false, placeholder
+		}
+		return false, expr
+	}
+	return true, expr
+}
+
+// VisitPost satisfies the Visitor interface.
+func (*replaceDatumPlaceholderVisitor) VisitPost(expr tree.Expr) tree.Expr { return expr }

pkg/internal/sqlsmith/tlp.go

@@ -12,19 +12,19 @@ package sqlsmith
 
 import (
 	"fmt"
-	"math/rand"
 	"strings"
 
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/errors"
 )
 
 // GenerateTLP returns two SQL queries as strings that can be used for Ternary
-// Logic Partitioning (TLP). TLP is a method for logically testing DBMSs which
-// is based on the logical guarantee that for a given predicate p, all rows must
-// satisfy exactly one of the following three predicates: p, NOT p, p IS NULL.
-// TLP can find bugs when an unpartitioned query and a query partitioned into
-// three sub-queries do not yield the same results.
+// Logic Partitioning (TLP). It also returns any placeholder arguments necessary
+// for the second partitioned query. TLP is a method for logically testing DBMSs
+// which is based on the logical guarantee that for a given predicate p, all
+// rows must satisfy exactly one of the following three predicates: p, NOT p, p
+// IS NULL. TLP can find bugs when an unpartitioned query and a query
+// partitioned into three sub-queries do not yield the same results.
 //
 // More information on TLP: https://www.manuelrigger.at/preprints/TLP.pdf.
 //
@@ -33,29 +33,32 @@ import (
 // possible to use TLP to test other aggregations, GROUP BY, and HAVING, which
 // have all been implemented in SQLancer. See:
 // https://github.com/sqlancer/sqlancer/tree/1.1.0/src/sqlancer/cockroachdb/oracle/tlp.
-func (s *Smither) GenerateTLP() (unpartitioned, partitioned string) {
+func (s *Smither) GenerateTLP() (unpartitioned, partitioned string, args []interface{}) {
 	// Set disableImpureFns to true so that generated predicates are immutable.
 	originalDisableImpureFns := s.disableImpureFns
 	s.disableImpureFns = true
 	defer func() {
 		s.disableImpureFns = originalDisableImpureFns
 	}()
 
-	switch tlpType := rand.Intn(4); tlpType {
+	switch tlpType := s.rnd.Intn(4); tlpType {
 	case 0:
 		return s.generateWhereTLP()
 	case 1:
-		return s.generateOuterJoinTLP()
+		partitioned, unpartitioned = s.generateOuterJoinTLP()
 	case 2:
-		return s.generateInnerJoinTLP()
+		partitioned, unpartitioned = s.generateInnerJoinTLP()
 	default:
-		return s.generateAggregationTLP()
+		partitioned, unpartitioned = s.generateAggregationTLP()
 	}
+	return partitioned, unpartitioned, nil
 }
 
 // generateWhereTLP returns two SQL queries as strings that can be used by the
 // GenerateTLP function. These queries make use of the WHERE clause to partition
-// the original query into three.
+// the original query into three. This function also returns a list of arguments
+// if the predicate contains placeholders. These arguments can be read
+// sequentially to match the placeholders.
 //
 // The first query returned is an unpartitioned query of the form:
 //
@@ -71,7 +74,7 @@ func (s *Smither) GenerateTLP() (unpartitioned, partitioned string) {
 //
 // If the resulting values of the two queries are not equal, there is a logical
 // bug.
-func (s *Smither) generateWhereTLP() (unpartitioned, partitioned string) {
+func (s *Smither) generateWhereTLP() (unpartitioned, partitioned string, args []interface{}) {
 	f := tree.NewFmtCtx(tree.FmtParsable)
 
 	table, _, _, cols, ok := s.getSchemaTable()
@@ -83,7 +86,12 @@ func (s *Smither) generateWhereTLP() (unpartitioned, partitioned string) {
 
 	unpartitioned = fmt.Sprintf("SELECT * FROM %s", tableName)
 
-	pred := makeBoolExpr(s, cols)
+	var pred tree.Expr
+	if s.coin() {
+		pred = makeBoolExpr(s, cols)
+	} else {
+		pred, args = makeBoolExprWithPlaceholders(s, cols)
+	}
 	pred.Format(f)
 	predicate := f.CloseAndGetString()
 
@@ -96,7 +104,7 @@ func (s *Smither) generateWhereTLP() (unpartitioned, partitioned string) {
 		part1, part2, part3,
 	)
 
-	return unpartitioned, partitioned
+	return unpartitioned, partitioned, args
 }
 
 // generateOuterJoinTLP returns two SQL queries as strings that can be used by the
@@ -304,7 +312,7 @@ func (s *Smither) generateAggregationTLP() (unpartitioned, partitioned string) {
 	tableNameAlias := strings.TrimSpace(strings.Split(tableName, "AS")[1])
 
 	var innerAgg, outerAgg string
-	switch aggType := rand.Intn(3); aggType {
+	switch aggType := s.rnd.Intn(3); aggType {
 	case 0:
 		innerAgg, outerAgg = "MAX", "MAX"
 	case 1:

pkg/sql/sem/tree/expr.go

@@ -57,7 +57,7 @@ type TypedExpr interface {
 	// encountered: Placeholder, VarName (and related UnqualifiedStar,
 	// UnresolvedName and AllColumnsSelector) or Subquery. These nodes
 	// should be replaced prior to expression evaluation by an
-	// appropriate WalkExpr. For example, Placeholder should be replace
+	// appropriate WalkExpr. For example, Placeholder should be replaced
 	// by the argument passed from the client.
 	Eval(*EvalContext) (Datum, error)
 	// ResolvedType provides the type of the TypedExpr, which is the type of Datum