Disallow dropping type cascade which is used in partition column

src/backend/distributed/commands/distribute_object_ops.c

@@ -1134,7 +1134,7 @@ static DistributeObjectOps View_AlterView = {
 static DistributeObjectOps Type_Drop = {
 	.deparse = DeparseDropTypeStmt,
 	.qualify = NULL,
-	.preprocess = PreprocessDropDistributedObjectStmt,
+	.preprocess = PreprocessDropTypeStmt,
 	.postprocess = NULL,
 	.operationType = DIST_OPS_DROP,
 	.address = NULL,

src/backend/distributed/commands/type.c

@@ -101,6 +101,7 @@ static CompositeTypeStmt * RecreateCompositeTypeStmt(Oid typeOid);
 static List * CompositeTypeColumnDefList(Oid typeOid);
 static CreateEnumStmt * RecreateEnumStmt(Oid typeOid);
 static List * EnumValsList(Oid typeOid);
+static ObjectAddress * typeDependentDistributionColumn(Oid typeId);
 
 /*
  * PreprocessRenameTypeAttributeStmt is called for changes of attribute names for composite
@@ -142,6 +143,63 @@ PreprocessRenameTypeAttributeStmt(Node *node, const char *queryString,
 }
 
 
+/*
+ * PreprocessDropTypeStmt is called to check whether this is a
+ * DROP TYPE ... CASCADE statement. If yes, it checks whether there are any
+ * columns part of distributed tables included. If yes, it errors out.
+ * After that it calls the general PreprocessDropDistributedObjectStmt
+ */
+List *
+PreprocessDropTypeStmt(Node *node, const char *queryString,
+					   ProcessUtilityContext processUtilityContext)
+{
+	DropStmt *stmt = castNode(DropStmt, node);
+	Assert(stmt->removeType == OBJECT_TYPE);
+
+	if (stmt->behavior == DROP_CASCADE)
+	{
+		ListCell *cell;
+		foreach(cell, stmt->objects)
+		{
+			Node *object = lfirst(cell);
+			Relation relation = NULL;
+
+			/* Get an ObjectAddress for the type. */
+			ObjectAddress typeAddress = get_object_address(OBJECT_TYPE,
+														   object,
+														   &relation,
+														   AccessExclusiveLock,
+														   stmt->missing_ok);
+
+			ObjectAddress *distTableAndColumn = typeDependentDistributionColumn(
+				typeAddress.objectId);
+			if (distTableAndColumn != NULL)
+			{
+				Oid distTableId = distTableAndColumn->objectId;
+				int32 attNum = distTableAndColumn->objectSubId;
+
+				HeapTuple attTuple = SearchSysCacheAttNum(distTableId, attNum);
+				Form_pg_attribute targetAttr = (Form_pg_attribute) GETSTRUCT(attTuple);
+				char *columnName = NameStr(targetAttr->attname);
+				ReleaseSysCache(attTuple);
+
+				TypeName *typeName = castNode(TypeName, object);
+				Oid typeOid = LookupTypeNameOid(NULL, typeName, false);
+				const char *identifier = format_type_be_qualified(typeOid);
+
+				ereport(ERROR,
+						(errmsg("Can't drop type \"%s\" "
+								"because it is used in the partition column \"%s\" "
+								"of Citus table \"%s\"",
+								identifier, columnName, get_rel_name(distTableId))));
+			}
+		}
+	}
+
+	return PreprocessDropDistributedObjectStmt(node, queryString, processUtilityContext);
+}
+
+
 /*
  * CreateTypeStmtByObjectAddress returns a parsetree for the CREATE TYPE statement to
  * recreate the type by its object address.
@@ -669,3 +727,57 @@ LookupNonAssociatedArrayTypeNameOid(ParseState *pstate, const TypeName *typeName
 
 	return typeOid;
 }
+
+
+/*
+ * typeDependentDistributionColumn checks whether the given type has a dependency
+ * to a Citus table as the type of the table's distribution column
+ * If there exists at least one, it returns the first object seen in the list.
+ */
+static ObjectAddress *
+typeDependentDistributionColumn(Oid typeId)
+{
+	ObjectAddress *distTableAndColumn = NULL;
+
+	ScanKeyData key[2];
+
+	Relation depRel = table_open(DependRelationId, AccessShareLock);
+
+	ScanKeyInit(&key[0],
+				Anum_pg_depend_refclassid,
+				BTEqualStrategyNumber, F_OIDEQ,
+				ObjectIdGetDatum(TypeRelationId));
+	ScanKeyInit(&key[1],
+				Anum_pg_depend_refobjid,
+				BTEqualStrategyNumber, F_OIDEQ,
+				ObjectIdGetDatum(typeId));
+
+	SysScanDesc scan = systable_beginscan(depRel, DependReferenceIndexId, true,
+										  NULL, 2, key);
+
+	HeapTuple tup;
+	while (HeapTupleIsValid(tup = systable_getnext(scan)))
+	{
+		Form_pg_depend pgDependEntry = (Form_pg_depend) GETSTRUCT(tup);
+
+		if (pgDependEntry->classid == RelationRelationId &&
+			get_rel_relkind(pgDependEntry->objid) == RELKIND_RELATION &&
+			IsCitusTable(pgDependEntry->objid))
+		{
+			Var *partitionColumn = DistPartitionKey(pgDependEntry->objid);
+			if (partitionColumn != NULL &&
+				partitionColumn->varattno == pgDependEntry->objsubid)
+			{
+				distTableAndColumn = palloc0(sizeof(ObjectAddress));
+				ObjectAddressSubSet(*distTableAndColumn, RelationRelationId,
+									pgDependEntry->objid, pgDependEntry->objsubid);
+				break;
+			}
+		}
+	}
+
+	systable_endscan(scan);
+	table_close(depRel, AccessShareLock);
+
+	return distTableAndColumn;
+}

src/include/distributed/commands.h

@@ -607,6 +607,8 @@ extern void PreprocessTruncateStatement(TruncateStmt *truncateStatement);
 extern List * PreprocessRenameTypeAttributeStmt(Node *stmt, const char *queryString,
 												ProcessUtilityContext
 												processUtilityContext);
+extern List * PreprocessDropTypeStmt(Node *stmt, const char *queryString,
+									 ProcessUtilityContext processUtilityContext);
 extern Node * CreateTypeStmtByObjectAddress(const ObjectAddress *address);
 extern List * CompositeTypeStmtObjectAddress(Node *stmt, bool missing_ok, bool
 											 isPostprocess);

src/test/regress/expected/distributed_types.out

@@ -578,6 +578,68 @@ DETAIL:  "type temp_type" will be created only locally
 CREATE TYPE pg_temp.temp_enum AS ENUM ('one', 'two', 'three');
 WARNING:  "type temp_enum" has dependency on unsupported object "schema pg_temp_xxx"
 DETAIL:  "type temp_enum" will be created only locally
+-- verify that dropping a type which is used in the distribution column
+-- of a distributed table fails
+-- create a custom type...
+CREATE TYPE my_type AS (
+    i integer,
+    i2 integer
+);
+CREATE FUNCTION equal_my_type_function(my_type, my_type) RETURNS boolean
+LANGUAGE 'internal'
+AS 'record_eq'
+IMMUTABLE
+RETURNS NULL ON NULL INPUT;
+-- ... use that function to create a custom equality operator...
+CREATE OPERATOR = (
+   LEFTARG = my_type,
+   RIGHTARG = my_type,
+   PROCEDURE = equal_my_type_function,
+   HASHES
+);
+-- ... create a test HASH function. Though it is a poor hash function,
+-- it is acceptable for our tests
+CREATE FUNCTION my_type_hash(my_type) RETURNS int
+AS 'SELECT hashtext( ($1.i + $1.i2)::text);'
+LANGUAGE SQL
+IMMUTABLE
+RETURNS NULL ON NULL INPUT;
+-- ... and create a custom operator family for hash indexes...
+CREATE OPERATOR FAMILY cats_op_fam USING hash;
+-- We need to define a default operator classes for my_type
+-- that uses HASH
+CREATE OPERATOR CLASS cats_op_fam_class
+DEFAULT FOR TYPE my_type USING HASH AS
+OPERATOR 1 = (my_type, my_type),
+FUNCTION 1 my_type_hash(my_type);
+CREATE TABLE tbl (a my_type);
+SELECT create_distributed_table('tbl', 'a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+DROP TYPE my_type CASCADE;
+ERROR:  Can't drop type "type_tests.my_type" because it is used in the partition column "a" of Citus table "tbl"
+ALTER TABLE tbl DROP COLUMN a;
+ERROR:  cannot execute ALTER TABLE command involving partition column
+SELECT undistribute_table('tbl');
+NOTICE:  creating a new table for type_tests.tbl
+NOTICE:  moving the data of type_tests.tbl
+NOTICE:  dropping the old type_tests.tbl
+NOTICE:  renaming the new table to type_tests.tbl
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
+DROP TYPE my_type CASCADE;
+NOTICE:  drop cascades to 5 other objects
+DETAIL:  drop cascades to function equal_my_type_function(my_type,my_type)
+drop cascades to operator =(my_type,my_type)
+drop cascades to function my_type_hash(my_type)
+drop cascades to operator class cats_op_fam_class for access method hash
+drop cascades to column a of table tbl
 -- clear objects
 SET client_min_messages TO error; -- suppress cascading objects dropping
 DROP SCHEMA type_tests CASCADE;

src/test/regress/sql/distributed_types.sql

@@ -350,6 +350,55 @@ SELECT create_distributed_table('table_text_local_def','id');
 CREATE TYPE pg_temp.temp_type AS (int_field int);
 CREATE TYPE pg_temp.temp_enum AS ENUM ('one', 'two', 'three');
 
+-- verify that dropping a type which is used in the distribution column
+-- of a distributed table fails
+
+-- create a custom type...
+CREATE TYPE my_type AS (
+    i integer,
+    i2 integer
+);
+
+CREATE FUNCTION equal_my_type_function(my_type, my_type) RETURNS boolean
+LANGUAGE 'internal'
+AS 'record_eq'
+IMMUTABLE
+RETURNS NULL ON NULL INPUT;
+
+-- ... use that function to create a custom equality operator...
+CREATE OPERATOR = (
+   LEFTARG = my_type,
+   RIGHTARG = my_type,
+   PROCEDURE = equal_my_type_function,
+   HASHES
+);
+
+-- ... create a test HASH function. Though it is a poor hash function,
+-- it is acceptable for our tests
+CREATE FUNCTION my_type_hash(my_type) RETURNS int
+AS 'SELECT hashtext( ($1.i + $1.i2)::text);'
+LANGUAGE SQL
+IMMUTABLE
+RETURNS NULL ON NULL INPUT;
+
+-- ... and create a custom operator family for hash indexes...
+CREATE OPERATOR FAMILY cats_op_fam USING hash;
+
+-- We need to define a default operator classes for my_type
+-- that uses HASH
+
+CREATE OPERATOR CLASS cats_op_fam_class
+DEFAULT FOR TYPE my_type USING HASH AS
+OPERATOR 1 = (my_type, my_type),
+FUNCTION 1 my_type_hash(my_type);
+
+CREATE TABLE tbl (a my_type);
+SELECT create_distributed_table('tbl', 'a');
+DROP TYPE my_type CASCADE;
+ALTER TABLE tbl DROP COLUMN a;
+SELECT undistribute_table('tbl');
+DROP TYPE my_type CASCADE;
+
 -- clear objects
 SET client_min_messages TO error; -- suppress cascading objects dropping
 DROP SCHEMA type_tests CASCADE;