From cac2b72155cad81a578b97c3331c44f8bef9b679 Mon Sep 17 00:00:00 2001
From: rharding6373 <rharding6373@users.noreply.github.com>
Date: Wed, 12 Apr 2023 13:11:31 -0700
Subject: [PATCH] sqlsmith: call udfs in sqlsmith queries

This PR adds UDFs created by sqlsmith to the list of available functions
that can then be called in other queries generated by sqlsmith.

Epic: CRDB-20370
Informs: #90782

Release note: None
---
 pkg/internal/sqlsmith/relational.go | 54 ++++++++++++++++++++++++-----
 pkg/internal/sqlsmith/scalar.go     |  3 ++
 2 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/pkg/internal/sqlsmith/relational.go b/pkg/internal/sqlsmith/relational.go
index dbbb0f021653..d65fba06a45e 100644
--- a/pkg/internal/sqlsmith/relational.go
+++ b/pkg/internal/sqlsmith/relational.go
@@ -17,6 +17,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/randgen"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree/treecmp"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/volatility"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util"
 )
@@ -914,16 +915,22 @@ func (s *Smither) makeCreateFunc() (cf *tree.CreateFunction, ok bool) {
 	// parameters are supported.
 	// TODO(100962): Set a param default value sometimes.
 	params := make(tree.FuncParams, paramCnt)
+	paramTypes := make(tree.ParamTypes, paramCnt)
 	for i := 0; i < paramCnt; i++ {
 		// Do not allow collated string types. These are not supported in UDFs.
 		ptyp := s.randType()
 		for ptyp.Family() == types.CollatedStringFamily {
 			ptyp = s.randType()
 		}
+		pname := fmt.Sprintf("p%d", i)
 		params[i] = tree.FuncParam{
-			Name: tree.Name(fmt.Sprintf("p%d", i)),
+			Name: tree.Name(pname),
 			Type: ptyp,
 		}
+		paramTypes[i] = tree.ParamType{
+			Name: pname,
+			Typ:  ptyp,
+		}
 	}
 
 	// There are up to 5 function options that may be applied to UDFs.
@@ -949,15 +956,18 @@ func (s *Smither) makeCreateFunc() (cf *tree.CreateFunction, ok bool) {
 	// ~17%: FunctionVolatile
 	// ~17%: FunctionImmutable
 	// ~17%: FunctionStable
-	immutable := false
+	funcVol := tree.FunctionVolatile
+	vol := volatility.Volatile
 	switch s.d6() {
 	case 1:
-		opts = append(opts, tree.FunctionVolatile)
+		funcVol = tree.FunctionImmutable
+		vol = volatility.Immutable
 	case 2:
-		opts = append(opts, tree.FunctionImmutable)
-		immutable = true
-	case 3:
-		opts = append(opts, tree.FunctionStable)
+		funcVol = tree.FunctionStable
+		vol = volatility.Stable
+	}
+	if funcVol != tree.FunctionVolatile || s.coin() {
+		opts = append(opts, funcVol)
 	}
 
 	// FunctionLeakproof
@@ -965,10 +975,13 @@ func (s *Smither) makeCreateFunc() (cf *tree.CreateFunction, ok bool) {
 	// immutable, also specify leakproof 50% of the time. Otherwise, specify
 	// not leakproof 50% of the time (default is not leakproof).
 	leakproof := false
-	if immutable {
+	if funcVol == tree.FunctionImmutable {
 		leakproof = s.coin()
 	}
 	if leakproof || s.coin() {
+		if leakproof {
+			vol = volatility.Leakproof
+		}
 		opts = append(opts, tree.FunctionLeakproof(leakproof))
 	}
 
@@ -985,11 +998,14 @@ func (s *Smither) makeCreateFunc() (cf *tree.CreateFunction, ok bool) {
 	stmts := make([]string, 0, stmtCnt)
 	// Disable CTEs temporarily, since they are not currently supported in UDFs.
 	// TODO(92961): Allow CTEs in generated statements in UDF bodies.
+	// TODO(93049): Allow UDFs to call other UDFs, as well as create other UDFs.
 	oldDisableWith := s.disableWith
 	defer func() {
 		s.disableWith = oldDisableWith
+		s.disableUDFs = false
 	}()
 	s.disableWith = true
+	s.disableUDFs = true
 	for i := 0; i < stmtCnt; i++ {
 		// UDFs currently only support SELECT statements.
 		// TODO(87289): Add mutations to the generated statements.
@@ -1017,7 +1033,27 @@ func (s *Smither) makeCreateFunc() (cf *tree.CreateFunction, ok bool) {
 		Params:     params,
 		Options:    opts,
 	}
-	// TODO(harding): Register existing functions so we can refer to them in queries.
+
+	// Add this function to the functions list so that we can use it in future
+	// queries. Unfortunately, if the function fails to be created, then any
+	// queries that reference it will also fail.
+	class := tree.NormalClass
+	if setof {
+		class = tree.GeneratorClass
+	}
+
+	// We only add overload fields that are necessary to generate functions.
+	ov := &tree.Overload{
+		Volatility: vol,
+		Types:      paramTypes,
+		Class:      class,
+		IsUDF:      true,
+	}
+
+	functions[class][rtyp.Oid()] = append(functions[class][rtyp.Oid()], function{
+		def:      tree.NewFunctionDefinition(name.String(), &tree.FunctionProperties{}, nil /* def */),
+		overload: ov,
+	})
 	return stmt, true
 }
 
diff --git a/pkg/internal/sqlsmith/scalar.go b/pkg/internal/sqlsmith/scalar.go
index f5ef2646f235..39ba7a52c4ae 100644
--- a/pkg/internal/sqlsmith/scalar.go
+++ b/pkg/internal/sqlsmith/scalar.go
@@ -415,6 +415,9 @@ func makeFunc(s *Smither, ctx Context, typ *types.T, refs colRefs) (tree.TypedEx
 		return nil, false
 	}
 	fn := fns[s.rnd.Intn(len(fns))]
+	if s.disableUDFs && fn.overload.IsUDF {
+		return nil, false
+	}
 	if s.disableNondeterministicFns && fn.overload.Volatility > volatility.Immutable {
 		return nil, false
 	}