Skip to content

Commit

Permalink
Merge #105465
Browse files Browse the repository at this point in the history 
105465: tree,optbuilder: store annotations for each UDF body statement r=DrewKimball a=rafiss

fixes #104242

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.

Co-authored-by: Rafi Shamim <[email protected]>
  • Loading branch information
@rafiss
craig[bot] and rafiss committed Jun 26, 2023
2 parents 19b67e3 + 22dabb0 commit f49b836
Show file tree
Hide file tree
Showing 7 changed files with 81 additions and 12 deletions.
2 changes: 1 addition & 1 deletion pkg/sql/function_resolver_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
}
Expand Down
2 changes: 1 addition & 1 deletion pkg/sql/logictest/testdata/logic_test/crdb_internal_catalog
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"}}
Expand Down
55 changes: 53 additions & 2 deletions pkg/sql/logictest/testdata/logic_test/udf
View file
Original file line number Diff line number Diff line change
Expand Up @@ -627,7 +627,7 @@ CREATE FUNCTION public.test_vf_f()
CALLED ON NULL INPUT
LANGUAGE SQL
AS $$
SELECT lower('hello');
SELECT lower('hello':::STRING);
$$


Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
20 changes: 16 additions & 4 deletions pkg/sql/opt/optbuilder/create_function.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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
Expand All @@ -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{}
Expand All @@ -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
}
}
Expand Down
4 changes: 2 additions & 2 deletions pkg/sql/opt/optbuilder/testdata/create_function
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
2 changes: 0 additions & 2 deletions pkg/sql/schemachanger/scbuild/builder_state.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
8 changes: 8 additions & 0 deletions pkg/sql/sem/tree/udf.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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 {
Expand All @@ -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
}

Expand Down

0 comments on commit f49b836

Please sign in to comment.