From 29f0196fdfe1c80b3e37e9ed5fe2bf7534401e6c Mon Sep 17 00:00:00 2001
From: songjinzhou <49380232+TsinghuaLucky912@users.noreply.github.com>
Date: Fri, 2 Dec 2022 00:45:32 +0800
Subject: [PATCH] Add support for SET ACCESS METHOD in altering a distributed
 table (#6525)

Co-authored-by: TsinghuaLucky912 <postgres@localhost.localdomain>
---
 src/backend/distributed/commands/table.c |  4 ++++
 src/test/regress/expected/pg15.out       | 30 ++++++++++++++++++++++++
 src/test/regress/sql/pg15.sql            | 15 ++++++++++++
 3 files changed, 49 insertions(+)

diff --git a/src/backend/distributed/commands/table.c b/src/backend/distributed/commands/table.c
index 1be61cbe2..8101dc964 100644
--- a/src/backend/distributed/commands/table.c
+++ b/src/backend/distributed/commands/table.c
@@ -2992,6 +2992,9 @@ ErrorIfUnsupportedAlterTableStmt(AlterTableStmt *alterTableStatement)
 				break;
 			}
 
+#if PG_VERSION_NUM >= PG_VERSION_15
+			case AT_SetAccessMethod:
+#endif
 			case AT_SetNotNull:
 			case AT_ReplicaIdentity:
 			case AT_ChangeOwner:
@@ -3007,6 +3010,7 @@ ErrorIfUnsupportedAlterTableStmt(AlterTableStmt *alterTableStatement)
 				{
 					/*
 					 * We will not perform any special check for:
+					 * ALTER TABLE .. SET ACCESS METHOD ..
 					 * ALTER TABLE .. ALTER COLUMN .. SET NOT NULL
 					 * ALTER TABLE .. REPLICA IDENTITY ..
 					 * ALTER TABLE .. VALIDATE CONSTRAINT ..
diff --git a/src/test/regress/expected/pg15.out b/src/test/regress/expected/pg15.out
index be14f70a7..5dcd2c04d 100644
--- a/src/test/regress/expected/pg15.out
+++ b/src/test/regress/expected/pg15.out
@@ -1456,6 +1456,36 @@ NOTICE:  drop cascades to 2 other objects
 CREATE DATABASE db_with_oid OID 987654;
 NOTICE:  Citus partially supports CREATE DATABASE for distributed databases
 DROP DATABASE db_with_oid;
+-- SET ACCESS METHOD
+-- Create a heap2 table am handler with heapam handler
+CREATE ACCESS METHOD heap2 TYPE TABLE HANDLER heap_tableam_handler;
+SELECT run_command_on_workers($$CREATE ACCESS METHOD heap2 TYPE TABLE HANDLER heap_tableam_handler$$);
+           run_command_on_workers
+---------------------------------------------------------------------
+ (localhost,57637,t,"CREATE ACCESS METHOD")
+ (localhost,57638,t,"CREATE ACCESS METHOD")
+(2 rows)
+
+CREATE TABLE mx_ddl_table2 (
+    key int primary key,
+    value int
+);
+SELECT create_distributed_table('mx_ddl_table2', 'key', 'hash', shard_count=> 4);
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ALTER TABLE mx_ddl_table2 SET ACCESS METHOD heap2;
+DROP TABLE mx_ddl_table2;
+DROP ACCESS METHOD heap2;
+SELECT run_command_on_workers($$DROP ACCESS METHOD heap2$$);
+          run_command_on_workers
+---------------------------------------------------------------------
+ (localhost,57637,t,"DROP ACCESS METHOD")
+ (localhost,57638,t,"DROP ACCESS METHOD")
+(2 rows)
+
 -- Clean up
 \set VERBOSITY terse
 SET client_min_messages TO ERROR;
diff --git a/src/test/regress/sql/pg15.sql b/src/test/regress/sql/pg15.sql
index 04dc46910..c334671bc 100644
--- a/src/test/regress/sql/pg15.sql
+++ b/src/test/regress/sql/pg15.sql
@@ -924,6 +924,21 @@ DROP SERVER foreign_server CASCADE;
 CREATE DATABASE db_with_oid OID 987654;
 DROP DATABASE db_with_oid;
 
+-- SET ACCESS METHOD
+-- Create a heap2 table am handler with heapam handler
+CREATE ACCESS METHOD heap2 TYPE TABLE HANDLER heap_tableam_handler;
+SELECT run_command_on_workers($$CREATE ACCESS METHOD heap2 TYPE TABLE HANDLER heap_tableam_handler$$);
+CREATE TABLE mx_ddl_table2 (
+    key int primary key,
+    value int
+);
+SELECT create_distributed_table('mx_ddl_table2', 'key', 'hash', shard_count=> 4);
+ALTER TABLE mx_ddl_table2 SET ACCESS METHOD heap2;
+
+DROP TABLE mx_ddl_table2;
+DROP ACCESS METHOD heap2;
+SELECT run_command_on_workers($$DROP ACCESS METHOD heap2$$);
+
 -- Clean up
 \set VERBOSITY terse
 SET client_min_messages TO ERROR;