From 4e01fb11deb3dd11b12da69e2b6458f2925ca94b Mon Sep 17 00:00:00 2001
From: Chengxiong Ruan <chengxiongruan@gmail.com>
Date: Wed, 17 Aug 2022 17:59:21 -0400
Subject: [PATCH] sql: add feature flag checking for UDF statements

Release note: None
Release justification: low risk feature flags for DDL statements.
---
 pkg/sql/alter_function.go                     | 32 +++++++++++++++++++
 pkg/sql/drop_function.go                      |  8 +++++
 .../logic_test/schema_change_feature_flags    | 27 ++++++++++++++++
 3 files changed, 67 insertions(+)

diff --git a/pkg/sql/alter_function.go b/pkg/sql/alter_function.go
index d452b2354388..4a78414ccc32 100644
--- a/pkg/sql/alter_function.go
+++ b/pkg/sql/alter_function.go
@@ -52,6 +52,14 @@ type alterFunctionDepExtensionNode struct {
 func (p *planner) AlterFunctionOptions(
 	ctx context.Context, n *tree.AlterFunctionOptions,
 ) (planNode, error) {
+	if err := checkSchemaChangeEnabled(
+		ctx,
+		p.ExecCfg(),
+		"ALTER FUNCTION",
+	); err != nil {
+		return nil, err
+	}
+
 	return &alterFunctionOptionsNode{n: n}, nil
 }
 
@@ -90,6 +98,14 @@ func (n *alterFunctionOptionsNode) Close(ctx context.Context)           {}
 func (p *planner) AlterFunctionRename(
 	ctx context.Context, n *tree.AlterFunctionRename,
 ) (planNode, error) {
+	if err := checkSchemaChangeEnabled(
+		ctx,
+		p.ExecCfg(),
+		"ALTER FUNCTION",
+	); err != nil {
+		return nil, err
+	}
+
 	return &alterFunctionRenameNode{n: n}, nil
 }
 
@@ -141,6 +157,14 @@ func (n *alterFunctionRenameNode) Close(ctx context.Context)           {}
 func (p *planner) AlterFunctionSetOwner(
 	ctx context.Context, n *tree.AlterFunctionSetOwner,
 ) (planNode, error) {
+	if err := checkSchemaChangeEnabled(
+		ctx,
+		p.ExecCfg(),
+		"ALTER FUNCTION",
+	); err != nil {
+		return nil, err
+	}
+
 	return &alterFunctionSetOwnerNode{n: n}, nil
 }
 
@@ -181,6 +205,14 @@ func (n *alterFunctionSetOwnerNode) Close(ctx context.Context)           {}
 func (p *planner) AlterFunctionSetSchema(
 	ctx context.Context, n *tree.AlterFunctionSetSchema,
 ) (planNode, error) {
+	if err := checkSchemaChangeEnabled(
+		ctx,
+		p.ExecCfg(),
+		"ALTER FUNCTION",
+	); err != nil {
+		return nil, err
+	}
+
 	return &alterFunctionSetSchemaNode{n: n}, nil
 }
 
diff --git a/pkg/sql/drop_function.go b/pkg/sql/drop_function.go
index 8a12851bcce8..3f726044fb5d 100644
--- a/pkg/sql/drop_function.go
+++ b/pkg/sql/drop_function.go
@@ -34,6 +34,14 @@ type dropFunctionNode struct {
 func (p *planner) DropFunction(
 	ctx context.Context, n *tree.DropFunction,
 ) (ret planNode, err error) {
+	if err := checkSchemaChangeEnabled(
+		ctx,
+		p.ExecCfg(),
+		"DROP FUNCTION",
+	); err != nil {
+		return nil, err
+	}
+
 	if n.DropBehavior == tree.DropCascade {
 		// TODO(chengxiong): remove this check when drop function cascade is supported.
 		return nil, unimplemented.Newf("DROP FUNCTION...CASCADE", "drop function cascade not supported")
diff --git a/pkg/sql/logictest/testdata/logic_test/schema_change_feature_flags b/pkg/sql/logictest/testdata/logic_test/schema_change_feature_flags
index 852921f89f13..567b04cd4c90 100644
--- a/pkg/sql/logictest/testdata/logic_test/schema_change_feature_flags
+++ b/pkg/sql/logictest/testdata/logic_test/schema_change_feature_flags
@@ -310,3 +310,30 @@ COMMENT ON TABLE t IS 'comment'
 # Reset feature flag to true so that test objects can be dropped.
 statement ok
 SET CLUSTER SETTING feature.schema_change.enabled = TRUE
+
+statement ok
+SET CLUSTER SETTING feature.schema_change.enabled = FALSE
+
+statement error pq: feature CREATE FUNCTION is part of the schema change category, which was disabled by the database administrator
+CREATE FUNCTION f() RETURNS INT LANGUAGE SQL AS $$ SELECT 1 $$;
+
+statement error pq: feature CREATE FUNCTION is part of the schema change category, which was disabled by the database administrator
+CREATE OR REPLACE FUNCTION f() RETURNS INT LANGUAGE SQL AS $$ SELECT 1 $$;
+
+statement error pq: feature DROP FUNCTION is part of the schema change category, which was disabled by the database administrator
+DROP FUNCTION f()
+
+statement error pq: feature ALTER FUNCTION is part of the schema change category, which was disabled by the database administrator
+ALTER FUNCTION f() IMMUTABLE;
+
+statement error pq: feature ALTER FUNCTION is part of the schema change category, which was disabled by the database administrator
+ALTER FUNCTION f() RENAME TO g;
+
+statement error pq: feature ALTER FUNCTION is part of the schema change category, which was disabled by the database administrator
+ALTER FUNCTION f() OWNER TO new_owner;
+
+statement error pq: feature ALTER FUNCTION is part of the schema change category, which was disabled by the database administrator
+ALTER FUNCTION f() SET SCHEMA new_schema;
+
+statement ok
+SET CLUSTER SETTING feature.schema_change.enabled = TRUE