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 (cockroachdb#71002).
In this PR, support for prepared queries is added.

Fixes cockroachdb#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. Generate TLP 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

@@ -73,6 +73,15 @@ func makeBoolExpr(s *Smither, refs colRefs) tree.TypedExpr {
 	return makeBoolExprContext(s, emptyCtx, refs)
 }
 
+func makeBoolExprPlaceholder(s *Smither, refs colRefs) (tree.Expr, []interface{}) {
+	expr := makeBoolExprContext(s, emptyCtx, refs)
+
+	// Replace constants with placeholders if the type is numeric or bool.
+	visitor := tree.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)
 }

pkg/internal/sqlsmith/tlp.go

@@ -33,24 +33,52 @@ 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 := rand.Intn(5); tlpType {
 	case 0:
-		return s.generateWhereTLP()
+		partitioned, unpartitioned = s.generateWhereTLP()
 	case 1:
-		return s.generateOuterJoinTLP()
+		partitioned, unpartitioned = s.generateOuterJoinTLP()
 	case 2:
-		return s.generateInnerJoinTLP()
+		partitioned, unpartitioned = s.generateInnerJoinTLP()
+	case 3:
+		partitioned, unpartitioned = s.generateAggregationTLP()
 	default:
-		return s.generateAggregationTLP()
+		return s.generatePrepareTLP()
 	}
+	return partitioned, unpartitioned, nil
+}
+
+func (s *Smither) generatePrepareTLP() (unpartitioned, partitioned string, args []interface{}) {
+	f := tree.NewFmtCtx(tree.FmtParsable)
+
+	table, _, _, cols, ok := s.getSchemaTable()
+	if !ok {
+		panic(errors.AssertionFailedf("failed to find random table"))
+	}
+	table.Format(f)
+	tableName := f.CloseAndGetString()
+
+	unpartitioned = fmt.Sprintf("SELECT * FROM %s", tableName)
+
+	pred, args := makeBoolExprPlaceholder(s, cols)
+	pred.Format(f)
+	predicate := f.CloseAndGetString()
+
+	part1 := fmt.Sprintf("SELECT * FROM %s WHERE %s", tableName, predicate)
+	part2 := fmt.Sprintf("SELECT * FROM %s WHERE NOT (%s)", tableName, predicate)
+	part3 := fmt.Sprintf("SELECT * FROM %s WHERE (%s) IS NULL", tableName, predicate)
+
+	partitioned = fmt.Sprintf("(%s) UNION ALL (%s) UNION ALL (%s)", part1, part2, part3)
+
+	return unpartitioned, partitioned, args
 }
 
 // generateWhereTLP returns two SQL queries as strings that can be used by the

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

pkg/sql/sem/tree/placeholders.go

@@ -14,6 +14,7 @@ import (
 	"bytes"
 	"fmt"
 	"math"
+	"strconv"
 
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
@@ -210,3 +211,29 @@ func (p *PlaceholderInfo) IsUnresolvedPlaceholder(expr Expr) bool {
 	}
 	return false
 }
+
+// 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 _ Visitor = &ReplaceDatumPlaceholderVisitor{}
+
+// VisitPre satisfies the tree.Visitor interface.
+func (v *ReplaceDatumPlaceholderVisitor) VisitPre(expr Expr) (recurse bool, newExpr Expr) {
+	switch t := expr.(type) {
+	case Datum:
+		if t.ResolvedType().IsNumeric() || t.ResolvedType() == types.Bool {
+			v.Args = append(v.Args, expr)
+			placeholder, _ := NewPlaceholder(strconv.Itoa(len(v.Args)))
+			return false, placeholder
+		}
+		return false, expr
+	}
+	return true, expr
+}
+
+// VisitPost satisfies the Visitor interface.
+func (*ReplaceDatumPlaceholderVisitor) VisitPost(expr Expr) Expr { return expr }