v3-0002-support-specify-tablespace-for-each-merged-split-.patch

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

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

As suggested by Sawada.

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

diff --git a/doc/src/sgml/ref/alter_table.sgml b/doc/src/sgml/ref/alter_table.sgml
index 6a2822adad7..d7f64b684d4 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 tablespace_name ],
+         PARTITION <replaceable class="parameter">partition_name2</replaceable> { FOR VALUES <replaceable class="parameter">partition_bound_spec</replaceable> | DEFAULT } [ TABLESPACE tablespace_name ] [, ...])
 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 tablespace_name ]
 
 <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 9e1eff19e07..f21884b83c2 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -20306,8 +20306,8 @@ moveSplitTableRows(Relation rel, Relation splitRel, List *partlist, List *newPar
  * 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;
@@ -20331,7 +20331,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);
 
@@ -20490,7 +20491,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);
 	}
 
@@ -20734,7 +20735,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 a043fd4c669..95bfe638ab4 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 85a62b538e5..4255e00d84f 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 26bf58b23db..62ed62e1ab7 100644
--- a/src/test/regress/expected/partition_merge.out
+++ b/src/test/regress/expected/partition_merge.out
@@ -883,6 +883,28 @@ 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
+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') 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') 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 1a8c95ad819..aa9e8abda23 100644
--- a/src/test/regress/expected/partition_split.out
+++ b/src/test/regress/expected/partition_split.out
@@ -1519,6 +1519,31 @@ 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') 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') 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 200bd3e7620..4f0e1c01fed 100644
--- a/src/test/regress/sql/partition_merge.sql
+++ b/src/test/regress/sql/partition_merge.sql
@@ -548,6 +548,17 @@ SELECT tablename, indexname, tablespace FROM pg_indexes
   WHERE tablename IN ('t', 'tp_0_2') ORDER BY tablename, indexname, tablespace;
 DROP TABLE t;
 
+-- Check the merged partition can be set to a different tablespace
+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') ORDER BY tablename, tablespace;
+SELECT tablename, indexname, tablespace FROM pg_indexes
+  WHERE tablename IN ('t', 'tp_0_2') 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 5143d85ec01..964c6ce2d61 100644
--- a/src/test/regress/sql/partition_split.sql
+++ b/src/test/regress/sql/partition_split.sql
@@ -893,6 +893,18 @@ SELECT tablename, indexname, tablespace FROM pg_indexes
   WHERE tablename IN ('t', 'tp_0_1', 'tp_1_2') 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') ORDER BY tablename, tablespace;
+SELECT tablename, indexname, tablespace FROM pg_indexes
+  WHERE tablename IN ('t', 'tp_0_1', 'tp_1_2') 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