sql: validate against array type usages when dropping enum values

pkg/sql/logictest/testdata/logic_test/alter_type

@@ -496,3 +496,81 @@ SELECT ARRAY['c']::_ab_58710;
 
 statement error invalid input value for enum ab_58710: "a"
 SELECT ARRAY['a']::_ab_58710
+
+subtest regression_60004_basic
+
+statement ok
+CREATE TYPE enum_60004 AS ENUM ('a', 'b', 'c')
+
+statement ok
+CREATE TABLE t_60004 (k INT PRIMARY KEY, v enum_60004[])
+
+statement ok
+INSERT INTO t_60004 VALUES (1, ARRAY['a'])
+
+statement error could not remove enum value "a" as it is being used by table "test.public.t_60004"
+ALTER TYPE enum_60004 DROP VALUE 'a'
+
+# Not used in a row, so this DROP should be fine.
+statement ok
+ALTER TYPE enum_60004 DROP VALUE 'b'
+
+statement ok
+CREATE VIEW v_60004 AS SELECT ARRAY['c']:::_enum_60004 AS v;
+
+statement error count-array-type-value-usage: enum value "c" is not yet public
+ALTER TYPE enum_60004 DROP VALUE 'c'
+
+subtest regression_60004_complex
+# Setup
+statement ok
+CREATE TYPE alphabets_60004 AS ENUM ('a', 'b', 'c', 'd')
+
+statement ok
+CREATE TABLE using_alphabets_60004(k INT PRIMARY KEY, v1 alphabets_60004[], v2 alphabets_60004[])
+
+statement ok
+INSERT INTO using_alphabets_60004 VALUES (1, ARRAY['a', 'b', 'c'], ARRAY['a','b']), (2, ARRAY['a', 'b', 'c'], ARRAY['a','d'])
+
+statement ok
+CREATE TABLE using_alphabets2_60004(k INT PRIMARY KEY, v1 alphabets_60004, v2 alphabets_60004[])
+
+statement ok
+INSERT INTO using_alphabets2_60004 VALUES (1, 'a', ARRAY['b', 'a'])
+
+statement error could not remove enum value "d" as it is being used by table "test.public.using_alphabets_60004"
+ALTER TYPE alphabets_60004 DROP VALUE 'd'
+
+# Remove the row that uses 'd' in it and then try dropping.
+statement ok
+DELETE FROM using_alphabets_60004 WHERE k = 2
+
+statement ok
+ALTER TYPE alphabets_60004 DROP VALUE 'd'
+
+statement error could not remove enum value "c" as it is being used by table "test.public.using_alphabets_60004"
+ALTER TYPE alphabets_60004 DROP VALUE 'c'
+
+statement error could not remove enum value "b" as it is being used by table "test.public.using_alphabets_60004"
+ALTER TYPE alphabets_60004 DROP VALUE 'b'
+
+statement ok
+TRUNCATE using_alphabets_60004
+
+statement ok
+ALTER TYPE alphabets_60004 DROP VALUE 'c'
+
+statement error could not remove enum value "a" as it is being used by "using_alphabets2_60004"
+ALTER TYPE alphabets_60004 DROP VALUE 'a'
+
+statement error could not remove enum value "b" as it is being used by table "test.public.using_alphabets2_60004"
+ALTER TYPE alphabets_60004 DROP VALUE 'b'
+
+statement ok
+TRUNCATE using_alphabets2_60004
+
+statement ok
+ALTER TYPE alphabets_60004 DROP VALUE 'a'
+
+statement ok
+ALTER TYPE alphabets_60004 DROP VALUE 'b'

pkg/sql/type_change.go

@@ -13,6 +13,7 @@ package sql
 import (
 	"bytes"
 	"context"
+	"encoding/hex"
 	"fmt"
 	"strings"
 	"time"
@@ -723,6 +724,19 @@ func (t *typeSchemaChanger) cleanupEnumValues(ctx context.Context) error {
 	return nil
 }
 
+// convertToSQLStringRepresentation takes an array of bytes (the physical
+// representation of an enum) and converts it into a string that can be used
+// in a SQL predicate.
+func convertToSQLStringRepresentation(bytes []byte) (string, error) {
+	var byteRep strings.Builder
+	byteRep.WriteString("x'")
+	if _, err := hex.NewEncoder(&byteRep).Write(bytes); err != nil {
+		return "", err
+	}
+	byteRep.WriteString("'")
+	return byteRep.String(), nil
+}
+
 // canRemoveEnumValue returns an error if the enum value is in use and therefore
 // can't be removed.
 func (t *typeSchemaChanger) canRemoveEnumValue(
@@ -732,19 +746,6 @@ func (t *typeSchemaChanger) canRemoveEnumValue(
 	member *descpb.TypeDescriptor_EnumMember,
 	descsCol *descs.Collection,
 ) error {
-	// convertToSQLStringRepresentation takes an array of bytes (the physical
-	// representation of an enum) and converts it into a string that can be used
-	// in a SQL predicate.
-	convertToSQLStringRepresentation := func(bytes []byte) string {
-		var byteRep strings.Builder
-		byteRep.WriteString("b'")
-		for _, b := range bytes {
-			byteRep.WriteByte(b)
-		}
-		byteRep.WriteString("'")
-		return byteRep.String()
-	}
-
 	for _, ID := range typeDesc.ReferencingDescriptorIDs {
 		desc, err := descsCol.GetImmutableTableByID(ctx, txn, ID, tree.ObjectLookupFlags{})
 		if err != nil {
@@ -762,8 +763,11 @@ func (t *typeSchemaChanger) canRemoveEnumValue(
 				if !firstClause {
 					query.WriteString(" OR")
 				}
-				query.WriteString(fmt.Sprintf(" t.%s = %s", col.GetName(),
-					convertToSQLStringRepresentation(member.PhysicalRepresentation)))
+				sqlPhysRep, err := convertToSQLStringRepresentation(member.PhysicalRepresentation)
+				if err != nil {
+					return err
+				}
+				query.WriteString(fmt.Sprintf(" t.%s = %s", col.GetName(), sqlPhysRep))
 				firstClause = false
 				validationQueryConstructed = true
 			}
@@ -817,8 +821,103 @@ func (t *typeSchemaChanger) canRemoveEnumValue(
 			}
 		}
 	}
-	// We have ascertained that the value is not in use, and can therefore be
-	// safely removed.
+
+	// Do validation for the array type now.
+	arrayTypeDesc, err := descsCol.GetImmutableTypeByID(
+		ctx, txn, typeDesc.ArrayTypeID, tree.ObjectLookupFlags{})
+	if err != nil {
+		return err
+	}
+
+	return t.canRemoveEnumValueFromArrayUsages(ctx, arrayTypeDesc, member, txn, descsCol)
+}
+
+// canRemoveEnumValueFromArrayUsages returns an error if the enum member is used
+// as a value by a table/view column which type resolves to a the given array
+// type.
+func (t *typeSchemaChanger) canRemoveEnumValueFromArrayUsages(
+	ctx context.Context,
+	arrayTypeDesc *typedesc.Immutable,
+	member *descpb.TypeDescriptor_EnumMember,
+	txn *kv.Txn,
+	descsCol *descs.Collection,
+) error {
+	const validationErr = "could not validate removal of enum value %q"
+	for _, ID := range arrayTypeDesc.ReferencingDescriptorIDs {
+		desc, err := descsCol.GetImmutableTableByID(ctx, txn, ID, tree.ObjectLookupFlags{})
+		if err != nil {
+			return errors.Wrapf(err, validationErr, member.LogicalRepresentation)
+		}
+		var unionUnnests strings.Builder
+		var query strings.Builder
+
+		// Construct a query of the form:
+		// SELECT unnest FROM (
+		//	SELECT unnest(c1) FROM [SELECT %d AS t]
+		//	UNION
+		//	SELECT unnest(c2) FROM [SELECT %d AS t]
+		// 		...
+		//	) WHERE unnest = 'enum_value'
+		firstClause := true
+		for _, col := range desc.PublicColumns() {
+			if arrayTypeDesc.ID == typedesc.GetTypeDescID(col.GetType()) {
+				if !firstClause {
+					unionUnnests.WriteString(" UNION ")
+				}
+				unionUnnests.WriteString(fmt.Sprintf("SELECT unnest(%s) FROM [%d AS t]", col.GetName(), ID))
+				firstClause = false
+			}
+		}
+		// Unfortunately, we install a backreference to both the type descriptor and
+		// its array alias type regardless of the actual type of the table column.
+		// This means we may not actually construct a valid query after going
+		// through the columns, in which case there's no validation to do.
+		if firstClause {
+			continue
+		}
+		query.WriteString("SELECT unnest FROM (")
+		query.WriteString(unionUnnests.String())
+
+		sqlPhysRep, err := convertToSQLStringRepresentation(member.PhysicalRepresentation)
+		if err != nil {
+			return err
+		}
+		query.WriteString(fmt.Sprintf(") WHERE unnest = %s", sqlPhysRep))
+
+		_, dbDesc, err := descsCol.GetImmutableDatabaseByID(
+			ctx, txn, arrayTypeDesc.ParentID, tree.DatabaseLookupFlags{Required: true})
+		if err != nil {
+			return errors.Wrapf(err, validationErr, member.LogicalRepresentation)
+		}
+		override := sessiondata.InternalExecutorOverride{
+			User:     security.RootUserName(),
+			Database: dbDesc.Name,
+		}
+		rows, err := t.execCfg.InternalExecutor.QueryRowEx(
+			ctx,
+			"count-array-type-value-usage",
+			txn,
+			override,
+			query.String(),
+		)
+		if err != nil {
+			return errors.Wrapf(err, validationErr, member.LogicalRepresentation)
+		}
+		if len(rows) > 0 {
+			// Use an FQN in the error message.
+			parentSchema, err := descsCol.GetImmutableSchemaByID(
+				ctx, txn, desc.GetParentSchemaID(), tree.SchemaLookupFlags{})
+			if err != nil {
+				return err
+			}
+			fqName := tree.MakeTableNameWithSchema(tree.Name(dbDesc.GetName()), tree.Name(parentSchema.Name), tree.Name(desc.GetName()))
+
+			return pgerror.Newf(pgcode.DependentObjectsStillExist, "could not remove enum value %q as it is being used by table %q",
+				member.LogicalRepresentation, fqName.FQString(),
+			)
+		}
+	}
+	// None of the tables use the enum member in their rows.
 	return nil
 }