From 4c066c98690544460863f85e944c50400013bb7d Mon Sep 17 00:00:00 2001
From: Zhao Junwang <zhjwpku@gmail.com>
Date: Fri, 12 Jul 2024 10:04:29 +0000
Subject: [PATCH v2 1/2] support specify tablespace for each merged/split

It might be a good idea that we can specify the tablespace
for each merged/split partition.

Signed-off-by: Zhao Junwang <zhjwpku@gmail.com>
---
 doc/src/sgml/ref/alter_table.sgml             |  6 ++---
 src/backend/commands/tablecmds.c              | 16 ++++++-----
 src/backend/parser/gram.y                     |  6 +++--
 src/include/nodes/parsenodes.h                |  2 ++
 src/test/regress/expected/partition_merge.out | 25 +++++++++++++++++
 src/test/regress/expected/partition_split.out | 27 +++++++++++++++++++
 src/test/regress/sql/partition_merge.sql      | 14 ++++++++++
 src/test/regress/sql/partition_split.sql      | 14 ++++++++++
 8 files changed, 98 insertions(+), 12 deletions(-)

diff --git a/doc/src/sgml/ref/alter_table.sgml b/doc/src/sgml/ref/alter_table.sgml
index 6a2822adad..259e8d9402 100644
--- a/doc/src/sgml/ref/alter_table.sgml
+++ b/doc/src/sgml/ref/alter_table.sgml
@@ -39,11 +39,11 @@ ALTER TABLE [ IF EXISTS ] <replaceable class="parameter">name</replaceable>
     DETACH PARTITION <replaceable class="parameter">partition_name</replaceable> [ CONCURRENTLY | FINALIZE ]
 ALTER TABLE [ IF EXISTS ] <replaceable class="parameter">name</replaceable>
     SPLIT PARTITION <replaceable class="parameter">partition_name</replaceable> INTO
-        (PARTITION <replaceable class="parameter">partition_name1</replaceable> { FOR VALUES <replaceable class="parameter">partition_bound_spec</replaceable> | DEFAULT },
-         PARTITION <replaceable class="parameter">partition_name2</replaceable> { FOR VALUES <replaceable class="parameter">partition_bound_spec</replaceable> | DEFAULT } [, ...])
+        (PARTITION <replaceable class="parameter">partition_name1</replaceable> { FOR VALUES <replaceable class="parameter">partition_bound_spec</replaceable> | DEFAULT } [ TABLESPACE <replaceable class="parameter">tablespace_name</replaceable> ],
+         PARTITION <replaceable class="parameter">partition_name2</replaceable> { FOR VALUES <replaceable class="parameter">partition_bound_spec</replaceable> | DEFAULT } [ TABLESPACE <replaceable class="parameter">tablespace_name</replaceable> ] [, ...])
 ALTER TABLE [ IF EXISTS ] <replaceable class="parameter">name</replaceable>
     MERGE PARTITIONS (<replaceable class="parameter">partition_name1</replaceable>, <replaceable class="parameter">partition_name2</replaceable> [, ...])
-        INTO <replaceable class="parameter">partition_name</replaceable>
+        INTO <replaceable class="parameter">partition_name</replaceable> [ TABLESPACE <replaceable class="parameter">tablespace_name</replaceable> ]
 
 <phrase>where <replaceable class="parameter">action</replaceable> is one of:</phrase>
 
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 0b2a52463f..d5f11b7601 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -20315,15 +20315,16 @@ moveSplitTableRows(Relation rel, Relation splitRel, List *partlist, List *newPar
  * Emulates command: CREATE [TEMP] TABLE <newPartName> (LIKE <modelRel's name>
  * INCLUDING ALL EXCLUDING INDEXES EXCLUDING IDENTITY EXCLUDING STATISTICS)
  *
- * Also, this function sets the new partition access method same as parent
- * table access methods (similarly to CREATE TABLE ... PARTITION OF).  It
+ * The new partition will be created in the same tablespace as the parent if not
+ * specified.  Also, this function sets the new partition access method same as
+ * parent table access methods (similarly to CREATE TABLE ... PARTITION OF).  It
  * checks that parent and child tables have compatible persistence.
  *
  * Function returns the created relation (locked in AccessExclusiveLock mode).
  */
 static Relation
-createPartitionTable(RangeVar *newPartName, Relation modelRel,
-					 AlterTableUtilityContext *context)
+createPartitionTable(RangeVar *newPartName, char *tablespacename,
+					 Relation modelRel, AlterTableUtilityContext *context)
 {
 	CreateStmt *createStmt;
 	TableLikeClause *tlc;
@@ -20347,7 +20348,8 @@ createPartitionTable(RangeVar *newPartName, Relation modelRel,
 	createStmt->constraints = NIL;
 	createStmt->options = NIL;
 	createStmt->oncommit = ONCOMMIT_NOOP;
-	createStmt->tablespacename = get_tablespace_name(modelRel->rd_rel->reltablespace);
+	createStmt->tablespacename = tablespacename ? tablespacename :
+		get_tablespace_name(modelRel->rd_rel->reltablespace);;
 	createStmt->if_not_exists = false;
 	createStmt->accessMethod = get_am_name(modelRel->rd_rel->relam);
 
@@ -20506,7 +20508,7 @@ ATExecSplitPartition(List **wqueue, AlteredTableInfo *tab, Relation rel,
 		SinglePartitionSpec *sps = (SinglePartitionSpec *) lfirst(listptr);
 		Relation	newPartRel;
 
-		newPartRel = createPartitionTable(sps->name, rel, context);
+		newPartRel = createPartitionTable(sps->name, sps->tablespacename, rel, context);
 		newPartRels = lappend(newPartRels, newPartRel);
 	}
 
@@ -20750,7 +20752,7 @@ ATExecMergePartitions(List **wqueue, AlteredTableInfo *tab, Relation rel,
 	}
 
 	/* Create table for new partition, use partitioned table as model. */
-	newPartRel = createPartitionTable(cmd->name, rel, context);
+	newPartRel = createPartitionTable(cmd->name, cmd->tablespacename, rel, context);
 
 	/* Copy data from merged partitions to new partition. */
 	moveMergedTablesRows(rel, mergingPartitionsList, newPartRel);
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index a043fd4c66..95bfe638ab 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -2317,12 +2317,13 @@ partitions_list:
 		;
 
 SinglePartitionSpec:
-			PARTITION qualified_name PartitionBoundSpec
+			PARTITION qualified_name PartitionBoundSpec OptTableSpace
 				{
 					SinglePartitionSpec *n = makeNode(SinglePartitionSpec);
 
 					n->name = $2;
 					n->bound = $3;
+					n->tablespacename = $4;
 
 					$$ = n;
 				}
@@ -2387,7 +2388,7 @@ partition_cmd:
 					$$ = (Node *) n;
 				}
 			/* ALTER TABLE <name> MERGE PARTITIONS () INTO <partition_name> */
-			| MERGE PARTITIONS '(' qualified_name_list ')' INTO qualified_name
+			| MERGE PARTITIONS '(' qualified_name_list ')' INTO qualified_name OptTableSpace
 				{
 					AlterTableCmd *n = makeNode(AlterTableCmd);
 					PartitionCmd *cmd = makeNode(PartitionCmd);
@@ -2397,6 +2398,7 @@ partition_cmd:
 					cmd->bound = NULL;
 					cmd->partlist = $4;
 					cmd->concurrent = false;
+					cmd->tablespacename = $8;
 					n->def = (Node *) cmd;
 					$$ = (Node *) n;
 				}
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 85a62b538e..9ccdcb9eff 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -945,6 +945,7 @@ typedef struct SinglePartitionSpec
 	NodeTag		type;
 
 	RangeVar   *name;			/* name of partition */
+	char	   *tablespacename; /* name of tablespace, or NULL for default */
 	PartitionBoundSpec *bound;	/* FOR VALUES, if attaching */
 } SinglePartitionSpec;
 
@@ -959,6 +960,7 @@ typedef struct PartitionCmd
 	List	   *partlist;		/* list of partitions, for MERGE/SPLIT
 								 * PARTITION command */
 	bool		concurrent;
+	char	   *tablespacename; /* name of tablespace, or NULL for default */
 } PartitionCmd;
 
 /****************************************************************************
diff --git a/src/test/regress/expected/partition_merge.out b/src/test/regress/expected/partition_merge.out
index 59836e2d35..c28836a241 100644
--- a/src/test/regress/expected/partition_merge.out
+++ b/src/test/regress/expected/partition_merge.out
@@ -886,6 +886,31 @@ SELECT tablename, indexname, tablespace FROM pg_indexes
  tp_0_2    | tp_0_2_pkey | regress_tblspace
 (2 rows)
 
+DROP TABLE t;
+-- Check the merged partition can be set to a different tablespace
+SET search_path = partitions_merge_schema, public;
+CREATE TABLE t (i int PRIMARY KEY) PARTITION BY RANGE (i);
+CREATE TABLE tp_0_1 PARTITION OF t FOR VALUES FROM (0) TO (1);
+CREATE TABLE tp_1_2 PARTITION OF t FOR VALUES FROM (1) TO (2);
+ALTER TABLE t MERGE PARTITIONS (tp_0_1, tp_1_2) INTO tp_0_2 TABLESPACE regress_tblspace;
+SELECT tablename, tablespace FROM pg_tables
+  WHERE tablename IN ('t', 'tp_0_2') AND schemaname = 'partitions_merge_schema'
+  ORDER BY tablename, tablespace;
+ tablename |    tablespace    
+-----------+------------------
+ t         | 
+ tp_0_2    | regress_tblspace
+(2 rows)
+
+SELECT tablename, indexname, tablespace FROM pg_indexes
+  WHERE tablename IN ('t', 'tp_0_2') AND schemaname = 'partitions_merge_schema'
+  ORDER BY tablename, indexname, tablespace;
+ tablename |  indexname  | tablespace 
+-----------+-------------+------------
+ t         | t_pkey      | 
+ tp_0_2    | tp_0_2_pkey | 
+(2 rows)
+
 DROP TABLE t;
 -- Check the new partition inherits parent's table access method
 SET search_path = partitions_merge_schema, public;
diff --git a/src/test/regress/expected/partition_split.out b/src/test/regress/expected/partition_split.out
index dc9a5130cc..d3853e4bef 100644
--- a/src/test/regress/expected/partition_split.out
+++ b/src/test/regress/expected/partition_split.out
@@ -1521,6 +1521,33 @@ SELECT tablename, indexname, tablespace FROM pg_indexes
  tp_1_2    | tp_1_2_pkey | regress_tblspace
 (3 rows)
 
+DROP TABLE t;
+-- Check the split partitions can be set to a different tablespace
+CREATE TABLE t (i int PRIMARY KEY) PARTITION BY RANGE (i);
+CREATE TABLE tp_0_2 PARTITION OF t FOR VALUES FROM (0) TO (2);
+ALTER TABLE t SPLIT PARTITION tp_0_2 INTO
+  (PARTITION tp_0_1 FOR VALUES FROM (0) TO (1) TABLESPACE regress_tblspace,
+   PARTITION tp_1_2 FOR VALUES FROM (1) TO (2));
+SELECT tablename, tablespace FROM pg_tables
+  WHERE tablename IN ('t', 'tp_0_1', 'tp_1_2') AND schemaname = 'partition_split_schema'
+  ORDER BY tablename, tablespace;
+ tablename |    tablespace    
+-----------+------------------
+ t         | 
+ tp_0_1    | regress_tblspace
+ tp_1_2    | 
+(3 rows)
+
+SELECT tablename, indexname, tablespace FROM pg_indexes
+  WHERE tablename IN ('t', 'tp_0_1', 'tp_1_2') AND schemaname = 'partition_split_schema'
+  ORDER BY tablename, indexname, tablespace;
+ tablename |  indexname  | tablespace 
+-----------+-------------+------------
+ t         | t_pkey      | 
+ tp_0_1    | tp_0_1_pkey | 
+ tp_1_2    | tp_1_2_pkey | 
+(3 rows)
+
 DROP TABLE t;
 -- Check new partitions inherits parent's table access method
 CREATE ACCESS METHOD partition_split_heap TYPE TABLE HANDLER heap_tableam_handler;
diff --git a/src/test/regress/sql/partition_merge.sql b/src/test/regress/sql/partition_merge.sql
index bede819af9..741c9696ec 100644
--- a/src/test/regress/sql/partition_merge.sql
+++ b/src/test/regress/sql/partition_merge.sql
@@ -551,6 +551,20 @@ SELECT tablename, indexname, tablespace FROM pg_indexes
   ORDER BY tablename, indexname, tablespace;
 DROP TABLE t;
 
+-- Check the merged partition can be set to a different tablespace
+SET search_path = partitions_merge_schema, public;
+CREATE TABLE t (i int PRIMARY KEY) PARTITION BY RANGE (i);
+CREATE TABLE tp_0_1 PARTITION OF t FOR VALUES FROM (0) TO (1);
+CREATE TABLE tp_1_2 PARTITION OF t FOR VALUES FROM (1) TO (2);
+ALTER TABLE t MERGE PARTITIONS (tp_0_1, tp_1_2) INTO tp_0_2 TABLESPACE regress_tblspace;
+SELECT tablename, tablespace FROM pg_tables
+  WHERE tablename IN ('t', 'tp_0_2') AND schemaname = 'partitions_merge_schema'
+  ORDER BY tablename, tablespace;
+SELECT tablename, indexname, tablespace FROM pg_indexes
+  WHERE tablename IN ('t', 'tp_0_2') AND schemaname = 'partitions_merge_schema'
+  ORDER BY tablename, indexname, tablespace;
+DROP TABLE t;
+
 -- Check the new partition inherits parent's table access method
 SET search_path = partitions_merge_schema, public;
 CREATE ACCESS METHOD partitions_merge_heap TYPE TABLE HANDLER heap_tableam_handler;
diff --git a/src/test/regress/sql/partition_split.sql b/src/test/regress/sql/partition_split.sql
index ef5ea07f74..aaaee1d876 100644
--- a/src/test/regress/sql/partition_split.sql
+++ b/src/test/regress/sql/partition_split.sql
@@ -895,6 +895,20 @@ SELECT tablename, indexname, tablespace FROM pg_indexes
   ORDER BY tablename, indexname, tablespace;
 DROP TABLE t;
 
+-- Check the split partitions can be set to a different tablespace
+CREATE TABLE t (i int PRIMARY KEY) PARTITION BY RANGE (i);
+CREATE TABLE tp_0_2 PARTITION OF t FOR VALUES FROM (0) TO (2);
+ALTER TABLE t SPLIT PARTITION tp_0_2 INTO
+  (PARTITION tp_0_1 FOR VALUES FROM (0) TO (1) TABLESPACE regress_tblspace,
+   PARTITION tp_1_2 FOR VALUES FROM (1) TO (2));
+SELECT tablename, tablespace FROM pg_tables
+  WHERE tablename IN ('t', 'tp_0_1', 'tp_1_2') AND schemaname = 'partition_split_schema'
+  ORDER BY tablename, tablespace;
+SELECT tablename, indexname, tablespace FROM pg_indexes
+  WHERE tablename IN ('t', 'tp_0_1', 'tp_1_2') AND schemaname = 'partition_split_schema'
+  ORDER BY tablename, indexname, tablespace;
+DROP TABLE t;
+
 -- Check new partitions inherits parent's table access method
 CREATE ACCESS METHOD partition_split_heap TYPE TABLE HANDLER heap_tableam_handler;
 CREATE TABLE t (i int) PARTITION BY RANGE (i) USING partition_split_heap;
-- 
2.39.2

