From 389a93c2fa7b65f0bd63166f46238e8cb0bf1b5a 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 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             | 27 ++++++++++++-------
 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, 113 insertions(+), 18 deletions(-)

diff --git a/doc/src/sgml/ref/alter_table.sgml b/doc/src/sgml/ref/alter_table.sgml
index 6a2822adad..d3e2f7dbfb 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>
 
@@ -1125,7 +1125,7 @@ WITH ( MODULUS <replaceable class="parameter">numeric_literal</replaceable>, REM
    </varlistentry>
 
    <varlistentry id="sql-altertable-split-partition">
-    <term><literal>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 } [, ...])</literal></term>
+    <term><literal>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 } [ 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> ] [, ...])</literal></term>
 
     <listitem>
      <para>
@@ -1163,8 +1163,8 @@ WITH ( MODULUS <replaceable class="parameter">numeric_literal</replaceable>, REM
       New partitions will have the same table access method as the parent.
       If the parent table is persistent then new partitions are created
       persistent.  If the parent table is temporary then new partitions
-      are also created temporary.  New partitions will also be created in
-      the same tablespace as the parent.
+      are also created temporary.  New partitions will be created in
+      the same tablespace as the parent if <literal>TABLESPACE</literal> not specified.
      </para>
      <note>
       <para>
@@ -1177,7 +1177,7 @@ WITH ( MODULUS <replaceable class="parameter">numeric_literal</replaceable>, REM
    </varlistentry>
 
    <varlistentry id="sql-altertable-merge-partitions">
-    <term><literal>MERGE PARTITIONS (<replaceable class="parameter">partition_name1</replaceable>, <replaceable class="parameter">partition_name2</replaceable> [, ...]) INTO <replaceable class="parameter">partition_name</replaceable></literal></term>
+    <term><literal>MERGE PARTITIONS (<replaceable class="parameter">partition_name1</replaceable>, <replaceable class="parameter">partition_name2</replaceable> [, ...]) INTO <replaceable class="parameter">partition_name</replaceable> [ TABLESPACE <replaceable class="parameter">tablespace_name</replaceable> ]</literal></term>
 
     <listitem>
      <para>
@@ -1236,8 +1236,8 @@ WITH ( MODULUS <replaceable class="parameter">numeric_literal</replaceable>, REM
       The new partition will have the same table access method as the parent.
       If the parent table is persistent then the new partition is created
       persistent.  If the parent table is temporary then the new partition
-      is also created temporary.  The new partition will also be created in
-      the same tablespace as the parent.
+      is also created temporary.  The new partition will be created in the
+      same tablespace as the parent if <literal>TABLESPACE</literal> not specified.
      </para>
      <note>
       <para>
@@ -1516,6 +1516,15 @@ WITH ( MODULUS <replaceable class="parameter">numeric_literal</replaceable>, REM
       </listitem>
      </varlistentry>
 
+     <varlistentry id="sql-altertable-parms-tablespace-name">
+      <term><replaceable class="parameter">tablespace_name</replaceable></term>
+      <listitem>
+       <para>
+        The name of the tablespace in which the split or merged partition will be created.
+       </para>
+      </listitem>
+     </varlistentry>
+
     </variablelist>
  </refsect1>
 
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

