From 332e8d228fa0fd5012156d4776111831333c14c6 Mon Sep 17 00:00:00 2001
From: Drew Kimball <drewk@cockroachlabs.com>
Date: Tue, 27 Aug 2024 04:00:57 +0000
Subject: [PATCH] opt: type check VALUES rows after building them

The type of a RECORD-returning UDF is not fully resolved until after it
is fully built within the optbuilder. Therefore, a VALUES expression with
a UDF can only determine its column types after building its rows. We
accounted for this in #103078, but that fix was incomplete - building the row
resolves the type of the UDF, but does not update the type annotation for
any expressions that wrap it (for example, a COALESCE).

This commit fully fixes the type resolution bug by type checking the row
once again after it is built. This ensures that the UDF's resolved type is
correctly propagated to parent expressions.

Fixes #117101

Release note (bug fix): Fixed a bug that would cause an internal error
when the result of a RECORD-returning UDF was wrapped by another expression
(such as COALESCE) within a VALUES clause.
---
 pkg/ccl/testccl/sqlccl/explain_test.go    |  3 +--
 pkg/sql/logictest/testdata/logic_test/udf | 22 ++++++++++++++++++++++
 pkg/sql/opt/optbuilder/values.go          |  8 ++++++++
 pkg/sql/sem/tree/type_check.go            |  6 ++++++
 4 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/pkg/ccl/testccl/sqlccl/explain_test.go b/pkg/ccl/testccl/sqlccl/explain_test.go
index 4426df71336f..4549395ff631 100644
--- a/pkg/ccl/testccl/sqlccl/explain_test.go
+++ b/pkg/ccl/testccl/sqlccl/explain_test.go
@@ -140,7 +140,7 @@ func TestExplainGist(t *testing.T) {
 
 	// Use the release-build panic-catching behavior instead of the
 	// crdb_test-build behavior. This is needed so that some known bugs like
-	// #117101 don't result in a test failure.
+	// #119045 and #133129 don't result in a test failure.
 	defer colexecerror.ProductionBehaviorForTests()()
 
 	ctx := context.Background()
@@ -195,7 +195,6 @@ func TestExplainGist(t *testing.T) {
 			if err != nil && strings.Contains(err.Error(), "internal error") {
 				// Ignore all errors except the internal ones.
 				for _, knownErr := range []string{
-					"invalid datum type given: RECORD, expected RECORD", // #117101
 					"expected equivalence dependants to be its closure", // #119045
 					"not in index", // #133129
 				} {
diff --git a/pkg/sql/logictest/testdata/logic_test/udf b/pkg/sql/logictest/testdata/logic_test/udf
index d63fcc1ba04f..480793e0ab80 100644
--- a/pkg/sql/logictest/testdata/logic_test/udf
+++ b/pkg/sql/logictest/testdata/logic_test/udf
@@ -925,4 +925,26 @@ SELECT f('0');
 statement ok
 DROP FUNCTION f;
 
+# Regression test for #117101 - update the resolved type of a VALUES row with a
+# UDF after the UDF is built.
+subtest regression_117101
+
+statement ok
+CREATE FUNCTION f117101() RETURNS RECORD CALLED ON NULL INPUT AS $funcbody$
+  SELECT ((42)::INT8, (43)::INT8)
+$funcbody$ LANGUAGE SQL;
+
+statement error pgcode 42804 pq: VALUES types tuple{int, int} and tuple{tuple{string, int}, unknown} cannot be matched
+SELECT
+  *
+FROM
+(
+  VALUES
+    (
+     (('aloha'::TEXT,
+     (44)::INT8), NULL)
+    ),
+    (COALESCE(f117101(), NULL))
+);
+
 subtest end
diff --git a/pkg/sql/opt/optbuilder/values.go b/pkg/sql/opt/optbuilder/values.go
index d5782929b69c..ea353ef610d4 100644
--- a/pkg/sql/opt/optbuilder/values.go
+++ b/pkg/sql/opt/optbuilder/values.go
@@ -75,6 +75,14 @@ func (b *Builder) buildValuesClause(
 			// resolving the column types.
 			elems[elemPos] = b.buildScalar(texpr, inScope, nil, nil, nil)
 			elemPos += numCols
+			// Type-check the expression once again in order to update expressions
+			// that wrap a UDF 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.
+			texpr, err = tree.TypeCheck(b.ctx, texpr, b.semaCtx, texpr.ResolvedType())
+			if err != nil {
+				panic(err)
+			}
 			if typ := texpr.ResolvedType(); typ.Family() != types.UnknownFamily {
 				if colTypes[colIdx].Family() == types.UnknownFamily {
 					colTypes[colIdx] = typ
diff --git a/pkg/sql/sem/tree/type_check.go b/pkg/sql/sem/tree/type_check.go
index ec1c302fc076..8a8eaf8d75fc 100644
--- a/pkg/sql/sem/tree/type_check.go
+++ b/pkg/sql/sem/tree/type_check.go
@@ -1160,6 +1160,12 @@ 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.
+		return expr, nil
+	}
+
 	searchPath := EmptySearchPath
 	var resolver FunctionReferenceResolver
 	if semaCtx != nil {