sql/sem/tree: do not Restore SemaRejectFlags during semantic analysis

This commit fixes a bug introduced in cockroachdb#105582 that caused
SemaRejectFlags to be restored during semantic analysis, preventing the
analysis from detecting some forms of invalid expressions.

Fixes cockroachdb#108166

There is no release note because the related bug does not exist in any
releases.

Release note: None

Author: mgartner

pkg/sql/logictest/testdata/logic_test/delete

@@ -576,3 +576,15 @@ statement error pgcode 42803 sum\(\): aggregate functions are not allowed in ORD
 DELETE FROM t107634 ORDER BY sum(a) LIMIT 1;
 
 subtest end
+
+# Regression test for #108166. Do not allow aggregate functions in ORDER BY when
+# the function is wrapped by a conditional expression.
+subtest regression_108166
+
+statement ok
+CREATE TABLE t108166 (a INT)
+
+statement error pgcode 42803 sum\(\): aggregate functions are not allowed in ORDER BY in DELETE
+DELETE FROM t108166 ORDER BY COALESCE(sum(a), 1) LIMIT 1;
+
+subtest end

pkg/sql/logictest/testdata/logic_test/srfs

@@ -1322,16 +1322,16 @@ subtest generator-syntax
 
 # Regression test for #97119 and #94890 - return syntax error when CASE or
 # COALESCE is used with a set-generating function as argument.
-statement error pq: set-returning functions are not allowed in CASE
+statement error pq: set-returning functions are not allowed in conditional expressions
 SELECT CASE generate_series(1, 3) WHEN 3 THEN 0 ELSE 1 END;
 
-statement error pq: set-returning functions are not allowed in CASE
+statement error pq: set-returning functions are not allowed in conditional expressions
 SELECT CASE WHEN true THEN generate_series(1, 3) ELSE 1 END;
 
-statement error pq: set-returning functions are not allowed in CASE
+statement error pq: set-returning functions are not allowed in conditional expressions
 SELECT CASE WHEN false THEN 1 ELSE generate_series(1, 3) END;
 
-statement error pq: set-returning functions are not allowed in COALESCE
+statement error pq: set-returning functions are not allowed in conditional expressions
 SELECT COALESCE(generate_series(1, 10));
 
 # A subquery with a generator function is allowed within CASE and COALESCE.
@@ -1376,12 +1376,12 @@ SELECT COALESCE(sum(x) OVER ()) FROM xy;
 15
 
 # IF does not allow generator functions.
-statement error pq: set-returning functions are not allowed in IF
+statement error pq: set-returning functions are not allowed in conditional expressions
 SELECT IF(x > y, generate_series(1, 3), 0) FROM xy;
 
 # IFNULL does not allow generator functions. Note that the error mentions
 # COALESCE because IFNULL is parsed directly as a COALESCE expression.
-statement error pq: set-returning functions are not allowed in COALESCE
+statement error pq: set-returning functions are not allowed in conditional expressions
 SELECT IFNULL(1, generate_series(1, 2));
 
 # NULLIF allows generator functions.

pkg/sql/logictest/testdata/logic_test/update

@@ -672,3 +672,15 @@ statement error pgcode 42803 sum\(\): aggregate functions are not allowed in ORD
 UPDATE t107634 SET a = 1 ORDER BY sum(a) LIMIT 1;
 
 subtest end
+
+# Regression test for #108166. Do not allow aggregate functions in ORDER BY when
+# the function is wrapped by a conditional expression.
+subtest regression_108166
+
+statement ok
+CREATE TABLE t108166 (a INT)
+
+statement error pgcode 42803 sum\(\): aggregate functions are not allowed in ORDER BY in UPDATE
+UPDATE t108166 SET a = 1 ORDER BY COALESCE(sum(a), 1) LIMIT 1;
+
+subtest end

pkg/sql/opt/optbuilder/srfs.go

@@ -50,6 +50,9 @@ func (s *srf) TypeCheck(
 		// invalid usage here.
 		return nil, tree.NewInvalidFunctionUsageError(tree.GeneratorClass, ctx.TypeCheckContext())
 	}
+	if ctx.Properties.Ancestors.Has(tree.ConditionalAncestor) {
+		return nil, tree.NewInvalidFunctionUsageError(tree.GeneratorClass, "conditional expressions")
+	}
 	if ctx.Properties.Derived.SeenGenerator {
 		// This error happens if this srf struct is nested inside a raw srf that
 		// has not yet been replaced. This is possible since scope.replaceSRF first

pkg/sql/sem/tree/type_check.go

@@ -209,6 +209,11 @@ const (
 	// checking the parameters of a window function in order to reject nested
 	// window functions.
 	WindowFuncAncestor
+
+	// ConditionalAncestor is temporarily added to ScalarAncestors while type
+	// checking condition expressions CASE, COALESCE, and IF. Used to reject
+	// set-returning functions within conditional expressions.
+	ConditionalAncestor
 )
 
 // Push adds the given ancestor to s.
@@ -424,11 +429,8 @@ func (expr *CaseExpr) TypeCheck(
 	ctx context.Context, semaCtx *SemaContext, desired *types.T,
 ) (TypedExpr, error) {
 	if semaCtx != nil {
-		// We need to save and restore the previous value of the field in
-		// semaCtx in case we are recursively called within a subquery
-		// context.
-		defer semaCtx.Properties.Restore(semaCtx.Properties)
-		semaCtx.Properties.Require("CASE", RejectGenerators)
+		defer semaCtx.Properties.Ancestors.PopTo(semaCtx.Properties.Ancestors)
+		semaCtx.Properties.Ancestors.Push(ConditionalAncestor)
 	}
 	var err error
 	tmpExprs := make([]Expr, 0, len(expr.Whens)+1)
@@ -877,11 +879,8 @@ func (expr *CoalesceExpr) TypeCheck(
 	ctx context.Context, semaCtx *SemaContext, desired *types.T,
 ) (TypedExpr, error) {
 	if semaCtx != nil {
-		// We need to save and restore the previous value of the field in
-		// semaCtx in case we are recursively called within a subquery
-		// context.
-		defer semaCtx.Properties.Restore(semaCtx.Properties)
-		semaCtx.Properties.Require("COALESCE", RejectGenerators)
+		defer semaCtx.Properties.Ancestors.PopTo(semaCtx.Properties.Ancestors)
+		semaCtx.Properties.Ancestors.Push(ConditionalAncestor)
 	}
 	typedSubExprs, retType, err := typeCheckSameTypedExprs(ctx, semaCtx, desired, expr.Exprs...)
 	if err != nil {
@@ -1029,6 +1028,9 @@ func (sc *SemaContext) checkFunctionUsage(expr *FuncExpr, def *ResolvedFunctionD
 		if sc.Properties.IsSet(RejectGenerators) {
 			return NewInvalidFunctionUsageError(GeneratorClass, sc.Properties.required.context)
 		}
+		if sc.Properties.Ancestors.Has(ConditionalAncestor) {
+			return NewInvalidFunctionUsageError(GeneratorClass, "conditional expressions")
+		}
 		sc.Properties.Derived.SeenGenerator = true
 	}
 	return nil
@@ -1355,11 +1357,8 @@ func (expr *IfExpr) TypeCheck(
 	ctx context.Context, semaCtx *SemaContext, desired *types.T,
 ) (TypedExpr, error) {
 	if semaCtx != nil {
-		// We need to save and restore the previous value of the field in
-		// semaCtx in case we are recursively called within a subquery
-		// context.
-		defer semaCtx.Properties.Restore(semaCtx.Properties)
-		semaCtx.Properties.Require("IF", RejectGenerators)
+		defer semaCtx.Properties.Ancestors.PopTo(semaCtx.Properties.Ancestors)
+		semaCtx.Properties.Ancestors.Push(ConditionalAncestor)
 	}
 	typedCond, err := typeCheckAndRequireBoolean(ctx, semaCtx, expr.Cond, "IF condition")
 	if err != nil {