Clear metadata_cache upon DROP EXTENSION

When we notice that pg_dist_partition is being invalidated we assume
that the citus extension is being dropped and drop state such as
extensionLoaded and the cached oids of all the metadata tables.

This frees the user from needing to reconnect after running DROP
EXTENSION, so we also no longer send a warning message.

src/backend/distributed/executor/multi_utility.c

@@ -66,7 +66,6 @@ static void ErrorIfUnsupportedAlterTableStmt(AlterTableStmt *alterTableStatement
 static void ErrorIfDistributedRenameStmt(RenameStmt *renameStatement);
 
 /* Local functions forward declarations for helper functions */
-static void WarnIfDropCitusExtension(DropStmt *dropStatement);
 static bool IsAlterTableRenameStmt(RenameStmt *renameStatement);
 static void ExecuteDistributedDDLCommand(Oid relationId, const char *ddlCommandString);
 static bool ExecuteCommandOnWorkerShards(Oid relationId, const char *commandString,
@@ -143,10 +142,6 @@ multi_ProcessUtility(Node *parsetree,
 		{
 			parsetree = ProcessDropIndexStmt(dropStatement, queryString);
 		}
-		else if (dropStatement->removeType == OBJECT_EXTENSION)
-		{
-			WarnIfDropCitusExtension(dropStatement);
-		}
 	}
 
 	if (IsA(parsetree, AlterTableStmt))
@@ -208,37 +203,6 @@ multi_ProcessUtility(Node *parsetree,
 }
 
 
-/*
- * WarnIfDropCitusExtension prints a WARNING if dropStatement includes dropping
- * citus extension.
- */
-static void
-WarnIfDropCitusExtension(DropStmt *dropStatement)
-{
-	ListCell *dropStatementObject = NULL;
-
-	Assert(dropStatement->removeType == OBJECT_EXTENSION);
-
-	foreach(dropStatementObject, dropStatement->objects)
-	{
-		List *objectNameList = lfirst(dropStatementObject);
-		char *objectName = NameListToString(objectNameList);
-
-		/* we're only concerned with the citus extension */
-		if (strncmp("citus", objectName, NAMEDATALEN) == 0)
-		{
-			/*
-			 * Warn the user about the possibility of invalid cache. Also, see
-			 * function comment of CachedRelationLookup().
-			 */
-			ereport(WARNING, (errmsg("could not clean the metadata cache on "
-									 "DROP EXTENSION command"),
-							  errhint("Reconnect to the server again.")));
-		}
-	}
-}
-
-
 /* Is the passed in statement a transmit statement? */
 static bool
 IsTransmitStmt(Node *parsetree)

src/backend/distributed/utils/metadata_cache.c

@@ -12,6 +12,7 @@
 #include "access/genam.h"
 #include "access/heapam.h"
 #include "access/htup_details.h"
+#include "access/xact.h"
 #include "catalog/indexing.h"
 #include "catalog/pg_namespace.h"
 #include "catalog/pg_type.h"
@@ -35,6 +36,17 @@
 #include "utils/syscache.h"
 
 
+/* state which should be cleared upon DROP EXTENSION */
+static bool extensionLoaded = false;
+static Oid distShardRelationId = InvalidOid;
+static Oid distShardPlacementRelationId = InvalidOid;
+static Oid distPartitionRelationId = InvalidOid;
+static Oid distPartitionLogicalRelidIndexId = InvalidOid;
+static Oid distShardLogicalRelidIndexId = InvalidOid;
+static Oid distShardShardidIndexId = InvalidOid;
+static Oid distShardPlacementShardidIndexId = InvalidOid;
+static Oid extraDataContainerFuncId = InvalidOid;
+
 /* Hash table for informations about each partition */
 static HTAB *DistTableCacheHash = NULL;
 
@@ -297,16 +309,10 @@ LookupDistTableCacheEntry(Oid relationId)
  * CitusHasBeenLoaded returns true if the citus extension has been created
  * in the current database and the extension script has been executed. Otherwise,
  * it returns false. The result is cached as this is called very frequently.
- *
- * NB: The way this is cached means the result will be wrong after the
- * extension is dropped. A reconnect fixes that though, so that seems
- * acceptable.
  */
 bool
 CitusHasBeenLoaded(void)
 {
-	static bool extensionLoaded = false;
-
 	/* recheck presence until citus has been loaded */
 	if (!extensionLoaded)
 	{
@@ -329,6 +335,19 @@ CitusHasBeenLoaded(void)
 		}
 
 		extensionLoaded = extensionPresent && extensionScriptExecuted;
+
+		if (extensionLoaded)
+		{
+			/*
+			 * InvalidateDistRelationCacheCallback resets state such as extensionLoaded
+			 * when it notices changes to pg_dist_partition (which usually indicate
+			 * `DROP EXTENSION citus;` has been run)
+			 *
+			 * Ensure InvalidateDistRelationCacheCallback will notice those changes
+			 * by caching pg_dist_partition's oid.
+			 */
+			DistPartitionRelationId();
+		}
 	}
 
 	return extensionLoaded;
@@ -339,11 +358,9 @@ CitusHasBeenLoaded(void)
 Oid
 DistShardRelationId(void)
 {
-	static Oid cachedOid = InvalidOid;
+	CachedRelationLookup("pg_dist_shard", &distShardRelationId);
 
-	CachedRelationLookup("pg_dist_shard", &cachedOid);
-
-	return cachedOid;
+	return distShardRelationId;
 }
 
 
@@ -351,11 +368,9 @@ DistShardRelationId(void)
 Oid
 DistShardPlacementRelationId(void)
 {
-	static Oid cachedOid = InvalidOid;
+	CachedRelationLookup("pg_dist_shard_placement", &distShardPlacementRelationId);
 
-	CachedRelationLookup("pg_dist_shard_placement", &cachedOid);
-
-	return cachedOid;
+	return distShardPlacementRelationId;
 }
 
 
@@ -363,11 +378,9 @@ DistShardPlacementRelationId(void)
 Oid
 DistPartitionRelationId(void)
 {
-	static Oid cachedOid = InvalidOid;
+	CachedRelationLookup("pg_dist_partition", &distPartitionRelationId);
 
-	CachedRelationLookup("pg_dist_partition", &cachedOid);
-
-	return cachedOid;
+	return distPartitionRelationId;
 }
 
 
@@ -375,11 +388,10 @@ DistPartitionRelationId(void)
 Oid
 DistPartitionLogicalRelidIndexId(void)
 {
-	static Oid cachedOid = InvalidOid;
+	CachedRelationLookup("pg_dist_partition_logical_relid_index",
+						 &distPartitionLogicalRelidIndexId);
 
-	CachedRelationLookup("pg_dist_partition_logical_relid_index", &cachedOid);
-
-	return cachedOid;
+	return distPartitionLogicalRelidIndexId;
 }
 
 
@@ -387,11 +399,10 @@ DistPartitionLogicalRelidIndexId(void)
 Oid
 DistShardLogicalRelidIndexId(void)
 {
-	static Oid cachedOid = InvalidOid;
+	CachedRelationLookup("pg_dist_shard_logical_relid_index",
+						 &distShardLogicalRelidIndexId);
 
-	CachedRelationLookup("pg_dist_shard_logical_relid_index", &cachedOid);
-
-	return cachedOid;
+	return distShardLogicalRelidIndexId;
 }
 
 
@@ -399,11 +410,9 @@ DistShardLogicalRelidIndexId(void)
 Oid
 DistShardShardidIndexId(void)
 {
-	static Oid cachedOid = InvalidOid;
+	CachedRelationLookup("pg_dist_shard_shardid_index", &distShardShardidIndexId);
 
-	CachedRelationLookup("pg_dist_shard_shardid_index", &cachedOid);
-
-	return cachedOid;
+	return distShardShardidIndexId;
 }
 
 
@@ -411,11 +420,10 @@ DistShardShardidIndexId(void)
 Oid
 DistShardPlacementShardidIndexId(void)
 {
-	static Oid cachedOid = InvalidOid;
+	CachedRelationLookup("pg_dist_shard_placement_shardid_index",
+						 &distShardPlacementShardidIndexId);
 
-	CachedRelationLookup("pg_dist_shard_placement_shardid_index", &cachedOid);
-
-	return cachedOid;
+	return distShardPlacementShardidIndexId;
 }
 
 
@@ -423,18 +431,17 @@ DistShardPlacementShardidIndexId(void)
 Oid
 CitusExtraDataContainerFuncId(void)
 {
-	static Oid cachedOid = 0;
 	List *nameList = NIL;
 	Oid paramOids[1] = { INTERNALOID };
 
-	if (cachedOid == InvalidOid)
+	if (extraDataContainerFuncId == InvalidOid)
 	{
 		nameList = list_make2(makeString("pg_catalog"),
 							  makeString("citus_extradata_container"));
-		cachedOid = LookupFuncName(nameList, 1, paramOids, false);
+		extraDataContainerFuncId = LookupFuncName(nameList, 1, paramOids, false);
 	}
 
-	return cachedOid;
+	return extraDataContainerFuncId;
 }
 
 
@@ -677,6 +684,24 @@ InvalidateDistRelationCacheCallback(Datum argument, Oid relationId)
 			cacheEntry->isValid = false;
 		}
 	}
+
+	/*
+	 * If pg_dist_partition is being invalidated drop all state
+	 * This happens pretty rarely, but most importantly happens during
+	 * DROP EXTENSION citus;
+	 */
+	if (relationId != InvalidOid && relationId == distPartitionRelationId)
+	{
+		extensionLoaded = false;
+		distShardRelationId = InvalidOid;
+		distShardPlacementRelationId = InvalidOid;
+		distPartitionRelationId = InvalidOid;
+		distPartitionLogicalRelidIndexId = InvalidOid;
+		distShardLogicalRelidIndexId = InvalidOid;
+		distShardShardidIndexId = InvalidOid;
+		distShardPlacementShardidIndexId = InvalidOid;
+		extraDataContainerFuncId = InvalidOid;
+	}
 }
 
 
@@ -878,10 +903,6 @@ TupleToShardInterval(HeapTuple heapTuple, TupleDesc tupleDescriptor, Oid interva
 /*
  * CachedRelationLookup performs a cached lookup for the relation
  * relationName, with the result cached in *cachedOid.
- *
- * NB: The way this is cached means the result will be wrong after the
- * extension is dropped and reconnect. A reconnect fixes that though, so that
- * seems acceptable.
  */
 static void
 CachedRelationLookup(const char *relationName, Oid *cachedOid)

src/test/regress/expected/multi_drop_extension.out

@@ -0,0 +1,39 @@
+--
+-- MULTI_DROP_EXTENSION
+--
+-- Tests around dropping and recreating the extension
+CREATE TABLE testtableddl(somecol int, distributecol text NOT NULL);
+SELECT master_create_distributed_table('testtableddl', 'distributecol', 'append');
+ master_create_distributed_table 
+---------------------------------
+ 
+(1 row)
+
+-- this emits a NOTICE message for every table we are dropping with our CASCADE. It would
+-- be nice to check that we get those NOTICE messages, but it's nicer to not have to
+-- change this test every time the previous tests change the set of tables they leave
+-- around.
+SET client_min_messages TO 'WARNING';
+DROP EXTENSION citus CASCADE;
+RESET client_min_messages;
+CREATE EXTENSION citus;
+-- verify that a table can be created after the extension has been dropped and recreated
+CREATE TABLE testtableddl(somecol int, distributecol text NOT NULL);
+SELECT master_create_distributed_table('testtableddl', 'distributecol', 'append');
+ master_create_distributed_table 
+---------------------------------
+ 
+(1 row)
+
+SELECT 1 FROM master_create_empty_shard('testtableddl');
+ ?column? 
+----------
+        1
+(1 row)
+
+SELECT * FROM testtableddl;
+ somecol | distributecol 
+---------+---------------
+(0 rows)
+
+DROP TABLE testtableddl;

src/test/regress/expected/multi_table_ddl.out

@@ -11,8 +11,6 @@ SELECT master_create_distributed_table('testtableddl', 'distributecol', 'append'
 
 -- verify that the citus extension can't be dropped while distributed tables exist
 DROP EXTENSION citus;
-WARNING:  could not clean the metadata cache on DROP EXTENSION command
-HINT:  Reconnect to the server again.
 ERROR:  cannot drop extension citus because other objects depend on it
 DETAIL:  table testtableddl depends on extension citus
 HINT:  Use DROP ... CASCADE to drop the dependent objects too.
@@ -62,11 +60,6 @@ SELECT * FROM pg_dist_shard_placement;
 ---------+------------+-------------+----------+----------
 (0 rows)
 
--- check that the extension now can be dropped (and recreated). We reconnect
--- before creating the extension to expire extension specific variables which
--- are cached for performance.
+-- check that the extension now can be dropped (and recreated)
 DROP EXTENSION citus;
-WARNING:  could not clean the metadata cache on DROP EXTENSION command
-HINT:  Reconnect to the server again.
-\c
 CREATE EXTENSION citus;

src/test/regress/multi_schedule

@@ -135,3 +135,8 @@ test: multi_router_planner
 # multi_large_shardid stages more shards into lineitem
 # ----------
 test: multi_large_shardid
+
+# ----------
+# multi_drop_extension makes sure we can safely drop and recreate the extension
+# ----------
+test: multi_drop_extension

src/test/regress/sql/multi_drop_extension.sql

@@ -0,0 +1,25 @@
+--
+-- MULTI_DROP_EXTENSION
+--
+-- Tests around dropping and recreating the extension
+
+
+CREATE TABLE testtableddl(somecol int, distributecol text NOT NULL);
+SELECT master_create_distributed_table('testtableddl', 'distributecol', 'append');
+
+-- this emits a NOTICE message for every table we are dropping with our CASCADE. It would
+-- be nice to check that we get those NOTICE messages, but it's nicer to not have to
+-- change this test every time the previous tests change the set of tables they leave
+-- around.
+SET client_min_messages TO 'WARNING';
+DROP EXTENSION citus CASCADE;
+RESET client_min_messages;
+
+CREATE EXTENSION citus;
+
+-- verify that a table can be created after the extension has been dropped and recreated
+CREATE TABLE testtableddl(somecol int, distributecol text NOT NULL);
+SELECT master_create_distributed_table('testtableddl', 'distributecol', 'append');
+SELECT 1 FROM master_create_empty_shard('testtableddl');
+SELECT * FROM testtableddl;
+DROP TABLE testtableddl;

src/test/regress/sql/multi_table_ddl.sql

@@ -34,9 +34,6 @@ SELECT * FROM pg_dist_partition;
 SELECT * FROM pg_dist_shard;
 SELECT * FROM pg_dist_shard_placement;
 
--- check that the extension now can be dropped (and recreated). We reconnect
--- before creating the extension to expire extension specific variables which
--- are cached for performance.
+-- check that the extension now can be dropped (and recreated)
 DROP EXTENSION citus;
-\c
 CREATE EXTENSION citus;