release-23.2: sql: fix statement tags for CREATE and DROP PROCEDURE

pkg/sql/logictest/testdata/logic_test/routine_schema_change

@@ -0,0 +1,56 @@
+# LogicTest: !local-mixed-23.1
+
+subtest statement-control
+
+statement ok
+SET use_declarative_schema_changer = 'on'
+
+# Disable CREATE FUNCTION.
+statement ok
+SET CLUSTER SETTING sql.schema.force_declarative_statements='!CREATE FUNCTION'
+
+# Verify that it is disabled.
+statement error pgcode 0A000 cannot explain a statement which is not supported by the declarative schema changer
+EXPLAIN (DDL) CREATE FUNCTION f() RETURNS INT LANGUAGE SQL AS 'SELECT 1'
+
+# Falling-back to the legacy schema changer works.
+statement ok
+CREATE FUNCTION f() RETURNS INT LANGUAGE SQL AS 'SELECT 1'
+
+# Disable CREATE PROCEDURE.
+statement ok
+SET CLUSTER SETTING sql.schema.force_declarative_statements='!CREATE PROCEDURE'
+
+# Verify that it is disabled.
+statement error pgcode 0A000 cannot explain a statement which is not supported by the declarative schema changer
+EXPLAIN (DDL) CREATE PROCEDURE p() LANGUAGE SQL as 'SELECT 1'
+
+# Falling-back to the legacy schema changer works.
+statement ok
+CREATE PROCEDURE p() LANGUAGE SQL as 'SELECT 1'
+
+# Disable DROP FUNCTION.
+statement ok
+SET CLUSTER SETTING sql.schema.force_declarative_statements='!DROP FUNCTION'
+
+# Verify that it is disabled.
+statement error pgcode 0A000 cannot explain a statement which is not supported by the declarative schema changer
+EXPLAIN (DDL) DROP FUNCTION f
+
+# Falling-back to the legacy schema changer works.
+statement ok
+DROP FUNCTION f
+
+# Disable DROP PROCEDURE.
+statement ok
+SET CLUSTER SETTING sql.schema.force_declarative_statements='!DROP PROCEDURE'
+
+# Verify that it is disabled.
+statement error pgcode 0A000 cannot explain a statement which is not supported by the declarative schema changer
+EXPLAIN (DDL) DROP PROCEDURE p
+
+# Falling-back to the legacy schema changer works.
+statement ok
+DROP PROCEDURE p
+
+subtest end

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/BUILD.bazel

@@ -88,6 +88,7 @@ go_test(
     embed = [":scbuildstmt"],
     deps = [
         "//pkg/settings",
+        "//pkg/sql/sem/tree",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/process.go

@@ -25,8 +25,8 @@ import (
 type supportedStatement struct {
 	// fn is a function to perform a schema change.
 	fn interface{}
-	// statementTag tag for this statement.
-	statementTag string
+	// statementTags contains tags for this statement.
+	statementTags []string
 	// checks contains a coarse-grained function to filter out most
 	// unsupported statements.
 	// It's possible for certain unsupported statements to pass it but will
@@ -46,26 +46,26 @@ var supportedStatements = map[reflect.Type]supportedStatement{
 	// Alter table will have commands individually whitelisted via the
 	// supportedAlterTableStatements list, so wwe will consider it fully supported
 	// here.
-	reflect.TypeOf((*tree.AlterTable)(nil)):          {fn: AlterTable, statementTag: tree.AlterTableTag, on: true, checks: alterTableChecks},
-	reflect.TypeOf((*tree.CreateIndex)(nil)):         {fn: CreateIndex, statementTag: tree.CreateIndexTag, on: true, checks: isV231Active},
-	reflect.TypeOf((*tree.DropDatabase)(nil)):        {fn: DropDatabase, statementTag: tree.DropDatabaseTag, on: true, checks: nil},
-	reflect.TypeOf((*tree.DropOwnedBy)(nil)):         {fn: DropOwnedBy, statementTag: tree.DropOwnedByTag, on: true, checks: isV222Active},
-	reflect.TypeOf((*tree.DropSchema)(nil)):          {fn: DropSchema, statementTag: tree.DropSchemaTag, on: true, checks: nil},
-	reflect.TypeOf((*tree.DropSequence)(nil)):        {fn: DropSequence, statementTag: tree.DropSequenceTag, on: true, checks: nil},
-	reflect.TypeOf((*tree.DropTable)(nil)):           {fn: DropTable, statementTag: tree.DropTableTag, on: true, checks: nil},
-	reflect.TypeOf((*tree.DropType)(nil)):            {fn: DropType, statementTag: tree.DropTypeTag, on: true, checks: nil},
-	reflect.TypeOf((*tree.DropView)(nil)):            {fn: DropView, statementTag: tree.DropViewTag, on: true, checks: nil},
-	reflect.TypeOf((*tree.CommentOnConstraint)(nil)): {fn: CommentOnConstraint, statementTag: tree.CommentOnConstraintTag, on: true, checks: isV222Active},
-	reflect.TypeOf((*tree.CommentOnDatabase)(nil)):   {fn: CommentOnDatabase, statementTag: tree.CommentOnDatabaseTag, on: true, checks: isV222Active},
-	reflect.TypeOf((*tree.CommentOnSchema)(nil)):     {fn: CommentOnSchema, statementTag: tree.CommentOnSchemaTag, on: true, checks: isV222Active},
-	reflect.TypeOf((*tree.CommentOnTable)(nil)):      {fn: CommentOnTable, statementTag: tree.CommentOnTableTag, on: true, checks: isV222Active},
-	reflect.TypeOf((*tree.CommentOnColumn)(nil)):     {fn: CommentOnColumn, statementTag: tree.CommentOnColumnTag, on: true, checks: isV222Active},
-	reflect.TypeOf((*tree.CommentOnIndex)(nil)):      {fn: CommentOnIndex, statementTag: tree.CommentOnIndexTag, on: true, checks: isV222Active},
-	reflect.TypeOf((*tree.DropIndex)(nil)):           {fn: DropIndex, statementTag: tree.DropIndexTag, on: true, checks: isV231Active},
-	reflect.TypeOf((*tree.DropRoutine)(nil)):         {fn: DropFunction, statementTag: tree.DropFunctionTag, on: true, checks: isV231Active},
-	reflect.TypeOf((*tree.CreateRoutine)(nil)):       {fn: CreateFunction, statementTag: tree.CreateRoutineTag, on: true, checks: isV231Active},
-	reflect.TypeOf((*tree.CreateSchema)(nil)):        {fn: CreateSchema, statementTag: tree.CreateSchemaTag, on: false, checks: isV232Active},
-	reflect.TypeOf((*tree.CreateSequence)(nil)):      {fn: CreateSequence, statementTag: tree.CreateSequenceTag, on: false, checks: isV232Active},
+	reflect.TypeOf((*tree.AlterTable)(nil)):          {fn: AlterTable, statementTags: []string{tree.AlterTableTag}, on: true, checks: alterTableChecks},
+	reflect.TypeOf((*tree.CreateIndex)(nil)):         {fn: CreateIndex, statementTags: []string{tree.CreateIndexTag}, on: true, checks: isV231Active},
+	reflect.TypeOf((*tree.DropDatabase)(nil)):        {fn: DropDatabase, statementTags: []string{tree.DropDatabaseTag}, on: true, checks: nil},
+	reflect.TypeOf((*tree.DropOwnedBy)(nil)):         {fn: DropOwnedBy, statementTags: []string{tree.DropOwnedByTag}, on: true, checks: isV222Active},
+	reflect.TypeOf((*tree.DropSchema)(nil)):          {fn: DropSchema, statementTags: []string{tree.DropSchemaTag}, on: true, checks: nil},
+	reflect.TypeOf((*tree.DropSequence)(nil)):        {fn: DropSequence, statementTags: []string{tree.DropSequenceTag}, on: true, checks: nil},
+	reflect.TypeOf((*tree.DropTable)(nil)):           {fn: DropTable, statementTags: []string{tree.DropTableTag}, on: true, checks: nil},
+	reflect.TypeOf((*tree.DropType)(nil)):            {fn: DropType, statementTags: []string{tree.DropTypeTag}, on: true, checks: nil},
+	reflect.TypeOf((*tree.DropView)(nil)):            {fn: DropView, statementTags: []string{tree.DropViewTag}, on: true, checks: nil},
+	reflect.TypeOf((*tree.CommentOnConstraint)(nil)): {fn: CommentOnConstraint, statementTags: []string{tree.CommentOnConstraintTag}, on: true, checks: isV222Active},
+	reflect.TypeOf((*tree.CommentOnDatabase)(nil)):   {fn: CommentOnDatabase, statementTags: []string{tree.CommentOnDatabaseTag}, on: true, checks: isV222Active},
+	reflect.TypeOf((*tree.CommentOnSchema)(nil)):     {fn: CommentOnSchema, statementTags: []string{tree.CommentOnSchemaTag}, on: true, checks: isV222Active},
+	reflect.TypeOf((*tree.CommentOnTable)(nil)):      {fn: CommentOnTable, statementTags: []string{tree.CommentOnTableTag}, on: true, checks: isV222Active},
+	reflect.TypeOf((*tree.CommentOnColumn)(nil)):     {fn: CommentOnColumn, statementTags: []string{tree.CommentOnColumnTag}, on: true, checks: isV222Active},
+	reflect.TypeOf((*tree.CommentOnIndex)(nil)):      {fn: CommentOnIndex, statementTags: []string{tree.CommentOnIndexTag}, on: true, checks: isV222Active},
+	reflect.TypeOf((*tree.DropIndex)(nil)):           {fn: DropIndex, statementTags: []string{tree.DropIndexTag}, on: true, checks: isV231Active},
+	reflect.TypeOf((*tree.DropRoutine)(nil)):         {fn: DropFunction, statementTags: []string{tree.DropFunctionTag, tree.DropProcedureTag}, on: true, checks: isV231Active},
+	reflect.TypeOf((*tree.CreateRoutine)(nil)):       {fn: CreateFunction, statementTags: []string{tree.CreateFunctionTag, tree.CreateProcedureTag}, on: true, checks: isV231Active},
+	reflect.TypeOf((*tree.CreateSchema)(nil)):        {fn: CreateSchema, statementTags: []string{tree.CreateSchemaTag}, on: false, checks: isV232Active},
+	reflect.TypeOf((*tree.CreateSequence)(nil)):      {fn: CreateSequence, statementTags: []string{tree.CreateSequenceTag}, on: false, checks: isV232Active},
 }
 
 // supportedStatementTags tracks statement tags which are implemented
@@ -107,7 +107,9 @@ func init() {
 		}
 		// Fetch the statement tag using the statement tag method on the type,
 		// we can use this as a blacklist of blocked schema changes.
-		supportedStatementTags[statementEntry.statementTag] = struct{}{}
+		for _, tag := range statementEntry.statementTags {
+			supportedStatementTags[tag] = struct{}{}
+		}
 	}
 }
 

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/process_test.go

@@ -16,10 +16,19 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/settings"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/stretchr/testify/require"
 )
 
 func TestSupportedStatements(t *testing.T) {
+	// Some statements n supportedStatements have multiple statementTags. This
+	// map contains concrete statements that will return each tag, corresponding
+	// one-to-one with each tag in statementTags.
+	var multiTagStmts = map[reflect.Type][]tree.Statement{
+		reflect.TypeOf((*tree.DropRoutine)(nil)):   {&tree.DropRoutine{}, &tree.DropRoutine{Procedure: true}},
+		reflect.TypeOf((*tree.CreateRoutine)(nil)): {&tree.CreateRoutine{}, &tree.CreateRoutine{IsProcedure: true}},
+	}
+
 	sv := &settings.Values{}
 	// Non-existent tags should error out.
 	require.Error(t, forceDeclarativeStatements.Validate(sv, "FAKE STATEMENT"))
@@ -28,25 +37,36 @@ func TestSupportedStatements(t *testing.T) {
 	noTags := strings.Builder{}
 	first := true
 	for typ, stmt := range supportedStatements {
-		require.Greaterf(t, len(stmt.statementTag), 0, "statement tag is missing %v %v", typ, stmt)
-		// Validate tags matches the statement tag
-		typTag, found := typ.MethodByName("StatementTag")
-		require.True(t, found, "unable to find stmt: %v %v", typ, stmt)
-		ret := typTag.Func.Call([]reflect.Value{reflect.New(typ.Elem())})
-		require.Equal(t, ret[0].String(), stmt.statementTag, "statement tag is different in AST")
-		// Validate all tags are supported.
-		require.NoError(t, forceDeclarativeStatements.Validate(sv, "+"+stmt.statementTag))
-		require.NoError(t, forceDeclarativeStatements.Validate(sv, "!"+stmt.statementTag))
-		// Validate all of them can be specified at once.
-		if !first {
-			allTags.WriteString(",")
-			noTags.WriteString(",")
+		for i, tag := range stmt.statementTags {
+			require.Greaterf(t, len(stmt.statementTags), 0, "statement tag is missing %v %v", typ, stmt)
+
+			// Validate tags matches the statement tag.
+			var expectedTag string
+			if concreteStmts, ok := multiTagStmts[typ]; ok {
+				expectedTag = concreteStmts[i].StatementTag()
+			} else {
+				// Otherwise, build a zero-value statement.
+				typTag, found := typ.MethodByName("StatementTag")
+				require.True(t, found, "unable to find stmt: %v %v", typ, stmt)
+				ret := typTag.Func.Call([]reflect.Value{reflect.New(typ.Elem())})
+				expectedTag = ret[0].String()
+			}
+			require.Equal(t, expectedTag, tag, "statement tag is different in AST")
+
+			// Validate all tags are supported.
+			require.NoError(t, forceDeclarativeStatements.Validate(sv, "+"+tag))
+			require.NoError(t, forceDeclarativeStatements.Validate(sv, "!"+tag))
+			// Validate all of them can be specified at once.
+			if !first {
+				allTags.WriteString(",")
+				noTags.WriteString(",")
+			}
+			first = false
+			allTags.WriteString("+")
+			allTags.WriteString(tag)
+			noTags.WriteString("!")
+			noTags.WriteString(tag)
 		}
-		first = false
-		allTags.WriteString("+")
-		allTags.WriteString(stmt.statementTag)
-		noTags.WriteString("!")
-		noTags.WriteString(stmt.statementTag)
 	}
 	require.NoError(t, forceDeclarativeStatements.Validate(sv, allTags.String()))
 	require.NoError(t, forceDeclarativeStatements.Validate(sv, noTags.String()))

pkg/sql/sem/tree/stmt.go

@@ -93,7 +93,8 @@ const (
 	AlterTableTag          = "ALTER TABLE"
 	BackupTag              = "BACKUP"
 	CreateIndexTag         = "CREATE INDEX"
-	CreateRoutineTag       = "CREATE FUNCTION"
+	CreateFunctionTag      = "CREATE FUNCTION"
+	CreateProcedureTag     = "CREATE PROCEDURE"
 	CreateSchemaTag        = "CREATE SCHEMA"
 	CreateSequenceTag      = "CREATE SEQUENCE"
 	CommentOnColumnTag     = "COMMENT ON COLUMN"
@@ -104,6 +105,7 @@ const (
 	CommentOnTableTag      = "COMMENT ON TABLE"
 	DropDatabaseTag        = "DROP DATABASE"
 	DropFunctionTag        = "DROP FUNCTION"
+	DropProcedureTag       = "DROP PROCEDURE"
 	DropIndexTag           = "DROP INDEX"
 	DropOwnedByTag         = "DROP OWNED BY"
 	DropSchemaTag          = "DROP SCHEMA"
@@ -1957,9 +1959,8 @@ func (*ShowRoutines) StatementType() StatementType { return TypeDML }
 func (n *ShowRoutines) StatementTag() string {
 	if n.Procedure {
 		return "SHOW PROCEDURES"
-	} else {
-		return "SHOW FUNCTIONS"
 	}
+	return "SHOW FUNCTIONS"
 }
 
 // StatementReturnType implements the Statement interface
@@ -2121,7 +2122,12 @@ func (*CreateRoutine) StatementReturnType() StatementReturnType { return DDL }
 func (*CreateRoutine) StatementType() StatementType { return TypeDDL }
 
 // StatementTag returns a short string identifying the type of statement.
-func (*CreateRoutine) StatementTag() string { return CreateRoutineTag }
+func (n *CreateRoutine) StatementTag() string {
+	if n.IsProcedure {
+		return CreateProcedureTag
+	}
+	return CreateFunctionTag
+}
 
 // StatementReturnType implements the Statement interface.
 func (*RoutineReturn) StatementReturnType() StatementReturnType { return Rows }
@@ -2139,7 +2145,12 @@ func (*DropRoutine) StatementReturnType() StatementReturnType { return DDL }
 func (*DropRoutine) StatementType() StatementType { return TypeDDL }
 
 // StatementTag returns a short string identifying the type of statement.
-func (*DropRoutine) StatementTag() string { return DropFunctionTag }
+func (n *DropRoutine) StatementTag() string {
+	if n.Procedure {
+		return DropProcedureTag
+	}
+	return DropFunctionTag
+}
 
 // StatementReturnType implements the Statement interface.
 func (*AlterFunctionOptions) StatementReturnType() StatementReturnType { return DDL }