From 276deacbe133af10aff2aeba482931d180ce9fe3 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Fri, 8 Mar 2024 13:24:14 +0900
Subject: [PATCH v5] Allow specifying access method for partitioned tables

Partitioned tables gain support for CREATE TABLE .. USING, where
specifying an access method can serve as a hint when creating partitions
on it.  This includes support for ALTER TABLE .. SET ACCESS METHOD,
where specifying DEFAULT is equivalent to resetting the partitioned
table's relam to 0.  Specifying a non-DEFAULT access method will set its
relam to not be 0.

The historical default of using pg_class.relam as InvalidOid is
preserved, corresponding to the case where USING clause is not used.  In
this case, like previous releases, a partition is created with the
access method set by default_table_access_method if its CREATE TABLE
query does not specify an access method with a USING clause.

The relcache of partitioned tables is not changed: rd_tableam is not
set, even if a partitioned table has a relam set.

Authors: Justin Pryzby, Soumyadeep Chakraborty
Reviewed-by: Michael Paquier, and a few others..
Discussion: https://postgr.es/m/
---
 doc/src/sgml/catalogs.sgml              |   8 +-
 doc/src/sgml/ref/alter_table.sgml       |   8 ++
 doc/src/sgml/ref/create_table.sgml      |   4 +
 src/backend/commands/tablecmds.c        | 165 +++++++++++++++++++-----
 src/backend/utils/cache/lsyscache.c     |  22 ++++
 src/backend/utils/cache/relcache.c      |   7 +
 src/bin/pg_dump/pg_dump.c               |   3 +-
 src/bin/pg_dump/t/002_pg_dump.pl        |  35 +++++
 src/include/catalog/pg_class.h          |   4 +-
 src/include/utils/lsyscache.h           |   1 +
 src/test/regress/expected/create_am.out | 158 ++++++++++++++++++++++-
 src/test/regress/sql/create_am.sql      |  91 ++++++++++++-
 12 files changed, 462 insertions(+), 44 deletions(-)

diff --git a/doc/src/sgml/catalogs.sgml b/doc/src/sgml/catalogs.sgml
index b7980eb499..03815e10e8 100644
--- a/doc/src/sgml/catalogs.sgml
+++ b/doc/src/sgml/catalogs.sgml
@@ -1989,8 +1989,12 @@ SCRAM-SHA-256$<replaceable>&lt;iteration count&gt;</replaceable>:<replaceable>&l
       </para>
       <para>
        If this is a table or an index, the access method used (heap,
-       B-tree, hash, etc.); otherwise zero (zero occurs for sequences,
-       as well as relations without storage, such as views)
+       B-tree, hash, etc.); for partitioned tables, the access method
+       for newly created partitions, when one is not specified in the
+       creation command, or zero to make creation fall back to
+       <varname>default_table_access_method</varname>.
+       Zero for sequences, as well as other relations without
+       storage, such as views.
       </para></entry>
      </row>
 
diff --git a/doc/src/sgml/ref/alter_table.sgml b/doc/src/sgml/ref/alter_table.sgml
index 96e3d77605..b6e794e6ce 100644
--- a/doc/src/sgml/ref/alter_table.sgml
+++ b/doc/src/sgml/ref/alter_table.sgml
@@ -737,6 +737,14 @@ WITH ( MODULUS <replaceable class="parameter">numeric_literal</replaceable>, REM
       <literal>DEFAULT</literal> changes the access method of the table
       to <xref linkend="guc-default-table-access-method"/>.
      </para>
+     <para>
+      When applied to a partitioned table, there is no data to rewrite, but any
+      partitions created afterwards will use that access method unless
+      overridden by a <literal>USING</literal> clause. Using
+      <varname>DEFAULT</varname> switches the partitioned table to use
+      <varname>default_table_access_method</varname> when new partitions are
+      created.
+     </para>
     </listitem>
    </varlistentry>
 
diff --git a/doc/src/sgml/ref/create_table.sgml b/doc/src/sgml/ref/create_table.sgml
index 4cbaaccaf7..b79081a5ec 100644
--- a/doc/src/sgml/ref/create_table.sgml
+++ b/doc/src/sgml/ref/create_table.sgml
@@ -1330,6 +1330,10 @@ WITH ( MODULUS <replaceable class="parameter">numeric_literal</replaceable>, REM
       method is chosen for the new table. See <xref
       linkend="guc-default-table-access-method"/> for more information.
      </para>
+     <para>
+      When creating a partition, the table access method is the access method
+      of its partitioned table, if set.
+     </para>
     </listitem>
    </varlistentry>
 
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 6c0c899210..9e36aae425 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -183,7 +183,9 @@ typedef struct AlteredTableInfo
 	List	   *afterStmts;		/* List of utility command parsetrees */
 	bool		verify_new_notnull; /* T if we should recheck NOT NULL */
 	int			rewrite;		/* Reason for forced rewrite, if any */
-	Oid			newAccessMethod;	/* new access method; 0 means no change */
+	bool		chgAccessMethod;	/* T if SET ACCESS METHOD is used */
+	Oid			newAccessMethod;	/* new access method; 0 means no change,
+									 * if above is true */
 	Oid			newTableSpace;	/* new tablespace; 0 means no change */
 	bool		chgPersistence; /* T if SET LOGGED/UNLOGGED is used */
 	char		newrelpersistence;	/* if above is true */
@@ -589,6 +591,7 @@ static ObjectAddress ATExecClusterOn(Relation rel, const char *indexName,
 									 LOCKMODE lockmode);
 static void ATExecDropCluster(Relation rel, LOCKMODE lockmode);
 static void ATPrepSetAccessMethod(AlteredTableInfo *tab, Relation rel, const char *amname);
+static void ATExecSetAccessMethodNoStorage(Relation rel, Oid newAccessMethod);
 static bool ATPrepChangePersistence(Relation rel, bool toLogged);
 static void ATPrepSetTableSpace(AlteredTableInfo *tab, Relation rel,
 								const char *tablespacename, LOCKMODE lockmode);
@@ -703,7 +706,6 @@ DefineRelation(CreateStmt *stmt, char relkind, Oid ownerId,
 	Oid			ofTypeId;
 	ObjectAddress address;
 	LOCKMODE	parentLockmode;
-	const char *accessMethod = NULL;
 	Oid			accessMethodId = InvalidOid;
 
 	/*
@@ -948,24 +950,22 @@ DefineRelation(CreateStmt *stmt, char relkind, Oid ownerId,
 	}
 
 	/*
-	 * If the statement hasn't specified an access method, but we're defining
-	 * a type of relation that needs one, use the default.
+	 * Select access method to use: an explicitly indicated one, or (in the
+	 * case of a partitioned table) the parent's, if it has one.
 	 */
 	if (stmt->accessMethod != NULL)
+		accessMethodId = get_table_am_oid(stmt->accessMethod, false);
+	else if (stmt->partbound)
 	{
-		accessMethod = stmt->accessMethod;
-
-		if (partitioned)
-			ereport(ERROR,
-					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-					 errmsg("specifying a table access method is not supported on a partitioned table")));
+		Assert(list_length(inheritOids) == 1);
+		accessMethodId = get_rel_relam(linitial_oid(inheritOids));
 	}
-	else if (RELKIND_HAS_TABLE_AM(relkind))
-		accessMethod = default_table_access_method;
+	else
+		accessMethodId = InvalidOid;
 
-	/* look up the access method, verify it is for a table */
-	if (accessMethod != NULL)
-		accessMethodId = get_table_am_oid(accessMethod, false);
+	/* still nothing? use the default */
+	if (RELKIND_HAS_TABLE_AM(relkind) && !OidIsValid(accessMethodId))
+		accessMethodId = get_table_am_oid(default_table_access_method, false);
 
 	/*
 	 * Create the relation.  Inherited defaults and constraints are passed in
@@ -5041,14 +5041,8 @@ ATPrepCmd(List **wqueue, Relation rel, AlterTableCmd *cmd,
 		case AT_SetAccessMethod:	/* SET ACCESS METHOD */
 			ATSimplePermissions(cmd->subtype, rel, ATT_TABLE | ATT_MATVIEW);
 
-			/* partitioned tables don't have an access method */
-			if (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
-				ereport(ERROR,
-						(errcode(ERRCODE_WRONG_OBJECT_TYPE),
-						 errmsg("cannot change access method of a partitioned table")));
-
 			/* check if another access method change was already requested */
-			if (OidIsValid(tab->newAccessMethod))
+			if (tab->chgAccessMethod)
 				ereport(ERROR,
 						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 						 errmsg("cannot have multiple SET ACCESS METHOD subcommands")));
@@ -5402,7 +5396,14 @@ ATExecCmd(List **wqueue, AlteredTableInfo *tab,
 			/* nothing to do here, oid columns don't exist anymore */
 			break;
 		case AT_SetAccessMethod:	/* SET ACCESS METHOD */
-			/* handled specially in Phase 3 */
+
+			/*
+			 * Only do this for partitioned tables, for which this is just a
+			 * catalog change.  Tables with storage are handled by Phase 3.
+			 */
+			if (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE &&
+				tab->chgAccessMethod)
+				ATExecSetAccessMethodNoStorage(rel, tab->newAccessMethod);
 			break;
 		case AT_SetTableSpace:	/* SET TABLESPACE */
 
@@ -5808,7 +5809,7 @@ ATRewriteTables(AlterTableStmt *parsetree, List **wqueue, LOCKMODE lockmode,
 			 * Select destination access method (same as original unless user
 			 * requested a change)
 			 */
-			if (OidIsValid(tab->newAccessMethod))
+			if (tab->chgAccessMethod)
 				NewAccessMethod = tab->newAccessMethod;
 			else
 				NewAccessMethod = OldHeap->rd_rel->relam;
@@ -6395,6 +6396,7 @@ ATGetQueueEntry(List **wqueue, Relation rel)
 	tab->relkind = rel->rd_rel->relkind;
 	tab->oldDesc = CreateTupleDescCopyConstr(RelationGetDescr(rel));
 	tab->newAccessMethod = InvalidOid;
+	tab->chgAccessMethod = false;
 	tab->newTableSpace = InvalidOid;
 	tab->newrelpersistence = RELPERSISTENCE_PERMANENT;
 	tab->chgPersistence = false;
@@ -15197,25 +15199,128 @@ ATExecDropCluster(Relation rel, LOCKMODE lockmode)
 /*
  * Preparation phase for SET ACCESS METHOD
  *
- * Check that access method exists.  If it is the same as the table's current
- * access method, it is a no-op.  Otherwise, a table rewrite is necessary.
- * If amname is NULL, select default_table_access_method as access method.
+ * Check that the access method exists and determine whether a change is
+ * actually needed.
  */
 static void
 ATPrepSetAccessMethod(AlteredTableInfo *tab, Relation rel, const char *amname)
 {
 	Oid			amoid;
 
-	/* Check that the table access method exists */
-	amoid = get_table_am_oid(amname ? amname : default_table_access_method,
-							 false);
+	/*
+	 * Look up the access method name and check that it differs from the
+	 * table's current AM.  If DEFAULT was specified for a partitioned table
+	 * (amname is NULL), set it to InvalidOid to reset the catalogued AM.
+	 */
+	if (amname != NULL)
+		amoid = get_table_am_oid(amname, false);
+	else if (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
+		amoid = InvalidOid;
+	else
+		amoid = get_table_am_oid(default_table_access_method, false);
 
+	/* if it's a match, phase 3 doesn't need to do anything */
 	if (rel->rd_rel->relam == amoid)
 		return;
 
 	/* Save info for Phase 3 to do the real work */
 	tab->rewrite |= AT_REWRITE_ACCESS_METHOD;
 	tab->newAccessMethod = amoid;
+	tab->chgAccessMethod = true;
+}
+
+/*
+ * Special handling of ALTER TABLE SET ACCESS METHOD for relations with no
+ * storage that have an interest in preserving AM.
+ *
+ * Since these have no storage, setting the access method is a catalog only
+ * operation.
+ */
+static void
+ATExecSetAccessMethodNoStorage(Relation rel, Oid newAccessMethodId)
+{
+	Relation	pg_class;
+	Oid			oldAccessMethodId;
+	HeapTuple	tuple;
+	Form_pg_class rd_rel;
+	Oid			reloid = RelationGetRelid(rel);
+
+	/*
+	 * Shouldn't be called on relations having storage; these are processed in
+	 * phase 3.
+	 */
+	Assert(!RELKIND_HAS_STORAGE(rel->rd_rel->relkind));
+
+	/* Get a modifiable copy of the relation's pg_class row. */
+	pg_class = table_open(RelationRelationId, RowExclusiveLock);
+
+	tuple = SearchSysCacheCopy1(RELOID, ObjectIdGetDatum(reloid));
+	if (!HeapTupleIsValid(tuple))
+		elog(ERROR, "cache lookup failed for relation %u", reloid);
+	rd_rel = (Form_pg_class) GETSTRUCT(tuple);
+
+	/* Update the pg_class row. */
+	oldAccessMethodId = rd_rel->relam;
+	rd_rel->relam = newAccessMethodId;
+
+	/* Leave if no update required */
+	if (rd_rel->relam == oldAccessMethodId)
+	{
+		heap_freetuple(tuple);
+		table_close(pg_class, RowExclusiveLock);
+		return;
+	}
+
+	CatalogTupleUpdate(pg_class, &tuple->t_self, tuple);
+
+	/*
+	 * Update the dependency on the new access method.  No dependency is added
+	 * if the new access method is InvalidOid (default case).  Be very careful
+	 * that this has to compare the previous value stored in pg_class with the
+	 * new one.
+	 */
+	if (!OidIsValid(oldAccessMethodId) && OidIsValid(rd_rel->relam))
+	{
+		ObjectAddress relobj,
+					referenced;
+
+		/*
+		 * New access method is defined and there was no dependency
+		 * previously, so record a new one.
+		 */
+		ObjectAddressSet(relobj, RelationRelationId, reloid);
+		ObjectAddressSet(referenced, AccessMethodRelationId, rd_rel->relam);
+		recordDependencyOn(&relobj, &referenced, DEPENDENCY_NORMAL);
+	}
+	else if (OidIsValid(oldAccessMethodId) &&
+			 !OidIsValid(rd_rel->relam))
+	{
+		/*
+		 * There was an access method defined, and no new one, so just remove
+		 * the existing dependency.
+		 */
+		deleteDependencyRecordsForClass(RelationRelationId, reloid,
+										AccessMethodRelationId,
+										DEPENDENCY_NORMAL);
+	}
+	else
+	{
+		Assert(OidIsValid(oldAccessMethodId) &&
+			   OidIsValid(rd_rel->relam));
+
+		/* Both are valid, so update the dependency */
+		changeDependencyFor(RelationRelationId, reloid,
+							AccessMethodRelationId,
+							oldAccessMethodId, rd_rel->relam);
+	}
+
+	/* make the relam and dependency changes visible */
+	CommandCounterIncrement();
+
+	InvokeObjectPostAlterHook(RelationRelationId, RelationGetRelid(rel), 0);
+
+	heap_freetuple(tuple);
+	table_close(pg_class, RowExclusiveLock);
 }
 
 /*
diff --git a/src/backend/utils/cache/lsyscache.c b/src/backend/utils/cache/lsyscache.c
index 6418d1c6eb..26368ffcc9 100644
--- a/src/backend/utils/cache/lsyscache.c
+++ b/src/backend/utils/cache/lsyscache.c
@@ -2069,6 +2069,28 @@ get_rel_persistence(Oid relid)
 	return result;
 }
 
+/*
+ * get_rel_relam
+ *
+ *		Returns the relam associated with a given relation.
+ */
+Oid
+get_rel_relam(Oid relid)
+{
+	HeapTuple	tp;
+	Form_pg_class reltup;
+	Oid			result;
+
+	tp = SearchSysCache1(RELOID, ObjectIdGetDatum(relid));
+	if (!HeapTupleIsValid(tp))
+		elog(ERROR, "cache lookup failed for relation %u", relid);
+	reltup = (Form_pg_class) GETSTRUCT(tp);
+	result = reltup->relam;
+	ReleaseSysCache(tp);
+
+	return result;
+}
+
 
 /*				---------- TRANSFORM CACHE ----------						 */
 
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 2cd19d603f..404861856f 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -1208,6 +1208,13 @@ retry:
 	else if (RELKIND_HAS_TABLE_AM(relation->rd_rel->relkind) ||
 			 relation->rd_rel->relkind == RELKIND_SEQUENCE)
 		RelationInitTableAccessMethod(relation);
+	else if (relation->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
+	{
+		/*
+		 * Do nothing: access methods are a setting that partitions can
+		 * inherit.
+		 */
+	}
 	else
 		Assert(relation->rd_rel->relam == InvalidOid);
 
diff --git a/src/bin/pg_dump/pg_dump.c b/src/bin/pg_dump/pg_dump.c
index d275b31605..b1c4c3ec7f 100644
--- a/src/bin/pg_dump/pg_dump.c
+++ b/src/bin/pg_dump/pg_dump.c
@@ -16656,7 +16656,8 @@ dumpTableSchema(Archive *fout, const TableInfo *tbinfo)
 		if (RELKIND_HAS_TABLESPACE(tbinfo->relkind))
 			tablespace = tbinfo->reltablespace;
 
-		if (RELKIND_HAS_TABLE_AM(tbinfo->relkind))
+		if (RELKIND_HAS_TABLE_AM(tbinfo->relkind) ||
+			tbinfo->relkind == RELKIND_PARTITIONED_TABLE)
 			tableam = tbinfo->amname;
 
 		ArchiveEntry(fout, tbinfo->dobj.catId, tbinfo->dobj.dumpId,
diff --git a/src/bin/pg_dump/t/002_pg_dump.pl b/src/bin/pg_dump/t/002_pg_dump.pl
index c8b489d94e..f0410ce6a1 100644
--- a/src/bin/pg_dump/t/002_pg_dump.pl
+++ b/src/bin/pg_dump/t/002_pg_dump.pl
@@ -4587,6 +4587,41 @@ my %tests = (
 			no_table_access_method => 1,
 			only_dump_measurement => 1,
 		},
+	},
+
+	# CREATE TABLE with partitioned table and various AMs.  One
+	# partition uses the same default as the parent, and a second
+	# uses its own AM.
+	'CREATE TABLE regress_pg_dump_table_part' => {
+		create_order => 19,
+		create_sql => '
+			CREATE TABLE dump_test.regress_pg_dump_table_am_parent (id int) PARTITION BY LIST (id);
+			ALTER TABLE dump_test.regress_pg_dump_table_am_parent SET ACCESS METHOD regress_table_am;
+			CREATE TABLE dump_test.regress_pg_dump_table_am_child_1
+			  PARTITION OF dump_test.regress_pg_dump_table_am_parent FOR VALUES IN (1) USING heap;
+			CREATE TABLE dump_test.regress_pg_dump_table_am_child_2
+			  PARTITION OF dump_test.regress_pg_dump_table_am_parent FOR VALUES IN (2);',
+		regexp => qr/^
+			\QSET default_table_access_method = regress_table_am;\E
+			(\n(?!SET[^;]+;)[^\n]*)*
+			\n\QCREATE TABLE dump_test.regress_pg_dump_table_am_parent (\E
+			(.*\n)*
+			\QSET default_table_access_method = heap;\E
+			(\n(?!SET[^;]+;)[^\n]*)*
+			\n\QCREATE TABLE dump_test.regress_pg_dump_table_am_child_1 (\E
+			(.*\n)*
+			\QSET default_table_access_method = regress_table_am;\E
+			(\n(?!SET[^;]+;)[^\n]*)*
+			\n\QCREATE TABLE dump_test.regress_pg_dump_table_am_child_2 (\E
+			(.*\n)*/xm,
+		like => {
+			%full_runs, %dump_test_schema_runs, section_pre_data => 1,
+		},
+		unlike => {
+			exclude_dump_test_schema => 1,
+			no_table_access_method => 1,
+			only_dump_measurement => 1,
+		},
 	});
 
 #########################################
diff --git a/src/include/catalog/pg_class.h b/src/include/catalog/pg_class.h
index 3b7533e7bb..0fc2c093b0 100644
--- a/src/include/catalog/pg_class.h
+++ b/src/include/catalog/pg_class.h
@@ -219,7 +219,9 @@ MAKE_SYSCACHE(RELNAMENSP, pg_class_relname_nsp_index, 128);
 /*
  * Relation kinds with a table access method (rd_tableam).  Although sequences
  * use the heap table AM, they are enough of a special case in most uses that
- * they are not included here.
+ * they are not included here.  Likewise, partitioned tables can have an access
+ * method defined so that their partitions can inherit it, but they do not set
+ * rd_tableam; hence, this is handled specially outside of this macro.
  */
 #define RELKIND_HAS_TABLE_AM(relkind) \
 	((relkind) == RELKIND_RELATION || \
diff --git a/src/include/utils/lsyscache.h b/src/include/utils/lsyscache.h
index e4a200b00e..35a8dec2b9 100644
--- a/src/include/utils/lsyscache.h
+++ b/src/include/utils/lsyscache.h
@@ -139,6 +139,7 @@ extern char get_rel_relkind(Oid relid);
 extern bool get_rel_relispartition(Oid relid);
 extern Oid	get_rel_tablespace(Oid relid);
 extern char get_rel_persistence(Oid relid);
+extern Oid	get_rel_relam(Oid relid);
 extern Oid	get_transform_fromsql(Oid typid, Oid langid, List *trftypes);
 extern Oid	get_transform_tosql(Oid typid, Oid langid, List *trftypes);
 extern bool get_typisdefined(Oid typid);
diff --git a/src/test/regress/expected/create_am.out b/src/test/regress/expected/create_am.out
index 8d73e21356..a27805a8f5 100644
--- a/src/test/regress/expected/create_am.out
+++ b/src/test/regress/expected/create_am.out
@@ -176,9 +176,16 @@ SELECT f1 FROM tableam_tblmv_heap2 ORDER BY f1;
   1
 (1 row)
 
--- CREATE TABLE ..  PARTITION BY doesn't not support USING
+-- CREATE TABLE ..  PARTITION BY supports USING.
 CREATE TABLE tableam_parted_heap2 (a text, b int) PARTITION BY list (a) USING heap2;
-ERROR:  specifying a table access method is not supported on a partitioned table
+SELECT a.amname FROM pg_class c, pg_am a
+  WHERE c.relname = 'tableam_parted_heap2' AND a.oid = c.relam;
+ amname 
+--------
+ heap2
+(1 row)
+
+DROP TABLE tableam_parted_heap2;
 CREATE TABLE tableam_parted_heap2 (a text, b int) PARTITION BY list (a);
 -- new partitions will inherit from the current default, rather the partition root
 SET default_table_access_method = 'heap';
@@ -336,12 +343,151 @@ ALTER MATERIALIZED VIEW heapmv SET ACCESS METHOD heap, SET ACCESS METHOD heap2;
 ERROR:  cannot have multiple SET ACCESS METHOD subcommands
 DROP MATERIALIZED VIEW heapmv;
 DROP TABLE heaptable;
--- No support for partitioned tables.
-CREATE TABLE am_partitioned(x INT, y INT)
-  PARTITION BY hash (x);
+-- Partition hierarchies with access methods
+BEGIN;
+SET LOCAL default_table_access_method = 'heap';
+CREATE TABLE am_partitioned(x INT, y INT) PARTITION BY hash (x);
+-- pg_class.relam is 0, no dependency recorded between the AM and the
+-- partitioned table.
+SELECT relam FROM pg_class WHERE relname = 'am_partitioned';
+ relam 
+-------
+     0
+(1 row)
+
+SELECT pg_describe_object(classid, objid, objsubid) AS obj,
+       pg_describe_object(refclassid, refobjid, refobjsubid) as refobj
+  FROM pg_depend, pg_am
+  WHERE pg_depend.refclassid = 'pg_am'::regclass
+    AND pg_am.oid = pg_depend.refobjid
+    AND pg_depend.objid = 'am_partitioned'::regclass;
+ obj | refobj 
+-----+--------
+(0 rows)
+
+-- New default is set, with dependency added.
 ALTER TABLE am_partitioned SET ACCESS METHOD heap2;
-ERROR:  cannot change access method of a partitioned table
+SELECT a.amname FROM pg_class c, pg_am a
+  WHERE c.relname = 'am_partitioned' AND a.oid = c.relam;
+ amname 
+--------
+ heap2
+(1 row)
+
+SELECT pg_describe_object(classid, objid, objsubid) AS obj,
+       pg_describe_object(refclassid, refobjid, refobjsubid) as refobj
+  FROM pg_depend, pg_am
+  WHERE pg_depend.refclassid = 'pg_am'::regclass
+    AND pg_am.oid = pg_depend.refobjid
+    AND pg_depend.objid = 'am_partitioned'::regclass;
+         obj          |       refobj        
+----------------------+---------------------
+ table am_partitioned | access method heap2
+(1 row)
+
+-- Default is set, with dependency updated.
+SET LOCAL default_table_access_method = 'heap2';
+ALTER TABLE am_partitioned SET ACCESS METHOD heap;
+SELECT a.amname FROM pg_class c, pg_am a
+  WHERE c.relname = 'am_partitioned' AND a.oid = c.relam;
+ amname 
+--------
+ heap
+(1 row)
+
+-- Dependency pinned, hence removed.
+SELECT pg_describe_object(classid, objid, objsubid) AS obj,
+       pg_describe_object(refclassid, refobjid, refobjsubid) as refobj
+  FROM pg_depend, pg_am
+  WHERE pg_depend.refclassid = 'pg_am'::regclass
+    AND pg_am.oid = pg_depend.refobjid
+    AND pg_depend.objid = 'am_partitioned'::regclass;
+ obj | refobj 
+-----+--------
+(0 rows)
+
+-- Default and AM set in the clause are the same, relam should be set.
+SET LOCAL default_table_access_method = 'heap2';
+ALTER TABLE am_partitioned SET ACCESS METHOD heap2;
+SELECT a.amname FROM pg_class c, pg_am a
+  WHERE c.relname = 'am_partitioned' AND a.oid = c.relam;
+ amname 
+--------
+ heap2
+(1 row)
+
+-- Reset to default
+ALTER TABLE am_partitioned SET ACCESS METHOD DEFAULT;
+SELECT relam FROM pg_class WHERE relname = 'am_partitioned';
+ relam 
+-------
+     0
+(1 row)
+
+-- Upon ALTER TABLE SET ACCESS METHOD on a partitioned table, new partitions
+-- will inherit the AM set.  Existing partitioned are unchanged.
+SELECT relam FROM pg_class WHERE relname = 'am_partitioned';
+ relam 
+-------
+     0
+(1 row)
+
+SET LOCAL default_table_access_method = 'heap';
+CREATE TABLE am_partitioned_0 PARTITION OF am_partitioned
+  FOR VALUES WITH (MODULUS 10, REMAINDER 0);
+SET LOCAL default_table_access_method = 'heap2';
+CREATE TABLE am_partitioned_1 PARTITION OF am_partitioned
+  FOR VALUES WITH (MODULUS 10, REMAINDER 1);
+SET LOCAL default_table_access_method = 'heap';
+ALTER TABLE am_partitioned SET ACCESS METHOD heap2;
+CREATE TABLE am_partitioned_2 PARTITION OF am_partitioned
+  FOR VALUES WITH (MODULUS 10, REMAINDER 2);
+ALTER TABLE am_partitioned SET ACCESS METHOD DEFAULT;
+SELECT relam FROM pg_class WHERE relname = 'am_partitioned';
+ relam 
+-------
+     0
+(1 row)
+
+CREATE TABLE am_partitioned_3 PARTITION OF am_partitioned
+  FOR VALUES WITH (MODULUS 10, REMAINDER 3);
+-- Partitioned table with relam at 0
+ALTER TABLE am_partitioned SET ACCESS METHOD DEFAULT;
+CREATE TABLE am_partitioned_5p PARTITION OF am_partitioned
+  FOR VALUES WITH (MODULUS 10, REMAINDER 5) PARTITION BY hash(y);
+-- Partitions of this partitioned table inherit default AM at creation
+-- time.
+CREATE TABLE am_partitioned_5p1 PARTITION OF am_partitioned_5p
+  FOR VALUES WITH (MODULUS 10, REMAINDER 1);
+-- Partitioned table with relam set.
+ALTER TABLE am_partitioned SET ACCESS METHOD heap2;
+CREATE TABLE am_partitioned_6p PARTITION OF am_partitioned
+  FOR VALUES WITH (MODULUS 10, REMAINDER 6) PARTITION BY hash(y);
+-- Partitions of this partitioned table inherit its AM.
+CREATE TABLE am_partitioned_6p1 PARTITION OF am_partitioned_6p
+  FOR VALUES WITH (MODULUS 10, REMAINDER 1);
+SELECT c.relname, a.amname FROM pg_class c, pg_am a
+  WHERE c.relam = a.oid AND
+        c.relname LIKE 'am_partitioned%'
+UNION ALL
+SELECT c.relname, 'default' FROM pg_class c
+  WHERE c.relam = 0
+        AND c.relname LIKE 'am_partitioned%' ORDER BY 1;
+      relname       | amname  
+--------------------+---------
+ am_partitioned     | heap2
+ am_partitioned_0   | heap
+ am_partitioned_1   | heap2
+ am_partitioned_2   | heap2
+ am_partitioned_3   | heap
+ am_partitioned_5p  | default
+ am_partitioned_5p1 | heap
+ am_partitioned_6p  | heap2
+ am_partitioned_6p1 | heap2
+(9 rows)
+
 DROP TABLE am_partitioned;
+COMMIT;
 -- Second, create objects in the new AM by changing the default AM
 BEGIN;
 SET LOCAL default_table_access_method = 'heap2';
diff --git a/src/test/regress/sql/create_am.sql b/src/test/regress/sql/create_am.sql
index 606ee4cb24..adff0079f9 100644
--- a/src/test/regress/sql/create_am.sql
+++ b/src/test/regress/sql/create_am.sql
@@ -124,8 +124,11 @@ CREATE SEQUENCE tableam_seq_heap2 USING heap2;
 CREATE MATERIALIZED VIEW tableam_tblmv_heap2 USING heap2 AS SELECT * FROM tableam_tbl_heap2;
 SELECT f1 FROM tableam_tblmv_heap2 ORDER BY f1;
 
--- CREATE TABLE ..  PARTITION BY doesn't not support USING
+-- CREATE TABLE ..  PARTITION BY supports USING.
 CREATE TABLE tableam_parted_heap2 (a text, b int) PARTITION BY list (a) USING heap2;
+SELECT a.amname FROM pg_class c, pg_am a
+  WHERE c.relname = 'tableam_parted_heap2' AND a.oid = c.relam;
+DROP TABLE tableam_parted_heap2;
 
 CREATE TABLE tableam_parted_heap2 (a text, b int) PARTITION BY list (a);
 -- new partitions will inherit from the current default, rather the partition root
@@ -213,11 +216,91 @@ ALTER TABLE heaptable SET ACCESS METHOD DEFAULT, SET ACCESS METHOD heap2;
 ALTER MATERIALIZED VIEW heapmv SET ACCESS METHOD heap, SET ACCESS METHOD heap2;
 DROP MATERIALIZED VIEW heapmv;
 DROP TABLE heaptable;
--- No support for partitioned tables.
-CREATE TABLE am_partitioned(x INT, y INT)
-  PARTITION BY hash (x);
+
+-- Partition hierarchies with access methods
+BEGIN;
+SET LOCAL default_table_access_method = 'heap';
+CREATE TABLE am_partitioned(x INT, y INT) PARTITION BY hash (x);
+-- pg_class.relam is 0, no dependency recorded between the AM and the
+-- partitioned table.
+SELECT relam FROM pg_class WHERE relname = 'am_partitioned';
+SELECT pg_describe_object(classid, objid, objsubid) AS obj,
+       pg_describe_object(refclassid, refobjid, refobjsubid) as refobj
+  FROM pg_depend, pg_am
+  WHERE pg_depend.refclassid = 'pg_am'::regclass
+    AND pg_am.oid = pg_depend.refobjid
+    AND pg_depend.objid = 'am_partitioned'::regclass;
+-- New default is set, with dependency added.
 ALTER TABLE am_partitioned SET ACCESS METHOD heap2;
+SELECT a.amname FROM pg_class c, pg_am a
+  WHERE c.relname = 'am_partitioned' AND a.oid = c.relam;
+SELECT pg_describe_object(classid, objid, objsubid) AS obj,
+       pg_describe_object(refclassid, refobjid, refobjsubid) as refobj
+  FROM pg_depend, pg_am
+  WHERE pg_depend.refclassid = 'pg_am'::regclass
+    AND pg_am.oid = pg_depend.refobjid
+    AND pg_depend.objid = 'am_partitioned'::regclass;
+-- Default is set, with dependency updated.
+SET LOCAL default_table_access_method = 'heap2';
+ALTER TABLE am_partitioned SET ACCESS METHOD heap;
+SELECT a.amname FROM pg_class c, pg_am a
+  WHERE c.relname = 'am_partitioned' AND a.oid = c.relam;
+-- Dependency pinned, hence removed.
+SELECT pg_describe_object(classid, objid, objsubid) AS obj,
+       pg_describe_object(refclassid, refobjid, refobjsubid) as refobj
+  FROM pg_depend, pg_am
+  WHERE pg_depend.refclassid = 'pg_am'::regclass
+    AND pg_am.oid = pg_depend.refobjid
+    AND pg_depend.objid = 'am_partitioned'::regclass;
+-- Default and AM set in the clause are the same, relam should be set.
+SET LOCAL default_table_access_method = 'heap2';
+ALTER TABLE am_partitioned SET ACCESS METHOD heap2;
+SELECT a.amname FROM pg_class c, pg_am a
+  WHERE c.relname = 'am_partitioned' AND a.oid = c.relam;
+-- Reset to default
+ALTER TABLE am_partitioned SET ACCESS METHOD DEFAULT;
+SELECT relam FROM pg_class WHERE relname = 'am_partitioned';
+-- Upon ALTER TABLE SET ACCESS METHOD on a partitioned table, new partitions
+-- will inherit the AM set.  Existing partitioned are unchanged.
+SELECT relam FROM pg_class WHERE relname = 'am_partitioned';
+SET LOCAL default_table_access_method = 'heap';
+CREATE TABLE am_partitioned_0 PARTITION OF am_partitioned
+  FOR VALUES WITH (MODULUS 10, REMAINDER 0);
+SET LOCAL default_table_access_method = 'heap2';
+CREATE TABLE am_partitioned_1 PARTITION OF am_partitioned
+  FOR VALUES WITH (MODULUS 10, REMAINDER 1);
+SET LOCAL default_table_access_method = 'heap';
+ALTER TABLE am_partitioned SET ACCESS METHOD heap2;
+CREATE TABLE am_partitioned_2 PARTITION OF am_partitioned
+  FOR VALUES WITH (MODULUS 10, REMAINDER 2);
+ALTER TABLE am_partitioned SET ACCESS METHOD DEFAULT;
+SELECT relam FROM pg_class WHERE relname = 'am_partitioned';
+CREATE TABLE am_partitioned_3 PARTITION OF am_partitioned
+  FOR VALUES WITH (MODULUS 10, REMAINDER 3);
+-- Partitioned table with relam at 0
+ALTER TABLE am_partitioned SET ACCESS METHOD DEFAULT;
+CREATE TABLE am_partitioned_5p PARTITION OF am_partitioned
+  FOR VALUES WITH (MODULUS 10, REMAINDER 5) PARTITION BY hash(y);
+-- Partitions of this partitioned table inherit default AM at creation
+-- time.
+CREATE TABLE am_partitioned_5p1 PARTITION OF am_partitioned_5p
+  FOR VALUES WITH (MODULUS 10, REMAINDER 1);
+-- Partitioned table with relam set.
+ALTER TABLE am_partitioned SET ACCESS METHOD heap2;
+CREATE TABLE am_partitioned_6p PARTITION OF am_partitioned
+  FOR VALUES WITH (MODULUS 10, REMAINDER 6) PARTITION BY hash(y);
+-- Partitions of this partitioned table inherit its AM.
+CREATE TABLE am_partitioned_6p1 PARTITION OF am_partitioned_6p
+  FOR VALUES WITH (MODULUS 10, REMAINDER 1);
+SELECT c.relname, a.amname FROM pg_class c, pg_am a
+  WHERE c.relam = a.oid AND
+        c.relname LIKE 'am_partitioned%'
+UNION ALL
+SELECT c.relname, 'default' FROM pg_class c
+  WHERE c.relam = 0
+        AND c.relname LIKE 'am_partitioned%' ORDER BY 1;
 DROP TABLE am_partitioned;
+COMMIT;
 
 -- Second, create objects in the new AM by changing the default AM
 BEGIN;
-- 
2.39.2

