From c29c242818819c8b5b8e7a5c7f317220d8adfa85 Mon Sep 17 00:00:00 2001
From: Chengxiong Ruan <chengxiongruan@gmail.com>
Date: Thu, 9 Feb 2023 17:30:31 -0500
Subject: [PATCH] backupccl: skip UDFs not found when rewriting IDs in Schema

Previously, if a BACKUP is created by doing BACKUP TABLE, and
the parent schema of the tables has UDFs in it, RESTOREing the
tables from the back crashes. This is because we store the
signatures of UDFs in schema descriptor. Restoring a table also
restores the parent schema if there is no schema with the same
name exists in the target database. So when trying to rewrite
the UDF ids in the schema descriptor, it crashes because
function descriptors are not backed up at all.

This pr adds logic to skip those rewrites if function descriptor
was not backed up.

Fixes: #96910

Release note (enterprise): this path fixes a bug where server
would crash if trying to restore a table from a backup generated
with `BACKUP TABLE` from a schema with user defined functions,
and the restore target database doesn't have a schema with the
same name.
---
 .../backup-restore/user-defined-functions     | 30 +++++++++++++++++++
 pkg/sql/catalog/rewrite/rewrite.go            | 20 +++++++++++--
 2 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/pkg/ccl/backupccl/testdata/backup-restore/user-defined-functions b/pkg/ccl/backupccl/testdata/backup-restore/user-defined-functions
index f71274a81483..27883682a548 100644
--- a/pkg/ccl/backupccl/testdata/backup-restore/user-defined-functions
+++ b/pkg/ccl/backupccl/testdata/backup-restore/user-defined-functions
@@ -291,3 +291,33 @@ exec-sql
 DROP TYPE sc1.enum1
 ----
 pq: cannot drop type "enum1" because other objects ([db1.sc1.f1]) still depend on it
+
+# Make sure that backup and restore individual tables from schema with UDF does
+# not crash.
+new-cluster name=s3
+----
+
+exec-sql cluster=s3
+CREATE DATABASE db1;
+CREATE TABLE t(a INT PRIMARY KEY);
+CREATE FUNCTION f() RETURNS INT LANGUAGE SQL AS $$ SELECT 1 $$;
+----
+
+exec-sql
+BACKUP TABLE t INTO 'nodelocal://0/test/'
+----
+
+exec-sql
+RESTORE TABLE t FROM LATEST IN 'nodelocal://0/test/' WITH into_db = 'db1';
+----
+
+exec-sql
+USE db1;
+----
+
+# Make sure proper error message is returned when trying to resolve the
+# function from the restore target db.
+query-sql
+SELECT f()
+----
+pq: unknown function: f(): function undefined
diff --git a/pkg/sql/catalog/rewrite/rewrite.go b/pkg/sql/catalog/rewrite/rewrite.go
index ace81b73be26..58bd849798cf 100644
--- a/pkg/sql/catalog/rewrite/rewrite.go
+++ b/pkg/sql/catalog/rewrite/rewrite.go
@@ -536,10 +536,19 @@ func SchemaDescs(schemas []*schemadesc.Mutable, descriptorRewrites jobspb.DescRe
 		sc.ParentID = rewrite.ParentID
 
 		// Rewrite function ID and types ID in function signatures.
-		for _, fn := range sc.GetFunctions() {
+		newFns := make(map[string]descpb.SchemaDescriptor_Function)
+		for fnName, fn := range sc.GetFunctions() {
+			newSigs := make([]descpb.SchemaDescriptor_FunctionSignature, 0, len(fn.Signatures))
 			for i := range fn.Signatures {
 				sig := &fn.Signatures[i]
-				sig.ID = descriptorRewrites[sig.ID].ID
+				// If the function is not found in the backup, we just skip. This only
+				// happens when restoring from a backup with `BACKUP TABLE` where the
+				// function descriptors are not backup.
+				fnDesc, ok := descriptorRewrites[sig.ID]
+				if !ok {
+					continue
+				}
+				sig.ID = fnDesc.ID
 				for _, typ := range sig.ArgTypes {
 					if err := rewriteIDsInTypesT(typ, descriptorRewrites); err != nil {
 						return err
@@ -548,6 +557,13 @@ func SchemaDescs(schemas []*schemadesc.Mutable, descriptorRewrites jobspb.DescRe
 				if err := rewriteIDsInTypesT(sig.ReturnType, descriptorRewrites); err != nil {
 					return err
 				}
+				newSigs = append(newSigs, *sig)
+			}
+			if len(newSigs) > 0 {
+				newFns[fnName] = descpb.SchemaDescriptor_Function{
+					Name:       fnName,
+					Signatures: newSigs,
+				}
 			}
 		}