Merge #149901

149901: backup/rewrite: handle UDF references from views when restoring r=rafiss a=rafiss

When we added support for views referencing UDFs, we did not account for these references in the RESTORE logic. This patch adds the logic to rewrite those references when needed.

No release note since this bug was not released.

fixes #149782
Release note: None

Co-authored-by: Rafi Shamim <[email protected]>

Author: craig[bot]

pkg/backup/backup_test.go

@@ -2984,7 +2984,7 @@ func TestBackupRestoreCrossTableReferences(t *testing.T) {
 		db := sqlutils.MakeSQLRunner(tc.Conns[0])
 		db.Exec(t, createStore)
 		db.ExpectErr(
-			t, `cannot restore view "early_customers" without restoring referenced table`,
+			t, `cannot restore view "early_customers" without restoring referenced object`,
 			`RESTORE TABLE store.early_customers FROM LATEST IN $1`, localFoo,
 		)
 		db.Exec(t, `RESTORE TABLE store.early_customers, store.customers, store.orders FROM LATEST IN $1`, localFoo)
@@ -3015,15 +3015,15 @@ func TestBackupRestoreCrossTableReferences(t *testing.T) {
 		db := sqlutils.MakeSQLRunner(tc.Conns[0])
 
 		db.ExpectErr(
-			t, `cannot restore view "ordercounts" without restoring referenced table`,
+			t, `cannot restore view "ordercounts" without restoring referenced object`,
 			`RESTORE DATABASE storestats FROM LATEST IN $1`, localFoo,
 		)
 
 		db.Exec(t, createStore)
 		db.Exec(t, createStoreStats)
 
 		db.ExpectErr(
-			t, `cannot restore view "ordercounts" without restoring referenced table`,
+			t, `cannot restore view "ordercounts" without restoring referenced object`,
 			`RESTORE TABLE storestats.ordercounts, store.customers FROM LATEST IN $1`, localFoo,
 		)
 
@@ -5404,7 +5404,7 @@ func TestBackupRestoreSequencesInViews(t *testing.T) {
 		sqlDB.Exec(t, `DROP VIEW v`)
 		// Restore v.
 		sqlDB.ExpectErr(
-			t, "pq: cannot restore view \"v\" without restoring referenced table \\(or \"skip_missing_views\" option\\)",
+			t, "pq: cannot restore view \"v\" without restoring referenced object \\(or \"skip_missing_views\" option\\)",
 			`RESTORE TABLE v FROM LATEST IN 'nodelocal://1/test/'`,
 		)
 	})

pkg/backup/restore_planning.go

@@ -107,22 +107,32 @@ var restoreCompactedBackups = settings.RegisterBoolSetting(
 func maybeFilterMissingViews(
 	tablesByID map[descpb.ID]*tabledesc.Mutable,
 	typesByID map[descpb.ID]*typedesc.Mutable,
+	functionsByID map[descpb.ID]*funcdesc.Mutable,
 	skipMissingViews bool,
+	skipMissingUDFs bool,
 ) (map[descpb.ID]*tabledesc.Mutable, error) {
 	// Function that recursively determines whether a given table, if it is a
 	// view, has valid dependencies. Dependencies are looked up in tablesByID.
+	missingOnlyFunctionDeps := true
 	var hasValidViewDependencies func(desc *tabledesc.Mutable) bool
 	hasValidViewDependencies = func(desc *tabledesc.Mutable) bool {
 		if !desc.IsView() {
 			return true
 		}
 		for _, id := range desc.DependsOn {
 			if depDesc, ok := tablesByID[id]; !ok || !hasValidViewDependencies(depDesc) {
+				missingOnlyFunctionDeps = false
 				return false
 			}
 		}
 		for _, id := range desc.DependsOnTypes {
 			if _, ok := typesByID[id]; !ok {
+				missingOnlyFunctionDeps = false
+				return false
+			}
+		}
+		for _, id := range desc.DependsOnFunctions {
+			if _, ok := functionsByID[id]; !ok {
 				return false
 			}
 		}
@@ -135,8 +145,12 @@ func maybeFilterMissingViews(
 			filteredTablesByID[id] = table
 		} else {
 			if !skipMissingViews {
+				if skipMissingUDFs && missingOnlyFunctionDeps {
+					// Skip this view since only function dependencies are missing.
+					continue
+				}
 				return nil, errors.Errorf(
-					"cannot restore view %q without restoring referenced table (or %q option)",
+					"cannot restore view %q without restoring referenced object (or %q option)",
 					table.Name, restoreOptSkipMissingViews,
 				)
 			}
@@ -1953,7 +1967,10 @@ func doRestorePlan(
 	filteredTablesByID, err := maybeFilterMissingViews(
 		tablesByID,
 		typesByID,
-		restoreStmt.Options.SkipMissingViews)
+		functionsByID,
+		restoreStmt.Options.SkipMissingViews,
+		restoreStmt.Options.SkipMissingUDFs,
+	)
 	if err != nil {
 		return err
 	}

pkg/backup/testdata/backup-restore/views

@@ -66,3 +66,174 @@ db1_new.sc1.v1 CREATE VIEW sc1.v1 (
 	a,
 	enum1
 ) AS SELECT a, 'Good':::sc1.enum1 FROM db1_new.sc1.tbl1;
+
+# Test backing up and restoring a view that references a UDF.
+# Regression test for issue #149782.
+
+subtest restore_table_view_with_skip_missing
+
+exec-sql
+CREATE DATABASE udf_db;
+----
+
+exec-sql
+USE udf_db;
+----
+
+exec-sql
+CREATE FUNCTION f1(x INT) RETURNS INT AS $$
+BEGIN
+  RETURN x + 1;
+END;
+$$ LANGUAGE PLpgSQL;
+----
+
+exec-sql
+CREATE TABLE t (id INT PRIMARY KEY, name VARCHAR(256), money INT);
+----
+
+exec-sql
+CREATE VIEW v AS SELECT id, name, f1(money) FROM t;
+----
+
+exec-sql
+INSERT INTO t VALUES (1, 'test', 100);
+----
+
+query-sql
+SELECT * FROM v;
+----
+1 test 101
+
+exec-sql
+BACKUP DATABASE udf_db INTO 'nodelocal://1/test_udf/';
+----
+
+exec-sql
+DROP VIEW v;
+----
+
+exec-sql
+DROP TABLE t;
+----
+
+exec-sql
+DROP FUNCTION f1;
+----
+
+# This should fail without one of the skip_missing_* options.
+exec-sql expect-error-regex=(cannot restore view "v" without restoring referenced object \(or "skip_missing_views" option\))
+RESTORE TABLE udf_db.v, udf_db.t FROM LATEST IN 'nodelocal://1/test_udf/'
+----
+regex matches error
+
+# With skip_missing_udfs option, the view should be skipped since the function is missing.
+exec-sql
+RESTORE TABLE udf_db.v, udf_db.t FROM LATEST IN 'nodelocal://1/test_udf/' WITH skip_missing_udfs;
+----
+
+# Check that the view was skipped.
+query-sql
+SELECT count(*) FROM information_schema.views WHERE table_name = 'v';
+----
+0
+
+# But the table should still be restored.
+query-sql
+SELECT * FROM t;
+----
+1 test 100
+
+exec-sql
+DROP TABLE t;
+----
+
+# With skip_missing_views option, the view should be skipped also.
+exec-sql
+RESTORE TABLE udf_db.v, udf_db.t FROM LATEST IN 'nodelocal://1/test_udf/' WITH skip_missing_views;
+----
+
+# Check that the view was skipped.
+query-sql
+SELECT count(*) FROM information_schema.views WHERE table_name = 'v';
+----
+0
+
+# But the table should still be restored.
+query-sql
+SELECT * FROM t;
+----
+1 test 100
+
+exec-sql
+DROP TABLE t;
+----
+
+exec-sql
+DROP DATABASE udf_db;
+----
+
+subtest end
+
+subtest restore_database_with_view_referencing_udf
+
+exec-sql
+CREATE DATABASE udf_db2;
+----
+
+exec-sql
+USE udf_db2;
+----
+
+exec-sql
+CREATE FUNCTION f1(x INT) RETURNS INT AS $$
+BEGIN
+  RETURN x + 1;
+END;
+$$ LANGUAGE PLpgSQL;
+----
+
+exec-sql
+CREATE TABLE t (id INT PRIMARY KEY, name VARCHAR(256), money INT);
+----
+
+exec-sql
+CREATE VIEW v AS SELECT id, name, f1(money) FROM t;
+----
+
+exec-sql
+INSERT INTO t VALUES (1, 'test', 100);
+----
+
+query-sql
+SELECT * FROM v;
+----
+1 test 101
+
+exec-sql
+BACKUP DATABASE udf_db2 INTO 'nodelocal://1/test_udf2/';
+----
+
+exec-sql
+DROP DATABASE udf_db2;
+----
+
+# This should restore the database with proper references between the view and the function.
+exec-sql
+RESTORE DATABASE udf_db2 FROM LATEST IN 'nodelocal://1/test_udf2/' WITH new_db_name='udf_db2_new';
+----
+
+exec-sql
+USE udf_db2_new;
+----
+
+query-sql
+SELECT * FROM v;
+----
+1 test 101
+
+exec-sql
+DROP DATABASE udf_db2_new;
+----
+
+subtest end

pkg/sql/catalog/rewrite/rewrite.go

@@ -218,6 +218,17 @@ func TableDescs(
 					table.Name, dest)
 			}
 		}
+		for i, dest := range table.DependsOnFunctions {
+			if depRewrite, ok := descriptorRewrites[dest]; ok {
+				table.DependsOnFunctions[i] = depRewrite.ID
+			} else {
+				// If skipMissingUDFs is set, views with missing function dependencies
+				// should have been filtered out in maybeFilterMissingViews.
+				return nil, errors.AssertionFailedf(
+					"cannot restore %q because referenced function %d was not found",
+					table.Name, dest)
+			}
+		}
 		origRefs := table.DependedOnBy
 		table.DependedOnBy = nil
 		for _, ref := range origRefs {