From 354ec1b5e86c673c27d1e578c8cddfa483fac659 Mon Sep 17 00:00:00 2001
From: Mats Kindahl <mats@timescale.com>
Date: Sun, 14 Feb 2021 18:58:43 +0100
Subject: Add callback to reset filenode to table access method

This commit adds a callback to `TableAccessMethod` that is called when
the table should be scheduled for unlinking and implement the method
for the heap access method.
---
 src/backend/access/heap/heapam_handler.c |  7 +++++++
 src/backend/access/table/tableamapi.c    |  1 +
 src/backend/catalog/heap.c               | 23 ++++++++++++++++++++++-
 src/include/access/tableam.h             | 20 ++++++++++++++++++++
 4 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/src/backend/access/heap/heapam_handler.c b/src/backend/access/heap/heapam_handler.c
index 4a70e20a14..287f8a47cf 100644
--- a/src/backend/access/heap/heapam_handler.c
+++ b/src/backend/access/heap/heapam_handler.c
@@ -613,6 +613,12 @@ heapam_relation_set_new_filenode(Relation rel,
 	smgrclose(srel);
 }
 
+static void
+heapam_relation_reset_filenode(Relation rel)
+{
+	RelationDropStorage(rel);
+}
+
 static void
 heapam_relation_nontransactional_truncate(Relation rel)
 {
@@ -2566,6 +2572,7 @@ static const TableAmRoutine heapam_methods = {
 	.index_delete_tuples = heap_index_delete_tuples,
 
 	.relation_set_new_filenode = heapam_relation_set_new_filenode,
+	.relation_reset_filenode = heapam_relation_reset_filenode,
 	.relation_nontransactional_truncate = heapam_relation_nontransactional_truncate,
 	.relation_copy_data = heapam_relation_copy_data,
 	.relation_copy_for_cluster = heapam_relation_copy_for_cluster,
diff --git a/src/backend/access/table/tableamapi.c b/src/backend/access/table/tableamapi.c
index 325ecdc122..30e070383c 100644
--- a/src/backend/access/table/tableamapi.c
+++ b/src/backend/access/table/tableamapi.c
@@ -83,6 +83,7 @@ GetTableAmRoutine(Oid amhandler)
 	Assert(routine->tuple_lock != NULL);
 
 	Assert(routine->relation_set_new_filenode != NULL);
+	Assert(routine->relation_reset_filenode != NULL);
 	Assert(routine->relation_nontransactional_truncate != NULL);
 	Assert(routine->relation_copy_data != NULL);
 	Assert(routine->relation_copy_for_cluster != NULL);
diff --git a/src/backend/catalog/heap.c b/src/backend/catalog/heap.c
index 9abc4a1f55..e1e8b93285 100644
--- a/src/backend/catalog/heap.c
+++ b/src/backend/catalog/heap.c
@@ -1990,7 +1990,28 @@ heap_drop_with_catalog(Oid relid)
 	 * Schedule unlinking of the relation's physical files at commit.
 	 */
 	if (RELKIND_HAS_STORAGE(rel->rd_rel->relkind))
-		RelationDropStorage(rel);
+		switch (rel->rd_rel->relkind)
+		{
+			case RELKIND_VIEW:
+			case RELKIND_COMPOSITE_TYPE:
+			case RELKIND_FOREIGN_TABLE:
+			case RELKIND_PARTITIONED_TABLE:
+			case RELKIND_PARTITIONED_INDEX:
+				Assert(false);
+				break;
+
+			case RELKIND_INDEX:
+			case RELKIND_SEQUENCE:
+				RelationDropStorage(rel);
+				break;
+
+			case RELKIND_RELATION:
+			case RELKIND_TOASTVALUE:
+			case RELKIND_MATVIEW:
+				table_relation_reset_filenode(rel);
+				break;
+
+		}
 
 	/*
 	 * Close relcache entry, but *keep* AccessExclusiveLock on the relation
diff --git a/src/include/access/tableam.h b/src/include/access/tableam.h
index 33bffb6815..b5bb44a470 100644
--- a/src/include/access/tableam.h
+++ b/src/include/access/tableam.h
@@ -550,6 +550,14 @@ typedef struct TableAmRoutine
 											  TransactionId *freezeXid,
 											  MultiXactId *minmulti);
 
+	/*
+	 * This callback needs to remove all associations with the relation `rel`
+	 * since the relation is being dropped.
+	 *
+	 * See also table_relation_reset_filenode().
+	 */
+	void        (*relation_reset_filenode) (Relation rel);
+
 	/*
 	 * This callback needs to remove all contents from `rel`'s current
 	 * relfilenode. No provisions for transactional behaviour need to be made.
@@ -1510,6 +1518,18 @@ table_relation_set_new_filenode(Relation rel,
 											   freezeXid, minmulti);
 }
 
+/*
+ * Remove all association with storage for the relation.
+ *
+ * This is used when a relation is about to be dropped and removed from the
+ * catalog.
+ */
+static inline void
+table_relation_reset_filenode(Relation rel)
+{
+	rel->rd_tableam->relation_reset_filenode(rel);
+}
+
 /*
  * Remove all table contents from `rel`, in a non-transactional manner.
  * Non-transactional meaning that there's no need to support rollbacks. This
-- 
2.25.1