Merge pull request #143638 from cockroachdb/blathers/backport-release-24.1-143632

release-24.1: sql: only re-use resolved routine type when flag is set

Author: DrewKimball

pkg/sql/logictest/testdata/logic_test/udf

@@ -974,6 +974,8 @@ LIMIT
 subtest end
 
 # Regression test for #104009.
+subtest regression_104009
+
 statement ok
 CREATE TABLE ab104009(a INT PRIMARY KEY, b INT)
 
@@ -997,3 +999,26 @@ PREPARE p AS SELECT f($1::REGCLASS::INT)
 
 statement ok
 EXECUTE p(10)
+
+# Regression test for #142615.
+subtest regression_142615
+
+statement ok
+create function app_to_db_id(app_id INT8) RETURNS INT8 LANGUAGE SQL AS $$ SELECT app_id * 2; $$;
+
+statement ok
+create sequence seq1;
+
+statement ok
+create table test (id int8 not null default app_to_db_id(nextval('seq1'::REGCLASS)));
+
+query TTITT rowsort
+select * from pg_catalog.pg_attrdef;
+----
+1508958170  132  2  unique_rowid()                                         unique_rowid()
+1202826234  151  1  gen_random_uuid()                                      gen_random_uuid()
+1202826232  151  3  now()                                                  now()
+2121222035  174  1  public.app_to_db_id(nextval('public.seq1'::REGCLASS))  public.app_to_db_id(nextval('public.seq1'::REGCLASS))
+2121222032  174  2  unique_rowid()                                         unique_rowid()
+
+subtest end

pkg/sql/opt/optbuilder/values.go

@@ -77,14 +77,15 @@ func (b *Builder) buildValuesClause(
 			elemPos += numCols
 			if texpr.ResolvedType().IsWildcardType() {
 				// Type-check the expression once again in order to update expressions
-				// that wrap a routine to reflect the modified type. Make sure to use
-				// the previously resolved type as the desired type, since the AST may
-				// have been modified to remove type annotations. This can happen when
-				// the routine's return type is unknown until its body is built.
-				texpr, err = tree.TypeCheck(b.ctx, texpr, b.semaCtx, texpr.ResolvedType())
-				if err != nil {
-					panic(err)
-				}
+				// that wrap a routine to reflect the modified type.
+				func() {
+					defer b.semaCtx.Properties.Restore(b.semaCtx.Properties)
+					b.semaCtx.Properties.RoutineUseResolvedType = true
+					texpr, err = tree.TypeCheck(b.ctx, texpr, b.semaCtx, desired)
+					if err != nil {
+						panic(err)
+					}
+				}()
 			}
 			if typ := texpr.ResolvedType(); typ.Family() != types.UnknownFamily {
 				if colTypes[colIdx].Family() == types.UnknownFamily {

pkg/sql/sem/tree/type_check.go

@@ -96,6 +96,12 @@ type SemaProperties struct {
 	// IgnoreUnpreferredOverloads is set to true when "unpreferred" overloads
 	// should not be used during type-checking and overload resolution.
 	IgnoreUnpreferredOverloads bool
+
+	// RoutineUseResolvedType is set to true when type-checking for a routine
+	// should reuse the already resolved type, if any. This is used to handle
+	// "re-type-checking" that occurs for a RECORD-returning routine, for which
+	// the return type is not known until the routine body is built.
+	RoutineUseResolvedType bool
 }
 
 type semaRequirements struct {
@@ -1154,9 +1160,10 @@ func (expr *FuncExpr) typeCheckWithFuncAncestor(semaCtx *SemaContext, fn func()
 func (expr *FuncExpr) TypeCheck(
 	ctx context.Context, semaCtx *SemaContext, desired *types.T,
 ) (TypedExpr, error) {
-	if expr.fn != nil && expr.fn.Type != BuiltinRoutine && expr.typ != nil {
-		// Don't overwrite the resolved properties for a user-defined routine if the
-		// routine has already been resolved.
+	if semaCtx != nil && semaCtx.Properties.RoutineUseResolvedType &&
+		expr.typ != nil && !expr.typ.IsWildcardType() {
+		// Don't overwrite the resolved properties for a routine if the routine has
+		// already been resolved (and we are in a context that needs this behavior).
 		return expr, nil
 	}