Test utility hook for dropping extension

This commit aims to test feasibility of using utility hook for
controlling DROP EXTENSION statement.

distribution_metadata.c

@@ -440,6 +440,40 @@ IsDistributedTable(Oid tableId)
 }
 
 
+/*
+ *  DistributedTablesExist returns true if there exists at least one distributed
+ *  table on metadata tables.
+ */
+bool
+DistributedTablesExist(void)
+{
+
+	bool distributedTablesExist = false;
+	RangeVar *heapRangeVar = NULL;
+	Relation heapRelation = NULL;
+	HeapScanDesc scanDesc = NULL;
+	HeapTuple heapTuple = NULL;
+
+	heapRangeVar = makeRangeVar(METADATA_SCHEMA_NAME, PARTITION_TABLE_NAME, -1);
+	heapRelation = relation_openrv(heapRangeVar, AccessShareLock);
+
+	scanDesc = heap_beginscan(heapRelation, SnapshotSelf, 0, NULL);
+
+	heapTuple = heap_getnext(scanDesc, ForwardScanDirection);
+
+	/*
+	 * Check if there exists any tuples in the partition table. If there are any tables,
+	 * we can conclude that there are at least one distributed table.
+	 */
+	distributedTablesExist = HeapTupleIsValid(heapTuple);
+
+	heap_endscan(scanDesc);
+	relation_close(heapRelation, AccessShareLock);
+
+	return distributedTablesExist;
+}
+
+
 /*
  * ColumnNameToColumn accepts a relation identifier and column name and returns
  * a Var that represents that column in that relation. This function throws an

distribution_metadata.h

@@ -139,6 +139,7 @@ extern List * LoadShardPlacementList(int64 shardId);
 extern Var * PartitionColumn(Oid distributedTableId);
 extern char PartitionType(Oid distributedTableId);
 extern bool IsDistributedTable(Oid tableId);
+extern bool DistributedTablesExist(void);
 extern void InsertPartitionRow(Oid distributedTableId, char partitionType,
 							   text *partitionKeyText);
 extern void InsertShardRow(Oid distributedTableId, uint64 shardId, char shardStorage,

pg_shard.c

@@ -38,6 +38,7 @@
 #include "catalog/namespace.h"
 #include "catalog/pg_class.h"
 #include "catalog/pg_type.h"
+#include "catalog/objectaddress.h"
 #include "commands/extension.h"
 #include "executor/execdesc.h"
 #include "executor/executor.h"
@@ -2007,6 +2008,71 @@ PgShardProcessUtility(Node *parsetree, const char *queryString,
 			}
 		}
 	}
+	else if (statementType == T_DropStmt)
+	{
+		DropStmt   *dropStatement = (DropStmt *) parsetree;
+
+		/* only apply hook for tables and extension */
+		if (dropStatement->removeType == OBJECT_TABLE)
+		{
+			ListCell *dropObjectCell = NULL;
+			foreach(dropObjectCell, dropStatement->objects)
+			{
+				List *tableNameList = (List *) lfirst(dropObjectCell);
+				RangeVar *rangeVar = makeRangeVarFromNameList(tableNameList);
+
+				Oid relationId = RangeVarGetRelid(rangeVar, AccessShareLock, true);
+				if (IsDistributedTable(relationId))
+				{
+					/*
+					 * Not implemented yet.
+					 */
+				}
+			}
+		}
+		else if (dropStatement->removeType == OBJECT_EXTENSION)
+		{
+			ListCell *dropStatementObject = NULL;
+
+			foreach(dropStatementObject, dropStatement->objects)
+			{
+				List *objectNameList = lfirst(dropStatementObject);
+				char *objectName = NameListToString(objectNameList);;
+
+				if (strcmp(PG_SHARD_EXTENSION_NAME, objectName) == 0)
+				{
+					/*
+					 * If pg_shard is dropped with CASCADE modifier, let it to be dropped.
+					 * However, inform the user that shards on the worker nodes will not
+					 * be dropeed.
+					 *
+					 * If CASCADE modifier is not used, check for existence of any
+					 * distributed tables. If there exists any distributed table, do not
+					 * let pg_shard to be dropped. Otherwise, let pg_shard to be dropped.
+					 *
+					 */
+					if (dropStatement->behavior == DROP_CASCADE)
+					{
+						ereport(INFO, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+										errmsg("shards on the worker nodes will not be "
+											   "dropped")));
+					}
+					else
+					{
+						bool distributedTableExists = DistributedTablesExist();
+						if (distributedTableExists == true)
+						{
+							ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+											errmsg("cannot drop %s because distributed"
+												   " table(s) exists", PG_SHARD_EXTENSION_NAME),
+											errhint("Try dropping the extension with"
+													    " CASCADE modifier.")));
+						}
+					}
+				}
+			}
+		}
+	}
 
 	if (PreviousProcessUtilityHook != NULL)
 	{
@@ -2019,3 +2085,4 @@ PgShardProcessUtility(Node *parsetree, const char *queryString,
 								params, dest, completionTag);
 	}
 }
+