Merge cockroachdb#85911

85911: sql: get create function statements from all dbs when db is empty r=rafiss a=chengxiong-ruan Fixes cockroachdb#85897 When db name is a quoted empty string, we need to loop through all databases to get function descriptor to construct the crdb_internal.create_function_statements table. Release note: None Co-authored-by: Chengxiong Ruan <[email protected]>
tbg · Aug 11, 2022 · fb8b849 · fb8b849
2 parents 9bd865d + be68623
commit fb8b849
Show file tree

Hide file tree

Showing 2 changed files with 99 additions and 34 deletions.
diff --git a/pkg/sql/crdb_internal.go b/pkg/sql/crdb_internal.go
@@ -2591,44 +2591,60 @@ CREATE TABLE crdb_internal.create_function_statements (
 )
 `,
 	populate: func(ctx context.Context, p *planner, db catalog.DatabaseDescriptor, addRow func(...tree.Datum) error) error {
-		return forEachSchema(ctx, p, db, func(sc catalog.SchemaDescriptor) error {
-			return sc.ForEachFunctionOverload(func(overload descpb.SchemaDescriptor_FunctionOverload) error {
-				fnDesc, err := p.Descriptors().GetImmutableFunctionByID(ctx, p.txn, overload.ID, tree.ObjectLookupFlags{})
-				if err != nil {
-					return err
-				}
-				treeNode, err := fnDesc.ToCreateExpr()
-				treeNode.FuncName.ObjectNamePrefix = tree.ObjectNamePrefix{
-					ExplicitSchema: true,
-					SchemaName:     tree.Name(sc.GetName()),
-				}
-				if err != nil {
-					return err
-				}
-				for i := range treeNode.Options {
-					if body, ok := treeNode.Options[i].(tree.FunctionBodyStr); ok {
-						stmtStrs := strings.Split(string(body), "\n")
-						for i := range stmtStrs {
-							stmtStrs[i] = "\t" + stmtStrs[i]
-						}
+		var dbDescs []catalog.DatabaseDescriptor
+		if db == nil {
+			var err error
+			dbDescs, err = p.Descriptors().GetAllDatabaseDescriptors(ctx, p.Txn())
+			if err != nil {
+				return err
+			}
+		} else {
+			dbDescs = append(dbDescs, db)
+		}
+		for _, db := range dbDescs {
+			err := forEachSchema(ctx, p, db, func(sc catalog.SchemaDescriptor) error {
+				return sc.ForEachFunctionOverload(func(overload descpb.SchemaDescriptor_FunctionOverload) error {
+					fnDesc, err := p.Descriptors().GetImmutableFunctionByID(ctx, p.txn, overload.ID, tree.ObjectLookupFlags{})
+					if err != nil {
+						return err
+					}
+					treeNode, err := fnDesc.ToCreateExpr()
+					treeNode.FuncName.ObjectNamePrefix = tree.ObjectNamePrefix{
+						ExplicitSchema: true,
+						SchemaName:     tree.Name(sc.GetName()),
+					}
+					if err != nil {
+						return err
+					}
+					for i := range treeNode.Options {
+						if body, ok := treeNode.Options[i].(tree.FunctionBodyStr); ok {
+							stmtStrs := strings.Split(string(body), "\n")
+							for i := range stmtStrs {
+								stmtStrs[i] = "\t" + stmtStrs[i]
+							}
 
-						p := &treeNode.Options[i]
-						// Add two new lines just for better formatting.
-						*p = "\n" + tree.FunctionBodyStr(strings.Join(stmtStrs, "\n")) + "\n"
+							p := &treeNode.Options[i]
+							// Add two new lines just for better formatting.
+							*p = "\n" + tree.FunctionBodyStr(strings.Join(stmtStrs, "\n")) + "\n"
+						}
 					}
-				}
 
-				return addRow(
-					tree.NewDInt(tree.DInt(db.GetID())),      // database_id
-					tree.NewDString(db.GetName()),            // database_name
-					tree.NewDInt(tree.DInt(sc.GetID())),      // schema_id
-					tree.NewDString(sc.GetName()),            // schema_name
-					tree.NewDInt(tree.DInt(fnDesc.GetID())),  // function_id
-					tree.NewDString(fnDesc.GetName()),        //function_name
-					tree.NewDString(tree.AsString(treeNode)), // create_statement
-				)
+					return addRow(
+						tree.NewDInt(tree.DInt(db.GetID())),      // database_id
+						tree.NewDString(db.GetName()),            // database_name
+						tree.NewDInt(tree.DInt(sc.GetID())),      // schema_id
+						tree.NewDString(sc.GetName()),            // schema_name
+						tree.NewDInt(tree.DInt(fnDesc.GetID())),  // function_id
+						tree.NewDString(fnDesc.GetName()),        //function_name
+						tree.NewDString(tree.AsString(treeNode)), // create_statement
+					)
+				})
 			})
-		})
+			if err != nil {
+				return err
+			}
+		}
+		return nil
 	},
 }
 

diff --git a/pkg/sql/logictest/testdata/logic_test/udf b/pkg/sql/logictest/testdata/logic_test/udf
@@ -439,6 +439,55 @@ CREATE FUNCTION sc.proc_f_2(IN STRING)
     SELECT 'hello';
 $$  104  test  127  sc  128  proc_f_2
 
+statement ok
+CREATE DATABASE test_cross_db;
+USE test_cross_db;
+CREATE FUNCTION f_cross_db() RETURNS INT LANGUAGE SQL AS $$ SELECT 1 $$;
+USE test;
+
+query TITITIT
+SELECT create_statement, database_id, database_name, schema_id, schema_name, function_id, function_name
+FROM "".crdb_internal.create_function_statements
+WHERE function_name IN ('proc_f', 'proc_f_2', 'f_cross_db')
+ORDER BY database_id, function_name;
+----
+CREATE FUNCTION public.proc_f(IN INT8)
+    RETURNS INT8
+    VOLATILE
+    NOT LEAKPROOF
+    CALLED ON NULL INPUT
+    LANGUAGE SQL
+    AS $$
+    SELECT 1;
+$$  104  test  105  public  124  proc_f
+CREATE FUNCTION public.proc_f(IN STRING, IN b INT8)
+    RETURNS SETOF STRING
+    IMMUTABLE
+    LEAKPROOF
+    STRICT
+    LANGUAGE SQL
+    AS $$
+    SELECT 'hello';
+$$  104  test  105  public  125  proc_f
+CREATE FUNCTION sc.proc_f_2(IN STRING)
+    RETURNS STRING
+    VOLATILE
+    NOT LEAKPROOF
+    CALLED ON NULL INPUT
+    LANGUAGE SQL
+    AS $$
+    SELECT 'hello';
+$$  104  test  127  sc  128  proc_f_2
+CREATE FUNCTION public.f_cross_db()
+    RETURNS INT8
+    VOLATILE
+    NOT LEAKPROOF
+    CALLED ON NULL INPUT
+    LANGUAGE SQL
+    AS $$
+    SELECT 1;
+$$  129  test_cross_db  130  public  131  f_cross_db
+
 subtest show_create_function
 
 query T