sql/opt: build CTEs for SQL expressions embedded in PL/pgSQL

This commit fixes an oversight from when we introduced support for
CTEs within routies. Namely, SQL expressions can contain subqueries,
which in turn can contain CTEs. The logic for building a SQL *statement*
within a PL/pgSQL routine correctly handles CTEs by building them into
the constructed `RelExpr`. The logic for SQL *expressions* was previously
missing the same.

To fix this issue, `plpgsqlBuilder.buildSQLExpr` now checks if the built
scalar expression contained any CTEs. If it did, the scalar is wrapped
in the built CTEs, which are in turn wrapped in a subquery, which is
returned as the final scalar result. This prevents issues where the
CTEs are built at an outer scope, which can produce invalid plans.

Fixes #138273

Release note (bug fix): Fixed a bug existing only in pre-release versions
of v25.1. The bug could cause creation of a PL/pgSQL routine with a CTE
to fail with an error like the following: `unexpected root expression: with`.

Author: DrewKimball

pkg/sql/logictest/testdata/logic_test/udf_cte

@@ -121,3 +121,42 @@ query T
 SELECT f();
 ----
 {10,20,30,40}
+
+subtest regression_138273
+
+statement ok
+CREATE SEQUENCE seq_1;
+
+# Verify that the function can be successfully created with a CTE inside a
+# scalar expression.
+statement ok
+CREATE FUNCTION f138273_1() RETURNS TIMESTAMP LANGUAGE PLpgSQL AS $$
+  DECLARE
+    decl TIMESTAMP;
+  BEGIN
+    WHILE false LOOP
+      IF true THEN
+      ELSIF EXISTS (WITH cte(col) AS (SELECT * FROM (VALUES (currval('seq_1'))) AS foo) SELECT 1 FROM cte) THEN
+        RETURN decl;
+      END IF;
+    END LOOP;
+  END;
+$$;
+
+# Test a similar function that actually executes the CTE.
+statement ok
+CREATE FUNCTION f138273_2() RETURNS INT LANGUAGE PLpgSQL AS $$
+  BEGIN
+    SELECT nextval('seq_1');
+    WHILE EXISTS (WITH cte(col) AS (SELECT * FROM (VALUES (currval('seq_1'))) AS foo) SELECT 1 FROM cte) LOOP
+      RETURN currval('seq_1');
+    END LOOP;
+  END;
+$$;
+
+query I
+SELECT f138273_2();
+----
+1
+
+subtest end

pkg/sql/opt/optbuilder/create_trigger.go

@@ -179,6 +179,15 @@ func (b *Builder) buildFunctionForTrigger(
 
 	// The trigger function can reference the NEW and OLD transition relations,
 	// aliased in the trigger definition.
+	if len(ct.Transitions) > 0 {
+		// Save the existing CTEs in the builder and restore them after the function
+		// body is built.
+		prevCTEs := b.ctes
+		b.ctes = nil
+		defer func() {
+			b.ctes = prevCTEs
+		}()
+	}
 	for _, transition := range ct.Transitions {
 		// Build a fake relational expression with a column corresponding to each
 		// column from the table.
@@ -203,12 +212,6 @@ func (b *Builder) buildFunctionForTrigger(
 		funcScope.ctes[string(transition.Name)] = cte
 		b.addCTE(cte)
 	}
-	if len(ct.Transitions) > 0 {
-		defer func() {
-			// Reset the CTEs in the builder after the function body is built.
-			b.ctes = nil
-		}()
-	}
 
 	// The trigger function takes a set of implicitly-defined parameters, two of
 	// which are determined by the table's record type. Add them to the trigger

pkg/sql/opt/optbuilder/plpgsql.go

@@ -2201,13 +2201,34 @@ func (b *plpgsqlBuilder) buildSQLExpr(expr ast.Expr, typ *types.T, s *scope) opt
 		// For lazy SQL evaluation, replace all expressions with NULL.
 		return memo.NullSingleton
 	}
+	// Save any outer CTEs before building the expression, which may have
+	// subqueries with inner CTEs.
+	prevCTEs := b.ob.ctes
+	b.ob.ctes = nil
+	defer func() {
+		b.ob.ctes = prevCTEs
+	}()
 	expr, _ = tree.WalkExpr(s, expr)
 	typedExpr, err := expr.TypeCheck(b.ob.ctx, b.ob.semaCtx, typ)
 	if err != nil {
 		panic(err)
 	}
 	scalar := b.ob.buildScalar(typedExpr, s, nil, nil, b.colRefs)
-	return b.coerceType(scalar, typ)
+	scalar = b.coerceType(scalar, typ)
+	if len(b.ob.ctes) == 0 {
+		return scalar
+	}
+	// There was at least one CTE within the scalar expression. It is possible to
+	// "hoist" them above this point, but building them eagerly here means that
+	// callers don't have to worry about CTE handling.
+	f := b.ob.factory
+	valuesCol := f.Metadata().AddColumn("", scalar.DataType())
+	valuesExpr := f.ConstructValues(
+		memo.ScalarListExpr{f.ConstructTuple(memo.ScalarListExpr{scalar}, scalar.DataType())},
+		&memo.ValuesPrivate{Cols: opt.ColList{valuesCol}, ID: f.Metadata().NextUniqueID()},
+	)
+	withExpr := b.ob.buildWiths(valuesExpr, b.ob.ctes)
+	return f.ConstructSubquery(withExpr, &memo.SubqueryPrivate{})
 }
 
 // buildSQLStatement type-checks and builds the given SQL statement into a

pkg/sql/opt/optbuilder/testdata/create_function

@@ -138,3 +138,36 @@ build
 CREATE FUNCTION f() RETURNS UNKNOWN LANGUAGE SQL AS $$ SELECT NULL; $$;
 ----
 error (42P13): SQL functions cannot return type unknown
+
+# Regression test for #138273 - correctly build CTEs in SQL expressions
+# embedded in a PL/pgSQL routine.
+build
+CREATE FUNCTION f138273_1() RETURNS TIMESTAMP LANGUAGE PLpgSQL AS $$
+  DECLARE
+    decl TIMESTAMP;
+  BEGIN
+    WHILE false LOOP
+      IF true THEN
+      ELSIF EXISTS (WITH cte(col) AS (SELECT * FROM (VALUES (random())) AS foo) SELECT 1 FROM cte) THEN
+        RETURN decl;
+      END IF;
+    END LOOP;
+  END;
+$$;
+----
+create-function
+ ├── CREATE FUNCTION f138273_1()
+ │   	RETURNS TIMESTAMP
+ │   	LANGUAGE plpgsql
+ │   	AS $$DECLARE
+ │   decl TIMESTAMP;
+ │   BEGIN
+ │   WHILE false LOOP
+ │   IF true THEN
+ │   ELSIF EXISTS (WITH cte (col) AS (SELECT foo.column1 FROM (VALUES (random())) AS foo) SELECT 1 FROM cte) THEN
+ │   	RETURN decl;
+ │   END IF;
+ │   END LOOP;
+ │   END;
+ │   $$
+ └── no dependencies