From 22dabb08e76decbe2fc1c091431e05ebfba5e73f Mon Sep 17 00:00:00 2001 From: Rafi Shamim Date: Fri, 23 Jun 2023 16:23:12 -0400 Subject: [PATCH] tree,optbuilder: store annotations for each UDF body statement This also fixes a related issue where user-defined types that was references in UDF bodies would not be shown in the CREATE FUNCTION statement. This problem was being masked by the main issue that is getting fixed. Release note (bug fix): Previously, referencing a user-defined type in the body of a user-defined function would result in an error at the time of creating the function. This is now fixed. --- pkg/sql/function_resolver_test.go | 2 +- .../testdata/logic_test/crdb_internal_catalog | 2 +- pkg/sql/logictest/testdata/logic_test/udf | 55 ++++++++++++++++++- pkg/sql/opt/optbuilder/create_function.go | 20 +++++-- .../opt/optbuilder/testdata/create_function | 4 +- .../schemachanger/scbuild/builder_state.go | 2 - pkg/sql/sem/tree/udf.go | 8 +++ 7 files changed, 81 insertions(+), 12 deletions(-) diff --git a/pkg/sql/function_resolver_test.go b/pkg/sql/function_resolver_test.go index f59b380ea83d..69249264c47b 100644 --- a/pkg/sql/function_resolver_test.go +++ b/pkg/sql/function_resolver_test.go @@ -379,7 +379,7 @@ CREATE FUNCTION sc1.lower(a STRING) RETURNS STRING VOLATILE LANGUAGE SQL AS $$ S exprStr: "lower('HI')", searchPath: []string{"sc1", "sc2", "pg_catalog"}, expectedFuncOID: 100109, - expectedFuncBody: "SELECT lower('HI');", + expectedFuncBody: "SELECT lower('HI':::STRING);", desiredType: types.String, }, } diff --git a/pkg/sql/logictest/testdata/logic_test/crdb_internal_catalog b/pkg/sql/logictest/testdata/logic_test/crdb_internal_catalog index 2e43cc49745d..b1ea948400e8 100644 --- a/pkg/sql/logictest/testdata/logic_test/crdb_internal_catalog +++ b/pkg/sql/logictest/testdata/logic_test/crdb_internal_catalog @@ -170,7 +170,7 @@ SELECT id, strip_volatile(descriptor) FROM crdb_internal.kv_catalog_descriptor O 110 {"type": {"alias": {"arrayContents": {"family": "EnumFamily", "oid": 100109, "udtMetadata": {"arrayTypeOid": 100110}}, "arrayElemType": "EnumFamily", "family": "ArrayFamily", "oid": 100110}, "id": 110, "kind": "ALIAS", "name": "_greeting", "parentId": 106, "parentSchemaId": 108, "privileges": {"ownerProto": "root", "users": [{"privileges": "2", "userProto": "admin", "withGrantOption": "2"}, {"privileges": "512", "userProto": "public"}, {"privileges": "2", "userProto": "root", "withGrantOption": "2"}], "version": 2}, "version": "1"}} 111 {"table": {"checks": [{"columnIds": [1], "constraintId": 2, "expr": "k > 0:::INT8", "name": "ck"}], "columns": [{"id": 1, "name": "k", "type": {"family": "IntFamily", "oid": 20, "width": 64}}, {"id": 2, "name": "v", "nullable": true, "type": {"family": "StringFamily", "oid": 25}}], "dependedOnBy": [{"columnIds": [1, 2], "id": 112}], "formatVersion": 3, "id": 111, "name": "kv", "nextColumnId": 3, "nextConstraintId": 3, "nextIndexId": 2, "nextMutationId": 1, "parentId": 106, "primaryIndex": {"constraintId": 1, "encodingType": 1, "foreignKey": {}, "geoConfig": {}, "id": 1, "interleave": {}, "keyColumnDirections": ["ASC"], "keyColumnIds": [1], "keyColumnNames": ["k"], "name": "kv_pkey", "partitioning": {}, "sharded": {}, "storeColumnIds": [2], "storeColumnNames": ["v"], "unique": true, "version": 4}, "privileges": {"ownerProto": "root", "users": [{"privileges": "2", "userProto": "admin", "withGrantOption": "2"}, {"privileges": "2", "userProto": "root", "withGrantOption": "2"}], "version": 2}, "replacementOf": {"time": {}}, "unexposedParentSchemaId": 107, "version": "4"}} 112 {"table": {"columns": [{"id": 1, "name": "k", "nullable": true, "type": {"family": "IntFamily", "oid": 20, "width": 64}}, {"id": 2, "name": "v", "nullable": true, "type": {"family": "StringFamily", "oid": 25}}, {"defaultExpr": "unique_rowid()", "hidden": true, "id": 3, "name": "rowid", "type": {"family": "IntFamily", "oid": 20, "width": 64}}], "dependsOn": [111], "formatVersion": 3, "id": 112, "indexes": [{"createdExplicitly": true, "foreignKey": {}, "geoConfig": {}, "id": 2, "interleave": {}, "keyColumnDirections": ["ASC"], "keyColumnIds": [2], "keyColumnNames": ["v"], "keySuffixColumnIds": [3], "name": "idx", "partitioning": {}, "sharded": {}, "version": 4}], "isMaterializedView": true, "name": "mv", "nextColumnId": 4, "nextConstraintId": 2, "nextIndexId": 4, "nextMutationId": 1, "parentId": 106, "primaryIndex": {"constraintId": 1, "encodingType": 1, "foreignKey": {}, "geoConfig": {}, "id": 1, "interleave": {}, "keyColumnDirections": ["ASC"], "keyColumnIds": [3], "keyColumnNames": ["rowid"], "name": "mv_pkey", "partitioning": {}, "sharded": {}, "storeColumnIds": [1, 2], "storeColumnNames": ["k", "v"], "unique": true, "version": 4}, "privileges": {"ownerProto": "root", "users": [{"privileges": "2", "userProto": "admin", "withGrantOption": "2"}, {"privileges": "2", "userProto": "root", "withGrantOption": "2"}], "version": 2}, "replacementOf": {"time": {}}, "unexposedParentSchemaId": 107, "version": "8", "viewQuery": "SELECT k, v FROM db.public.kv"}} -113 {"function": {"functionBody": "SELECT json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(d, ARRAY['table', 'families']), ARRAY['table', 'nextFamilyId']), ARRAY['table', 'indexes', '0', 'createdAtNanos']), ARRAY['table', 'indexes', '1', 'createdAtNanos']), ARRAY['table', 'indexes', '2', 'createdAtNanos']), ARRAY['table', 'primaryIndex', 'createdAtNanos']), ARRAY['table', 'createAsOfTime']), ARRAY['table', 'modificationTime']), ARRAY['function', 'modificationTime']), ARRAY['type', 'modificationTime']), ARRAY['schema', 'modificationTime']), ARRAY['database', 'modificationTime']);", "id": 113, "lang": "SQL", "name": "strip_volatile", "nullInputBehavior": "CALLED_ON_NULL_INPUT", "params": [{"class": "IN", "name": "d", "type": {"family": "JsonFamily", "oid": 3802}}], "parentId": 104, "parentSchemaId": 105, "privileges": {"ownerProto": "root", "users": [{"privileges": "2", "userProto": "admin", "withGrantOption": "2"}, {"privileges": "2", "userProto": "root", "withGrantOption": "2"}], "version": 2}, "returnType": {"type": {"family": "JsonFamily", "oid": 3802}}, "version": "1", "volatility": "STABLE"}} +113 {"function": {"functionBody": "SELECT json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(json_remove_path(d, ARRAY['table':::STRING, 'families':::STRING]:::STRING[]), ARRAY['table':::STRING, 'nextFamilyId':::STRING]:::STRING[]), ARRAY['table':::STRING, 'indexes':::STRING, '0':::STRING, 'createdAtNanos':::STRING]:::STRING[]), ARRAY['table':::STRING, 'indexes':::STRING, '1':::STRING, 'createdAtNanos':::STRING]:::STRING[]), ARRAY['table':::STRING, 'indexes':::STRING, '2':::STRING, 'createdAtNanos':::STRING]:::STRING[]), ARRAY['table':::STRING, 'primaryIndex':::STRING, 'createdAtNanos':::STRING]:::STRING[]), ARRAY['table':::STRING, 'createAsOfTime':::STRING]:::STRING[]), ARRAY['table':::STRING, 'modificationTime':::STRING]:::STRING[]), ARRAY['function':::STRING, 'modificationTime':::STRING]:::STRING[]), ARRAY['type':::STRING, 'modificationTime':::STRING]:::STRING[]), ARRAY['schema':::STRING, 'modificationTime':::STRING]:::STRING[]), ARRAY['database':::STRING, 'modificationTime':::STRING]:::STRING[]);", "id": 113, "lang": "SQL", "name": "strip_volatile", "nullInputBehavior": "CALLED_ON_NULL_INPUT", "params": [{"class": "IN", "name": "d", "type": {"family": "JsonFamily", "oid": 3802}}], "parentId": 104, "parentSchemaId": 105, "privileges": {"ownerProto": "root", "users": [{"privileges": "2", "userProto": "admin", "withGrantOption": "2"}, {"privileges": "2", "userProto": "root", "withGrantOption": "2"}], "version": 2}, "returnType": {"type": {"family": "JsonFamily", "oid": 3802}}, "version": "1", "volatility": "STABLE"}} 4294966979 {"table": {"columns": [{"id": 1, "name": "srid", "nullable": true, "type": {"family": "IntFamily", "oid": 20, "width": 64}}, {"id": 2, "name": "auth_name", "nullable": true, "type": {"family": "StringFamily", "oid": 1043, "visibleType": 7, "width": 256}}, {"id": 3, "name": "auth_srid", "nullable": true, "type": {"family": "IntFamily", "oid": 20, "width": 64}}, {"id": 4, "name": "srtext", "nullable": true, "type": {"family": "StringFamily", "oid": 1043, "visibleType": 7, "width": 2048}}, {"id": 5, "name": "proj4text", "nullable": true, "type": {"family": "StringFamily", "oid": 1043, "visibleType": 7, "width": 2048}}], "formatVersion": 3, "id": 4294966979, "name": "spatial_ref_sys", "nextColumnId": 6, "nextConstraintId": 2, "nextIndexId": 2, "nextMutationId": 1, "primaryIndex": {"constraintId": 1, "foreignKey": {}, "geoConfig": {}, "id": 1, "interleave": {}, "partitioning": {}, "sharded": {}}, "privileges": {"ownerProto": "node", "users": [{"privileges": "32", "userProto": "public"}], "version": 2}, "replacementOf": {"time": {}}, "unexposedParentSchemaId": 4294966982, "version": "1"}} 4294966980 {"table": {"columns": [{"id": 1, "name": "f_table_catalog", "nullable": true, "type": {"family": 11, "oid": 19}}, {"id": 2, "name": "f_table_schema", "nullable": true, "type": {"family": 11, "oid": 19}}, {"id": 3, "name": "f_table_name", "nullable": true, "type": {"family": 11, "oid": 19}}, {"id": 4, "name": "f_geometry_column", "nullable": true, "type": {"family": 11, "oid": 19}}, {"id": 5, "name": "coord_dimension", "nullable": true, "type": {"family": "IntFamily", "oid": 20, "width": 64}}, {"id": 6, "name": "srid", "nullable": true, "type": {"family": "IntFamily", "oid": 20, "width": 64}}, {"id": 7, "name": "type", "nullable": true, "type": {"family": "StringFamily", "oid": 25}}], "formatVersion": 3, "id": 4294966980, "name": "geometry_columns", "nextColumnId": 8, "nextConstraintId": 2, "nextIndexId": 2, "nextMutationId": 1, "primaryIndex": {"constraintId": 1, "foreignKey": {}, "geoConfig": {}, "id": 1, "interleave": {}, "partitioning": {}, "sharded": {}}, "privileges": {"ownerProto": "node", "users": [{"privileges": "32", "userProto": "public"}], "version": 2}, "replacementOf": {"time": {}}, "unexposedParentSchemaId": 4294966982, "version": "1"}} 4294966981 {"table": {"columns": [{"id": 1, "name": "f_table_catalog", "nullable": true, "type": {"family": 11, "oid": 19}}, {"id": 2, "name": "f_table_schema", "nullable": true, "type": {"family": 11, "oid": 19}}, {"id": 3, "name": "f_table_name", "nullable": true, "type": {"family": 11, "oid": 19}}, {"id": 4, "name": "f_geography_column", "nullable": true, "type": {"family": 11, "oid": 19}}, {"id": 5, "name": "coord_dimension", "nullable": true, "type": {"family": "IntFamily", "oid": 20, "width": 64}}, {"id": 6, "name": "srid", "nullable": true, "type": {"family": "IntFamily", "oid": 20, "width": 64}}, {"id": 7, "name": "type", "nullable": true, "type": {"family": "StringFamily", "oid": 25}}], "formatVersion": 3, "id": 4294966981, "name": "geography_columns", "nextColumnId": 8, "nextConstraintId": 2, "nextIndexId": 2, "nextMutationId": 1, "primaryIndex": {"constraintId": 1, "foreignKey": {}, "geoConfig": {}, "id": 1, "interleave": {}, "partitioning": {}, "sharded": {}}, "privileges": {"ownerProto": "node", "users": [{"privileges": "32", "userProto": "public"}], "version": 2}, "replacementOf": {"time": {}}, "unexposedParentSchemaId": 4294966982, "version": "1"}} diff --git a/pkg/sql/logictest/testdata/logic_test/udf b/pkg/sql/logictest/testdata/logic_test/udf index 855b6ea5805d..fd1cb12e42fc 100644 --- a/pkg/sql/logictest/testdata/logic_test/udf +++ b/pkg/sql/logictest/testdata/logic_test/udf @@ -627,7 +627,7 @@ CREATE FUNCTION public.test_vf_f() CALLED ON NULL INPUT LANGUAGE SQL AS $$ - SELECT lower('hello'); + SELECT lower('hello':::STRING); $$ @@ -3096,7 +3096,7 @@ ORDER BY oid; ---- 100273 f_93314 105 1546506610 14 false false false v 0 100272 · {} NULL SELECT i, e FROM test.public.t_93314 ORDER BY i LIMIT 1; 100275 f_93314_alias 105 1546506610 14 false false false v 0 100274 · {} NULL SELECT i, e FROM test.public.t_93314_alias ORDER BY i LIMIT 1; -100279 f_93314_comp 105 1546506610 14 false false false v 0 100276 · {} NULL SELECT (1, 2); +100279 f_93314_comp 105 1546506610 14 false false false v 0 100276 · {} NULL SELECT (1:::INT8, 2:::INT8); 100280 f_93314_comp_t 105 1546506610 14 false false false v 0 100278 · {} NULL SELECT a, c FROM test.public.t_93314_comp LIMIT 1; # Regression test for #95240. Strict UDFs that are inlined should result in NULL @@ -3574,3 +3574,54 @@ query T SELECT f104927() ---- [{"i": 1, "s": "foo"}] + +subtest end + +# Regression test for https://github.com/cockroachdb/cockroach/issues/104242. +# Verify that statements that use UDFs with unresolved names that require +# AST annotations work correctly. +subtest udf_with_unresolved_names + +statement ok +CREATE TABLE tab104242 (a INT); + +statement ok +CREATE TYPE typ104242 AS ENUM ('foo'); + +statement ok +CREATE FUNCTION func104242() RETURNS INT LANGUAGE SQL AS $$ + SELECT 1 FROM tab104242 WHERE NULL::typ104242 IN () +$$; + +query T +SELECT create_statement FROM [SHOW CREATE FUNCTION func104242] +---- +CREATE FUNCTION public.func104242() + RETURNS INT8 + VOLATILE + NOT LEAKPROOF + CALLED ON NULL INPUT + LANGUAGE SQL + AS $$ + SELECT 1 FROM test.public.tab104242 WHERE NULL IN (); +$$ + +statement ok +CREATE FUNCTION func104242_not_null() RETURNS INT LANGUAGE SQL AS $$ + SELECT 1 FROM tab104242 WHERE 'foo'::typ104242 IN () +$$; + +query T +SELECT create_statement FROM [SHOW CREATE FUNCTION func104242_not_null] +---- +CREATE FUNCTION public.func104242_not_null() + RETURNS INT8 + VOLATILE + NOT LEAKPROOF + CALLED ON NULL INPUT + LANGUAGE SQL + AS $$ + SELECT 1 FROM test.public.tab104242 WHERE 'foo':::test.public.typ104242 IN (); +$$ + +subtest end diff --git a/pkg/sql/opt/optbuilder/create_function.go b/pkg/sql/opt/optbuilder/create_function.go index 72697c5e6fbf..3dd139d80133 100644 --- a/pkg/sql/opt/optbuilder/create_function.go +++ b/pkg/sql/opt/optbuilder/create_function.go @@ -49,12 +49,16 @@ func (b *Builder) buildCreateFunction(cf *tree.CreateFunction, inScope *scope) ( b.trackSchemaDeps = true // Make sure datasource names are qualified. b.qualifyDataSourceNamesInAST = true + oldEvalCtxAnn := b.evalCtx.Annotations + oldSemaCtxAnn := b.semaCtx.Annotations defer func() { b.insideFuncDef = false b.trackSchemaDeps = false b.schemaDeps = nil b.schemaTypeDeps = intsets.Fast{} b.qualifyDataSourceNamesInAST = false + b.evalCtx.Annotations = oldEvalCtxAnn + b.semaCtx.Annotations = oldSemaCtxAnn b.semaCtx.FunctionResolver = preFuncResolver switch recErr := recover().(type) { @@ -162,8 +166,18 @@ func (b *Builder) buildCreateFunction(cf *tree.CreateFunction, inScope *scope) ( targetVolatility := tree.GetFuncVolatility(cf.Options) // Validate each statement and collect the dependencies. - fmtCtx := tree.NewFmtCtx(tree.FmtSimple) + fmtCtx := tree.NewFmtCtx(tree.FmtSerializable) for i, stmt := range stmts { + // Add statement ast into CreateFunction node for logging purpose, and set + // the annotations for this statement so names can be resolved. + cf.BodyStatements = append(cf.BodyStatements, stmt.AST) + ann := tree.MakeAnnotations(stmt.NumAnnotations) + cf.BodyAnnotations = append(cf.BodyAnnotations, &ann) + + // The defer logic will reset the annotations to the old value. + b.semaCtx.Annotations = ann + b.evalCtx.Annotations = &ann + var stmtScope *scope // We need to disable stable function folding because we want to catch the // volatility of stable functions. If folded, we only get a scalar and lose @@ -190,8 +204,6 @@ func (b *Builder) buildCreateFunction(cf *tree.CreateFunction, inScope *scope) ( deps = append(deps, b.schemaDeps...) typeDeps.UnionWith(b.schemaTypeDeps) - // Add statement ast into CreateFunction node for logging purpose. - cf.BodyStatements = append(cf.BodyStatements, stmt.AST) // Reset the tracked dependencies for next statement. b.schemaDeps = nil b.schemaTypeDeps = intsets.Fast{} @@ -209,7 +221,7 @@ func (b *Builder) buildCreateFunction(cf *tree.CreateFunction, inScope *scope) ( // Override the function body so that references are fully qualified. for i, option := range cf.Options { if _, ok := option.(tree.FunctionBodyStr); ok { - cf.Options[i] = tree.FunctionBodyStr(fmtCtx.String()) + cf.Options[i] = tree.FunctionBodyStr(fmtCtx.CloseAndGetString()) break } } diff --git a/pkg/sql/opt/optbuilder/testdata/create_function b/pkg/sql/opt/optbuilder/testdata/create_function index d0b607ef956e..dfdd8975ba26 100644 --- a/pkg/sql/opt/optbuilder/testdata/create_function +++ b/pkg/sql/opt/optbuilder/testdata/create_function @@ -54,7 +54,7 @@ create-function ├── CREATE FUNCTION f() │ RETURNS STRING │ LANGUAGE SQL - │ AS $$SELECT 'MON'::STRING;$$ + │ AS $$SELECT x'40':::@100001::STRING;$$ └── dependencies └── workday @@ -102,7 +102,7 @@ create-function │ RETURNS INT8 │ LANGUAGE SQL │ AS $$SELECT a FROM t.public.ab; - │ SELECT nextval('s');$$ + │ SELECT nextval('s':::STRING);$$ └── dependencies ├── ab [columns: a] └── s diff --git a/pkg/sql/schemachanger/scbuild/builder_state.go b/pkg/sql/schemachanger/scbuild/builder_state.go index 2cd0fadfe774..57a40a4cd79f 100644 --- a/pkg/sql/schemachanger/scbuild/builder_state.go +++ b/pkg/sql/schemachanger/scbuild/builder_state.go @@ -1539,8 +1539,6 @@ func (b *builderState) serializeUserDefinedTypes(queryStr string) string { default: return true, expr, nil } - // semaCtx may be nil if this is a virtual view being created at - // init time. var typ *types.T typ, err = tree.ResolveType(b.ctx, typRef, b.semaCtx.TypeResolver) if err != nil { diff --git a/pkg/sql/sem/tree/udf.go b/pkg/sql/sem/tree/udf.go index 4fae92fb8bff..7e4c96715150 100644 --- a/pkg/sql/sem/tree/udf.go +++ b/pkg/sql/sem/tree/udf.go @@ -94,6 +94,9 @@ type CreateFunction struct { // format. That is sequence names and type name in expressions are not // rewritten with OIDs. BodyStatements Statements + // BodyAnnotations is not assigned during initial parsing of user input. It's + // assigned by the opt builder when the optimizer parses the body statements. + BodyAnnotations []*Annotations } // Format implements the NodeFormatter interface. @@ -131,8 +134,11 @@ func (node *CreateFunction) Format(ctx *FmtCtx) { if i > 0 { ctx.WriteString(" ") } + oldAnn := ctx.ann + ctx.ann = node.BodyAnnotations[i] ctx.FormatNode(stmt) ctx.WriteString(";") + ctx.ann = oldAnn } ctx.WriteString("$$") } else if node.RoutineBody != nil { @@ -149,6 +155,8 @@ func (node *CreateFunction) Format(ctx *FmtCtx) { // RoutineBody represent a list of statements in a UDF body. type RoutineBody struct { + // Stmts is populated during parsing. Unlike BodyStatements, we don't need + // to create separate Annotations for each statement. Stmts Statements }