sql/schemachanger: add support for setting schema of enum types

pkg/sql/alter_type.go

@@ -355,8 +355,8 @@ func (p *planner) setTypeSchema(ctx context.Context, n *alterTypeNode, schema st
 		return err
 	}
 
-	// The prepareSetSchema checks the primary type's name. The name of the
-	// companion array is collision checked below.
+	// The prepareSetSchema checks the primary type's name while the name of the
+	// companion array is collision-checked below.
 	desiredSchemaID, err := p.prepareSetSchema(ctx, n.prefix.Database, typeDesc, schema)
 	if err != nil {
 		return err
@@ -373,12 +373,12 @@ func (p *planner) setTypeSchema(ctx context.Context, n *alterTypeNode, schema st
 		return err
 	}
 
-	// The CheckObjectNameCollision checks that the companion array can be moved
-	// without colliding with
+	// CheckObjectNameCollision checks that the companion array can be moved
+	// without a name collision.
 	//
 	// This is consistent with the PG behavior which itself is inconsistent:
-	// `SET SCHEMA` errors on collision while `ALTER TYPE ... RENAME` auto-resolves
-	// conflicts on companion array names.
+	// `SET SCHEMA` errors on collision while `ALTER TYPE ... RENAME`
+	// auto-resolves conflicts on companion array names.
 	if err := descs.CheckObjectNameCollision(
 		ctx,
 		p.Descriptors(),

pkg/sql/logictest/testdata/logic_test/alter_type

@@ -928,3 +928,104 @@ t_rename_computed  CREATE TABLE public.t_rename_computed (
                    );
 
 subtest end
+
+subtest alter_type_set_schema
+
+statement ok
+CREATE SCHEMA s1;
+CREATE SCHEMA s2;
+CREATE TYPE sc_type AS ENUM ('a', 'b', 'c');
+
+# Move type from public to s1.
+statement ok
+ALTER TYPE sc_type SET SCHEMA s1
+
+# The type should be resolvable in the new schema.
+query T
+SELECT 'a'::s1.sc_type
+----
+a
+
+# The array type should have moved too.
+query T
+SELECT ARRAY['a']::s1._sc_type
+----
+{a}
+
+# Setting the same schema should be a no-op.
+statement ok
+ALTER TYPE s1.sc_type SET SCHEMA s1
+
+# Move from s1 to s2.
+statement ok
+ALTER TYPE s1.sc_type SET SCHEMA s2
+
+query T
+SELECT 'b'::s2.sc_type
+----
+b
+
+query T
+SELECT ARRAY['b']::s2._sc_type
+----
+{b}
+
+# The type should no longer exist in s1.
+query error pq: type "s1.sc_type" does not exist
+SELECT 'a'::s1.sc_type
+
+# Error when target schema doesn't exist.
+statement error pq: unknown schema "does_not_exist"
+ALTER TYPE s2.sc_type SET SCHEMA does_not_exist
+
+# Error when there's a name conflict in the target schema.
+statement ok
+CREATE TYPE s1.sc_type AS ENUM ('x', 'y')
+
+statement error pq: type "test.s1.sc_type" already exists
+ALTER TYPE s2.sc_type SET SCHEMA s1
+
+# Clean up.
+statement ok
+DROP TYPE s1.sc_type;
+DROP TYPE s2.sc_type;
+DROP SCHEMA s1;
+DROP SCHEMA s2
+
+subtest end
+
+subtest alter_type_set_schema_with_table_ref
+
+# Moving a type that is referenced by a table column should succeed
+# because types are referenced by ID, not name.
+statement ok
+CREATE SCHEMA s_ref;
+CREATE TYPE ref_type AS ENUM ('x', 'y', 'z');
+CREATE TABLE t_ref (c ref_type);
+INSERT INTO t_ref VALUES ('x')
+
+statement ok
+ALTER TYPE ref_type SET SCHEMA s_ref
+
+# The table should still work with the moved type.
+query T
+SELECT c FROM t_ref
+----
+x
+
+statement ok
+INSERT INTO t_ref VALUES ('y')
+
+query T rowsort
+SELECT c FROM t_ref
+----
+x
+y
+
+# Clean up.
+statement ok
+DROP TABLE t_ref;
+DROP TYPE s_ref.ref_type;
+DROP SCHEMA s_ref
+
+subtest end

pkg/sql/logictest/testdata/logic_test/event_log_legacy

@@ -937,6 +937,9 @@ DROP USER u
 # event log.
 subtest regression_57734
 
+statement ok
+USE defaultdb
+
 statement ok
 CREATE TYPE eventlog AS ENUM ('event', 'log')
 

pkg/sql/logictest/testdata/logic_test/set_schema

@@ -28,8 +28,7 @@ statement error pq: relation "test.s2.t" already exists
 ALTER TABLE t SET SCHEMA s2
 
 # Ensure we cannot set schema to a virtual schema.
-# use regex to accept both the legacy and the declarative schema changer format
-statement error pq: (cannot move objects into or out of virtual schemas)?(user root does not have CREATE privilege on schema information_schema)?
+statement error pq: cannot move objects into or out of virtual schemas
 ALTER TABLE t SET SCHEMA information_schema
 
 # Ensure we cannot set schema for a table in a virtual schema.

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

@@ -6,10 +6,13 @@
 package scbuildstmt
 
 import (
+	"strings"
+
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/catconstants"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/catid"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/util/log/eventpb"
@@ -49,10 +52,10 @@ func AlterTableSetSchema(b BuildCtx, n *tree.AlterTableSetSchema) {
 	// Get the schema ID directly from the namespace element
 	currNamespace := mustRetrieveNamespaceElem(b, descID)
 	currSchemaID := currNamespace.SchemaID
+	// Ensure that new schema is neither temporary nor virtual.
+	panicIfSchemaIsTemporaryOrVirtual(n.Schema)
 	newSchema := resolveSchemaByName(b, n.Schema, currNamespace.DatabaseID)
 	newSchemaID := newSchema.SchemaID
-	// Ensure that new schema is not temporary or virtual
-	panicIfSchemaIsTemporaryOrVirtual(newSchema)
 	// If new schema is the same as the curr schema, do a no-op
 	if currSchemaID == newSchemaID {
 		return
@@ -63,20 +66,7 @@ func AlterTableSetSchema(b BuildCtx, n *tree.AlterTableSetSchema) {
 	// Check for name conflicts
 	checkTableNameConflicts(b, currName, newName, currNamespace)
 
-	// drop the old namespace and add a new one
-	newNamespace := *currNamespace
-	newNamespace.SchemaID = newSchemaID
-	b.Drop(currNamespace)
-	b.Add(&newNamespace)
-
-	// drop old schema child and add new one
-	currSchemaChild := b.QueryByID(descID).FilterSchemaChild().MustGetOneElement()
-	newSchemaChild := scpb.SchemaChild{
-		ChildObjectID: descID,
-		SchemaID:      newSchemaID,
-	}
-	b.Drop(currSchemaChild)
-	b.Add(&newSchemaChild)
+	newNamespace, _ := moveDescriptorToSchema(b, descID, currNamespace, newSchemaID)
 
 	// Log event for audit logging.
 	kind := tree.GetTableType(n.IsSequence, n.IsView, n.IsMaterialized)
@@ -85,7 +75,7 @@ func AlterTableSetSchema(b BuildCtx, n *tree.AlterTableSetSchema) {
 		NewDescriptorName: newName.FQString(),
 		DescriptorType:    kind,
 	}
-	b.LogEventForExistingPayload(&newNamespace, setSchemaEvent)
+	b.LogEventForExistingPayload(newNamespace, setSchemaEvent)
 }
 
 func resolveSchemaByName(b BuildCtx, schemaName tree.Name, databaseID catid.DescID) *scpb.Schema {
@@ -104,13 +94,38 @@ func resolveSchemaByName(b BuildCtx, schemaName tree.Name, databaseID catid.Desc
 	return newSchema
 }
 
-func panicIfSchemaIsTemporaryOrVirtual(newSchema *scpb.Schema) {
-	if newSchema.IsTemporary {
+func panicIfSchemaIsTemporaryOrVirtual(schemaName tree.Name) {
+	name := string(schemaName)
+	if _, ok := catconstants.VirtualSchemaNames[name]; ok {
 		panic(pgerror.Newf(pgcode.FeatureNotSupported,
-			"cannot move objects into or out of temporary schemas"))
+			"cannot move objects into or out of virtual schemas"))
 	}
-	if newSchema.IsVirtual {
+	if strings.HasPrefix(name, catconstants.PgTempSchemaName) {
 		panic(pgerror.Newf(pgcode.FeatureNotSupported,
-			"cannot move objects into or out of virtual schemas"))
+			"cannot move objects into or out of temporary schemas"))
+	}
+}
+
+// moveDescriptorToSchema drops the old Namespace and SchemaChild elements for
+// the given descriptor and adds new ones with the updated schema ID.
+// It returns the new Namespace and SchemaChild elements.
+func moveDescriptorToSchema(
+	b BuildCtx, descID catid.DescID, currNamespace *scpb.Namespace, newSchemaID catid.DescID,
+) (*scpb.Namespace, *scpb.SchemaChild) {
+	// drop the old namespace and add a new one
+	newNamespace := *currNamespace
+	newNamespace.SchemaID = newSchemaID
+	b.Drop(currNamespace)
+	b.Add(&newNamespace)
+
+	// drop old schema child and add new one
+	currSchemaChild := b.QueryByID(descID).FilterSchemaChild().MustGetOneElement()
+	newSchemaChild := &scpb.SchemaChild{
+		ChildObjectID: descID,
+		SchemaID:      newSchemaID,
 	}
+	b.Drop(currSchemaChild)
+	b.Add(newSchemaChild)
+
+	return &newNamespace, newSchemaChild
 }

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

@@ -37,7 +37,7 @@ var supportedAlterTypeStatements = map[reflect.Type]supportedAlterTypeCommand{
 	reflect.TypeOf((*tree.AlterTypeAddValue)(nil)):    {fn: alterTypeAddValue, on: true, checks: isV263Active},
 	reflect.TypeOf((*tree.AlterTypeRenameValue)(nil)): {fn: alterTypeRenameValue, on: true, checks: isV263Active},
 	reflect.TypeOf((*tree.AlterTypeRename)(nil)):      {fn: alterTypeRename, on: true, checks: isV263Active},
-	reflect.TypeOf((*tree.AlterTypeSetSchema)(nil)):   {fn: alterTypeSetSchema, on: false, checks: nil},
+	reflect.TypeOf((*tree.AlterTypeSetSchema)(nil)):   {fn: alterTypeSetSchema, on: true, checks: isV263Active},
 	reflect.TypeOf((*tree.AlterTypeOwner)(nil)):       {fn: alterTypeOwner, on: true, checks: isV263Active},
 	reflect.TypeOf((*tree.AlterTypeDropValue)(nil)):   {fn: alterTypeDropValue, on: true, checks: isV263Active},
 }
@@ -280,7 +280,40 @@ func alterTypeRenameValue(
 func alterTypeSetSchema(
 	b BuildCtx, tn *tree.TypeName, enumType *scpb.EnumType, t *tree.AlterTypeSetSchema,
 ) {
-	panic(scerrors.NotImplementedErrorf(t, "ALTER TYPE SET SCHEMA is not supported"))
+	typeID := enumType.TypeID
+
+	currNamespace := mustRetrieveNamespaceElem(b, typeID)
+	currSchemaID := currNamespace.SchemaID
+	panicIfSchemaIsTemporaryOrVirtual(t.Schema)
+	newSchema := resolveSchemaByName(b, t.Schema, currNamespace.DatabaseID)
+	newSchemaID := newSchema.SchemaID
+
+	if currSchemaID == newSchemaID {
+		return
+	}
+
+	currName := tree.MakeTableNameFromPrefix(b.NamePrefix(enumType), tree.Name(currNamespace.Name))
+	newName := currName
+	newName.SchemaName = t.Schema
+
+	checkTableNameConflicts(b, currName, newName, currNamespace)
+
+	arrayNamespace := mustRetrieveNamespaceElem(b, enumType.ArrayTypeID)
+	arrayName := tree.MakeTableNameFromPrefix(
+		b.NamePrefix(enumType), tree.Name(arrayNamespace.Name),
+	)
+	newArrayName := arrayName
+	newArrayName.SchemaName = t.Schema
+	checkTableNameConflicts(b, arrayName, newArrayName, arrayNamespace)
+
+	newNS, _ := moveDescriptorToSchema(b, typeID, currNamespace, newSchemaID)
+	moveDescriptorToSchema(b, enumType.ArrayTypeID, arrayNamespace, newSchemaID)
+
+	b.LogEventForExistingPayload(newNS, &eventpb.SetSchema{
+		DescriptorName:    currName.FQString(),
+		NewDescriptorName: newName.FQString(),
+		DescriptorType:    "type",
+	})
 }
 
 func alterTypeOwner(

pkg/sql/schemachanger/scbuild/testdata/alter_type

@@ -33,3 +33,29 @@ ALTER TYPE defaultdb.climes RENAME TO environs
   {databaseId: 100, descriptorId: 105, name: _climes, schemaId: 101}
 - [[Namespace:{DescID: 105, Name: _environs, ReferencedDescID: 100, IntValue: 101}, PUBLIC], ABSENT]
   {databaseId: 100, descriptorId: 105, name: _environs, schemaId: 101}
+ALTER TYPE defaultdb.climes SET SCHEMA public
+----
+
+setup
+CREATE SCHEMA s;
+----
+
+build
+ALTER TYPE defaultdb.climes SET SCHEMA s
+----
+- [[Namespace:{DescID: 104, Name: climes, ReferencedDescID: 100, IntValue: 101}, ABSENT], PUBLIC]
+  {databaseId: 100, descriptorId: 104, name: climes, schemaId: 101}
+- [[SchemaChild:{DescID: 104, ReferencedDescID: 101}, ABSENT], PUBLIC]
+  {childObjectId: 104, schemaId: 101}
+- [[Namespace:{DescID: 105, Name: _climes, ReferencedDescID: 100, IntValue: 101}, ABSENT], PUBLIC]
+  {databaseId: 100, descriptorId: 105, name: _climes, schemaId: 101}
+- [[SchemaChild:{DescID: 105, ReferencedDescID: 101}, ABSENT], PUBLIC]
+  {childObjectId: 105, schemaId: 101}
+- [[Namespace:{DescID: 104, Name: climes, ReferencedDescID: 100, IntValue: 106}, PUBLIC], ABSENT]
+  {databaseId: 100, descriptorId: 104, name: climes, schemaId: 106}
+- [[SchemaChild:{DescID: 104, ReferencedDescID: 106}, PUBLIC], ABSENT]
+  {childObjectId: 104, schemaId: 106}
+- [[Namespace:{DescID: 105, Name: _climes, ReferencedDescID: 100, IntValue: 106}, PUBLIC], ABSENT]
+  {databaseId: 100, descriptorId: 105, name: _climes, schemaId: 106}
+- [[SchemaChild:{DescID: 105, ReferencedDescID: 106}, PUBLIC], ABSENT]
+  {childObjectId: 105, schemaId: 106}