From 81db8df7b012211685b1851106acf3db83df4c4d Mon Sep 17 00:00:00 2001
From: reshke kirill <reshke@double.cloud>
Date: Sun, 27 Oct 2024 07:48:00 +0000
Subject: [PATCH v2] Change RELKIND_HAS_TABLE_AM to support partitioned tables

---
 src/backend/commands/tablecmds.c     | 7 +++++--
 src/backend/optimizer/util/plancat.c | 2 +-
 src/backend/storage/buffer/bufmgr.c  | 2 +-
 src/backend/utils/cache/relcache.c   | 2 +-
 src/bin/pg_dump/pg_dump.c            | 3 +--
 src/include/catalog/pg_class.h       | 1 +
 6 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index e14bc0c054..c76294e6cb 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -979,10 +979,10 @@ DefineRelation(CreateStmt *stmt, char relkind, Oid ownerId,
 	 */
 	if (stmt->accessMethod != NULL)
 	{
-		Assert(RELKIND_HAS_TABLE_AM(relkind) || relkind == RELKIND_PARTITIONED_TABLE);
+		Assert(RELKIND_HAS_TABLE_AM(relkind));
 		accessMethodId = get_table_am_oid(stmt->accessMethod, false);
 	}
-	else if (RELKIND_HAS_TABLE_AM(relkind) || relkind == RELKIND_PARTITIONED_TABLE)
+	else if (RELKIND_HAS_TABLE_AM(relkind))
 	{
 		if (stmt->partbound)
 		{
@@ -15029,6 +15029,9 @@ ATExecSetAccessMethodNoStorage(Relation rel, Oid newAccessMethodId)
 
 	/* Update the pg_class row. */
 	oldAccessMethodId = rd_rel->relam;
+	if (!OidIsValid(newAccessMethodId))
+		newAccessMethodId = get_table_am_oid(default_table_access_method, false);
+
 	rd_rel->relam = newAccessMethodId;
 
 	/* Leave if no update required */
diff --git a/src/backend/optimizer/util/plancat.c b/src/backend/optimizer/util/plancat.c
index b913f91ff0..b8e3d9f5bd 100644
--- a/src/backend/optimizer/util/plancat.c
+++ b/src/backend/optimizer/util/plancat.c
@@ -1072,7 +1072,7 @@ estimate_rel_size(Relation rel, int32 *attr_widths,
 	BlockNumber relallvisible;
 	double		density;
 
-	if (RELKIND_HAS_TABLE_AM(rel->rd_rel->relkind))
+	if (RELKIND_HAS_TABLE_AM(rel->rd_rel->relkind) && rel->rd_rel->relkind != RELKIND_PARTITIONED_TABLE)
 	{
 		table_relation_estimate_size(rel, attr_widths, pages, tuples,
 									 allvisfrac);
diff --git a/src/backend/storage/buffer/bufmgr.c b/src/backend/storage/buffer/bufmgr.c
index 0f02bf62fa..1c18df7b80 100644
--- a/src/backend/storage/buffer/bufmgr.c
+++ b/src/backend/storage/buffer/bufmgr.c
@@ -3922,7 +3922,7 @@ FlushBuffer(BufferDesc *buf, SMgrRelation reln, IOObject io_object,
 BlockNumber
 RelationGetNumberOfBlocksInFork(Relation relation, ForkNumber forkNum)
 {
-	if (RELKIND_HAS_TABLE_AM(relation->rd_rel->relkind))
+	if (RELKIND_HAS_TABLE_AM(relation->rd_rel->relkind) && relation->rd_rel->relkind != RELKIND_PARTITIONED_TABLE)
 	{
 		/*
 		 * Not every table AM uses BLCKSZ wide fixed size blocks. Therefore
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index c326f687eb..feb90c5ee7 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -3863,7 +3863,7 @@ RelationSetNewRelfilenumber(Relation relation, char persistence)
 	newrlocator = relation->rd_locator;
 	newrlocator.relNumber = newrelfilenumber;
 
-	if (RELKIND_HAS_TABLE_AM(relation->rd_rel->relkind))
+	if (RELKIND_HAS_TABLE_AM(relation->rd_rel->relkind) && relation->rd_rel->relkind != RELKIND_PARTITIONED_TABLE)
 	{
 		table_relation_set_new_filelocator(relation, &newrlocator,
 										   persistence,
diff --git a/src/bin/pg_dump/pg_dump.c b/src/bin/pg_dump/pg_dump.c
index 1b47c388ce..665fb4c27a 100644
--- a/src/bin/pg_dump/pg_dump.c
+++ b/src/bin/pg_dump/pg_dump.c
@@ -16573,8 +16573,7 @@ dumpTableSchema(Archive *fout, const TableInfo *tbinfo)
 		if (RELKIND_HAS_TABLESPACE(tbinfo->relkind))
 			tablespace = tbinfo->reltablespace;
 
-		if (RELKIND_HAS_TABLE_AM(tbinfo->relkind) ||
-			tbinfo->relkind == RELKIND_PARTITIONED_TABLE)
+		if (RELKIND_HAS_TABLE_AM(tbinfo->relkind))
 			tableam = tbinfo->amname;
 
 		ArchiveEntry(fout, tbinfo->dobj.catId, tbinfo->dobj.dumpId,
diff --git a/src/include/catalog/pg_class.h b/src/include/catalog/pg_class.h
index 0fc2c093b0..66b8748e3e 100644
--- a/src/include/catalog/pg_class.h
+++ b/src/include/catalog/pg_class.h
@@ -225,6 +225,7 @@ MAKE_SYSCACHE(RELNAMENSP, pg_class_relname_nsp_index, 128);
  */
 #define RELKIND_HAS_TABLE_AM(relkind) \
 	((relkind) == RELKIND_RELATION || \
+	 (relkind) == RELKIND_PARTITIONED_TABLE || \
 	 (relkind) == RELKIND_TOASTVALUE || \
 	 (relkind) == RELKIND_MATVIEW)
 
-- 
2.34.1