Merge #97038

97038: sql: use udf in check constraint r=chengxiong-ruan a=chengxiong-ruan

Informs #87699

**sql: use UDFs in check constraints**

This commit turns on UDF usages in CHECK constraints.
It does a few things:
1. Add a concept of `SchemaExprContext`, which used to be a
pure string. This commit makes it a type so that we can
switch on different context of a expression. This allow fine
grain control over in which circumstance that a UDF is allowed.
2. Serialize Function Expression to use OID reference to UDF.
This allows objects reference UDFs by ID so that functions can
be renamed.
3. Refine cross reference validation in function descriptor
and table descriptor for more safety.
4. Fix legacy schema changer, so that cross references between
UDF and tables are properly tracked when CHECK constraints are
added and dropped.
5. Fix declarative schema changer, so that UDFs are recognized
in check constraint elements, and ops are added to add and remove
UDF back references.

Release note (sql change): previously UDFs are not allowed in
tables and any other object. This patch enables UDF usage in
CHECK constraints of tables in both legacy schema changer and
delcarative schema changer. Dependency circles are not allowed,
namely if a UDF depends on a table, then the table can't use that
UDF.


**sql: version gate udf usage check**

This commit adds version gate for udf usage check, so that
cross reference won't be broken because there won't be cases
like a constraint is added in new version, but dropped in
an old version.

Release note (sql change): This patch adds version gate so the
UDF usage in CHECK constraints are not allowed until cluster is
fully upgraded to 23.1.


**backupccl: fix backup and restore to work with objects referencing UDFs**

Release note (enterprise change): Previously UDFs can't be referenced
from other objects, so backup and restore oly needs to read and
write UDF descriptors without worrying about missing UDFs when
doing `RESTORE TABLE`. Now that we turned UDF's usage in table
CHECK constraints. We need to be able to handle all function ID
rewrites in CHECK constraint expressions during `RESTORE`. A new
`RSTORE` command option `skip_missing_udfs` is added, so that
when the option is given, UDF dependencies will be skipped if UDF
descriptors are missing. The main use case, as of this patch is that
 currently when doing `RESTORE TABLE`, we don't restore referenced
UDFs together, so UDFs are consider missing, which blocks the table
restore unless the `skip_missing_udfs` option is specified, so that
those check constraints are dropped to remove the dependency to
unblock `RESTORE TABLE`.


Co-authored-by: Chengxiong Ruan <[email protected]>

docs/generated/sql/bnf/stmt_block.bnf

@@ -1360,6 +1360,7 @@ unreserved_keyword ::=
 	| 'SKIP_MISSING_SEQUENCES'
 	| 'SKIP_MISSING_SEQUENCE_OWNERS'
 	| 'SKIP_MISSING_VIEWS'
+	| 'SKIP_MISSING_UDFS'
 	| 'SNAPSHOT'
 	| 'SPLIT'
 	| 'SQL'
@@ -2615,6 +2616,7 @@ restore_options ::=
 	| 'SKIP_MISSING_SEQUENCES'
 	| 'SKIP_MISSING_SEQUENCE_OWNERS'
 	| 'SKIP_MISSING_VIEWS'
+	| 'SKIP_MISSING_UDFS'
 	| 'DETACHED'
 	| 'SKIP_LOCALITIES_CHECK'
 	| 'DEBUG_PAUSE_ON' '=' string_or_placeholder
@@ -3438,6 +3440,7 @@ bare_label_keywords ::=
 	| 'TRANSFORM'
 	| 'VOLATILE'
 	| 'SETOF'
+	| 'SKIP_MISSING_VIEWS'
 
 opt_col_def_list_no_types ::=
 	'(' col_def_list_no_types ')'

pkg/ccl/backupccl/backup_telemetry.go

@@ -65,6 +65,7 @@ const (
 	telemetryOptionSkipMissingViews          = "skip_missing_views"
 	telemetryOptionSkipLocalitiesCheck       = "skip_localities_check"
 	telemetryOptionSchemaOnly                = "schema_only"
+	telemetryOptionSkipMissingUDFs           = "skip_missing_udfs"
 )
 
 // logBackupTelemetry publishes an eventpb.RecoveryEvent about a manually
@@ -396,6 +397,9 @@ func logRestoreTelemetry(
 	if opts.SkipMissingSequenceOwners {
 		options = append(options, telemetryOptionSkipMissingSequenceOwners)
 	}
+	if opts.SkipMissingUDFs {
+		options = append(options, telemetryOptionSkipMissingUDFs)
+	}
 	if opts.Detached {
 		options = append(options, telemetryOptionDetached)
 	}

pkg/ccl/backupccl/restore_planning.go

@@ -73,6 +73,7 @@ const (
 	restoreOptIntoDB                    = "into_db"
 	restoreOptSkipMissingFKs            = "skip_missing_foreign_keys"
 	restoreOptSkipMissingSequences      = "skip_missing_sequences"
+	restoreOptSkipMissingUDFs           = "skip_missing_udfs"
 	restoreOptSkipMissingSequenceOwners = "skip_missing_sequence_owners"
 	restoreOptSkipMissingViews          = "skip_missing_views"
 	restoreOptSkipLocalitiesCheck       = "skip_localities_check"
@@ -192,6 +193,22 @@ func allocateDescriptorRewrites(
 			}
 		}
 
+		// Check that functions referenced in check constraints exist.
+		fnIDs, err := table.GetAllReferencedFunctionIDs()
+		if err != nil {
+			return nil, err
+		}
+		for _, fnID := range fnIDs.Ordered() {
+			if _, ok := functionsByID[fnID]; !ok {
+				if !opts.SkipMissingUDFs {
+					return nil, errors.Errorf(
+						"cannot restore table %q without referenced function %d (or %q option)",
+						table.Name, fnID, restoreOptSkipMissingUDFs,
+					)
+				}
+			}
+		}
+
 		// Check that referenced sequences exist.
 		for i := range table.Columns {
 			col := &table.Columns[i]
@@ -869,6 +886,7 @@ func resolveOptionsForRestoreJobDescription(
 		SkipMissingSequences:      opts.SkipMissingSequences,
 		SkipMissingSequenceOwners: opts.SkipMissingSequenceOwners,
 		SkipMissingViews:          opts.SkipMissingViews,
+		SkipMissingUDFs:           opts.SkipMissingUDFs,
 		Detached:                  opts.Detached,
 		SchemaOnly:                opts.SchemaOnly,
 		VerifyData:                opts.VerifyData,

pkg/ccl/backupccl/testdata/backup-restore/user-defined-functions-in-checks

@@ -0,0 +1,228 @@
+# Test backing up and restoring a database with user defined functions.
+new-cluster name=s
+----
+
+exec-sql
+CREATE DATABASE db1;
+USE db1;
+CREATE SCHEMA sc1;
+CREATE FUNCTION sc1.f1(a INT) RETURNS INT LANGUAGE SQL AS $$
+  SELECT a + 1;
+$$;
+CREATE TABLE sc1.t1(a INT PRIMARY KEY, b INT CHECK(sc1.f1(b) > 1));
+----
+
+exec-sql
+BACKUP DATABASE db1 INTO 'nodelocal://0/test/'
+----
+
+query-sql
+WITH descs AS (
+  SHOW BACKUP LATEST IN 'nodelocal://0/test/'
+)
+SELECT database_name, parent_schema_name, object_name, object_type, is_full_cluster FROM descs
+----
+<nil> <nil> db1 database false
+db1 <nil> public schema false
+db1 <nil> sc1 schema false
+db1 sc1 f1 function false
+db1 sc1 t1 table false
+
+exec-sql
+RESTORE DATABASE db1 FROM LATEST IN 'nodelocal://0/test/' WITH new_db_name = db1_new
+----
+
+exec-sql
+USE db1_new
+----
+
+# Make sure function ids in CHECK constraint are rewritten.
+query-sql
+SELECT create_statement FROM [SHOW CREATE TABLE sc1.t1]
+----
+CREATE TABLE sc1.t1 (
+	a INT8 NOT NULL,
+	b INT8 NULL,
+	CONSTRAINT t1_pkey PRIMARY KEY (a ASC),
+	CONSTRAINT check_b CHECK (sc1.f1(b) > 1:::INT8)
+)
+
+# Make sure that the CHECK constraint still applies correctly
+query-sql
+INSERT INTO sc1.t1 VALUES (1, 0)
+----
+pq: failed to satisfy CHECK constraint (sc1.f1(b) > 1:::INT8)
+
+# Make sure dependency IDs are rewritten.
+# Note that technically this only tests forward-reference IDs in depended-on
+# objects are rewritten. But since we have cross-references validation, so this
+# also means back-references in UDF descriptor are good.
+exec-sql
+DROP FUNCTION sc1.f1
+----
+pq: cannot drop function "f1" because other objects ([db1_new.sc1.t1]) still depend on it
+
+# Test backing up and restoring a full cluster with user defined function.
+new-cluster name=s1
+----
+
+exec-sql cluster=s1
+CREATE DATABASE db1;
+USE db1;
+CREATE SCHEMA sc1;
+CREATE FUNCTION sc1.f1(a INT) RETURNS INT LANGUAGE SQL AS $$
+  SELECT a + 1;
+$$;
+CREATE TABLE sc1.t1(a INT PRIMARY KEY, b INT CHECK(sc1.f1(b) > 1));
+----
+
+exec-sql
+BACKUP INTO 'nodelocal://0/test/'
+----
+
+query-sql
+WITH descs AS (
+  SHOW BACKUP LATEST IN 'nodelocal://0/test/'
+)
+SELECT
+  database_name, parent_schema_name, object_name, object_type, is_full_cluster
+FROM
+  descs
+WHERE
+  database_name = 'db1'
+----
+db1 <nil> public schema true
+db1 <nil> sc1 schema true
+db1 sc1 f1 function true
+db1 sc1 t1 table true
+
+# Start a new cluster with the same IO dir.
+new-cluster name=s2 share-io-dir=s1
+----
+
+# Restore into the new cluster.
+exec-sql cluster=s2
+RESTORE FROM LATEST IN 'nodelocal://0/test/'
+----
+
+exec-sql
+USE db1
+----
+
+# Make sure function ids in CHECK constraint are rewritten.
+query-sql
+SELECT create_statement FROM [SHOW CREATE TABLE sc1.f1]
+----
+pq: relation "sc1.f1" does not exist
+
+# Make sure that CHECK constraint still applies correctly
+query-sql
+INSERT INTO sc1.t1 VALUES (1, 0)
+----
+pq: failed to satisfy CHECK constraint (sc1.f1(b) > 1:::INT8)
+
+# Make sure dependency IDs are rewritten.
+# Note that technically this only tests forward-reference IDs in depended-on
+# objects are rewritten. But since we have cross-references validation, so this
+# also means back-references in UDF descriptor are good.
+exec-sql
+DROP FUNCTION sc1.f1
+----
+pq: cannot drop function "f1" because other objects ([db1.sc1.t1]) still depend on it
+
+# Make sure that backup and restore individual tables referencing UDFs able to
+# drop check constraints.
+new-cluster name=s3
+----
+
+exec-sql cluster=s3
+CREATE DATABASE db1;
+CREATE DATABASE db2;
+CREATE DATABASE db3;
+USE db1;
+CREATE SCHEMA sc1;
+CREATE FUNCTION sc1.f1(a INT) RETURNS INT LANGUAGE SQL AS $$
+  SELECT a + 1;
+$$;
+CREATE TABLE sc1.t1(a INT PRIMARY KEY, b INT CHECK(sc1.f1(b) > 1));
+----
+
+exec-sql
+BACKUP DATABASE db1 INTO 'nodelocal://0/test/'
+----
+
+query-sql
+WITH descs AS (
+  SHOW BACKUP LATEST IN 'nodelocal://0/test/'
+)
+SELECT database_name, parent_schema_name, object_name, object_type, is_full_cluster FROM descs
+----
+<nil> <nil> db1 database false
+db1 <nil> public schema false
+db1 <nil> sc1 schema false
+db1 sc1 f1 function false
+db1 sc1 t1 table false
+
+exec-sql
+RESTORE TABLE sc1.t1 FROM LATEST IN 'nodelocal://0/test/' WITH into_db = 'db2';
+----
+pq: cannot restore table "t1" without referenced function 114 (or "skip_missing_udfs" option)
+
+exec-sql
+RESTORE TABLE sc1.t1 FROM LATEST IN 'nodelocal://0/test/' WITH into_db = 'db2', skip_missing_udfs;
+----
+
+exec-sql
+USE db2
+----
+
+# Make sure CHECK constraint is dropped.
+query-sql
+SELECT create_statement FROM [SHOW CREATE TABLE sc1.t1]
+----
+CREATE TABLE sc1.t1 (
+	a INT8 NOT NULL,
+	b INT8 NULL,
+	CONSTRAINT t1_pkey PRIMARY KEY (a ASC)
+)
+
+exec-sql
+USE db1
+----
+
+exec-sql
+BACKUP TABLE sc1.t1 INTO 'nodelocal://0/test/'
+----
+
+query-sql
+WITH descs AS (
+  SHOW BACKUP LATEST IN 'nodelocal://0/test/'
+)
+SELECT database_name, parent_schema_name, object_name, object_type, is_full_cluster FROM descs
+----
+<nil> <nil> db1 database false
+db1 <nil> sc1 schema false
+db1 sc1 t1 table false
+
+exec-sql
+RESTORE TABLE sc1.t1 FROM LATEST IN 'nodelocal://0/test/' WITH into_db = 'db3';
+----
+pq: cannot restore table "t1" without referenced function 114 (or "skip_missing_udfs" option)
+
+exec-sql
+RESTORE TABLE sc1.t1 FROM LATEST IN 'nodelocal://0/test/' WITH into_db = 'db3', skip_missing_udfs;
+----
+
+exec-sql
+USE db3
+----
+
+# Make sure CHECK constraint is dropped.
+query-sql
+SELECT create_statement FROM [SHOW CREATE TABLE sc1.t1]
+----
+CREATE TABLE sc1.t1 (
+	a INT8 NOT NULL,
+	b INT8 NULL,
+	CONSTRAINT t1_pkey PRIMARY KEY (a ASC)
+)

pkg/ccl/schemachangerccl/testdata/decomp/multiregion

@@ -278,6 +278,7 @@ ElementState:
     expr: unique_rowid()
     referencedColumnIds: []
     tableId: 110
+    usesFunctionIds: []
     usesSequenceIds: []
     usesTypeIds: []
   Status: PUBLIC
@@ -555,6 +556,7 @@ ElementState:
     expr: unique_rowid()
     referencedColumnIds: []
     tableId: 109
+    usesFunctionIds: []
     usesSequenceIds: []
     usesTypeIds: []
   Status: PUBLIC
@@ -883,6 +885,7 @@ ElementState:
     expr: default_to_database_primary_region(gateway_region())::@100106
     referencedColumnIds: []
     tableId: 108
+    usesFunctionIds: []
     usesSequenceIds: []
     usesTypeIds:
     - 106

pkg/sql/BUILD.bazel

@@ -118,6 +118,7 @@ go_library(
         "explain_vec.go",
         "export.go",
         "filter.go",
+        "function_references.go",
         "generate_objects.go",
         "gossip.go",
         "grant_revoke.go",

pkg/sql/add_column.go

@@ -139,7 +139,8 @@ func (p *planner) addColumnImpl(
 
 	if d.IsComputed() {
 		serializedExpr, _, err := schemaexpr.ValidateComputedColumnExpression(
-			params.ctx, n.tableDesc, d, tn, "computed column", params.p.SemaCtx(),
+			params.ctx, n.tableDesc, d, tn, tree.ComputedColumnExprContext(d.IsVirtual()), params.p.SemaCtx(),
+			params.ExecCfg().Settings.Version.ActiveVersion(params.ctx),
 		)
 		if err != nil {
 			return err