From d7d2cd8adce2e4949494bbc65e888681bd9f7da5 Mon Sep 17 00:00:00 2001
From: "dgrowley@gmail.com" <dgrowley@gmail.com>
Date: Fri, 23 Nov 2018 15:58:38 +1300
Subject: [PATCH v2] Allow newly created partitions to inherit their parent's
 tablespace

Previously a partitioned table's tablespace option was never stored in the
catalog and partitions which were newly created would always default to
the default tablespace.  Here we change things so that the tablespace
option for a partitioned table is stored in the catalog.  This is used to
determine which tablespace any newly created partitions which are attached
to the partitioned table are stored.  Of course, the tablespace of newly
created partitions can still be overwritten with the TABLESPACE option
during the CREATE TABLE, we're only changing the behavior when the
tablespace is not otherwise specified.
---
 doc/src/sgml/ref/create_table.sgml        |  6 ++++-
 src/backend/catalog/heap.c                | 13 ---------
 src/backend/commands/tablecmds.c          | 44 ++++++++++++++++++++++++-------
 src/test/regress/input/tablespace.source  | 11 ++++++++
 src/test/regress/output/tablespace.source | 17 ++++++++++++
 5 files changed, 68 insertions(+), 23 deletions(-)

diff --git a/doc/src/sgml/ref/create_table.sgml b/doc/src/sgml/ref/create_table.sgml
index 50d5597002..073c24fbbe 100644
--- a/doc/src/sgml/ref/create_table.sgml
+++ b/doc/src/sgml/ref/create_table.sgml
@@ -1215,7 +1215,11 @@ WITH ( MODULUS <replaceable class="parameter">numeric_literal</replaceable>, REM
       of the tablespace in which the new table is to be created.
       If not specified,
       <xref linkend="guc-default-tablespace"/> is consulted, or
-      <xref linkend="guc-temp-tablespaces"/> if the table is temporary.
+      <xref linkend="guc-temp-tablespaces"/> if the table is temporary.  For
+      partitioned tables, since no storage is required for the table itself,
+      the tablespace specified here only serves to mark the default tablespace
+      for any newly created partitions when no other tablespace is explicitly
+      specified.
      </para>
     </listitem>
    </varlistentry>
diff --git a/src/backend/catalog/heap.c b/src/backend/catalog/heap.c
index 11debaa780..b718fba540 100644
--- a/src/backend/catalog/heap.c
+++ b/src/backend/catalog/heap.c
@@ -334,20 +334,7 @@ heap_create(const char *relname,
 		case RELKIND_COMPOSITE_TYPE:
 		case RELKIND_FOREIGN_TABLE:
 		case RELKIND_PARTITIONED_TABLE:
-			create_storage = false;
-
-			/*
-			 * Force reltablespace to zero if the relation has no physical
-			 * storage.  This is mainly just for cleanliness' sake.
-			 */
-			reltablespace = InvalidOid;
-			break;
-
 		case RELKIND_PARTITIONED_INDEX:
-			/*
-			 * Preserve tablespace so that it's used as tablespace for indexes
-			 * on future partitions.
-			 */
 			create_storage = false;
 			break;
 
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 843ed48aa7..3d1b5b561f 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -446,7 +446,7 @@ static bool ATPrepChangePersistence(Relation rel, bool toLogged);
 static void ATPrepSetTableSpace(AlteredTableInfo *tab, Relation rel,
 					const char *tablespacename, LOCKMODE lockmode);
 static void ATExecSetTableSpace(Oid tableOid, Oid newTableSpace, LOCKMODE lockmode);
-static void ATExecPartedIdxSetTableSpace(Relation rel, Oid newTableSpace);
+static void ATExecSetTableSpaceNoStorage(Relation rel, Oid newTableSpace);
 static void ATExecSetRelOptions(Relation rel, List *defList,
 					AlterTableType operation,
 					LOCKMODE lockmode);
@@ -588,6 +588,28 @@ DefineRelation(CreateStmt *stmt, char relkind, Oid ownerId,
 	{
 		tablespaceId = get_tablespace_oid(stmt->tablespacename, false);
 	}
+	else if (stmt->partbound)
+	{
+		RangeVar   *parent;
+		Relation	parentrel;
+
+		/*
+		 * For partitions, when no other tablespace is specified, we default
+		 * the tablespace to the parent partitioned table's.
+		 */
+		Assert(list_length(stmt->inhRelations) == 1);
+
+		parent = linitial(stmt->inhRelations);
+
+		parentrel = heap_openrv(parent, AccessExclusiveLock);
+
+		if (OidIsValid(parentrel->rd_rel->reltablespace))
+			tablespaceId = parentrel->rd_rel->reltablespace;
+		else
+			tablespaceId = GetDefaultTablespace(stmt->relation->relpersistence);
+
+		relation_close(parentrel, NoLock);
+	}
 	else
 	{
 		tablespaceId = GetDefaultTablespace(stmt->relation->relpersistence);
@@ -4147,11 +4169,13 @@ ATExecCmd(List **wqueue, AlteredTableInfo *tab, Relation rel,
 			break;
 		case AT_SetTableSpace:	/* SET TABLESPACE */
 			/*
-			 * Only do this for partitioned indexes, for which this is just
-			 * a catalog change.  Other relation types are handled by Phase 3.
+			 * Only do this for partitioned tables and indexes, for which this
+			 * is just a catalog change.  Other relation types are handled by
+			 * Phase 3.
 			 */
-			if (rel->rd_rel->relkind == RELKIND_PARTITIONED_INDEX)
-				ATExecPartedIdxSetTableSpace(rel, tab->newTableSpace);
+			if (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE ||
+				rel->rd_rel->relkind == RELKIND_PARTITIONED_INDEX)
+				ATExecSetTableSpaceNoStorage(rel, tab->newTableSpace);
 
 			break;
 		case AT_SetRelOptions:	/* SET (...) */
@@ -10925,11 +10949,12 @@ ATExecSetTableSpace(Oid tableOid, Oid newTableSpace, LOCKMODE lockmode)
 }
 
 /*
- * Special handling of ALTER TABLE SET TABLESPACE for partitioned indexes,
- * which have no storage (so not handled in Phase 3 like other relation types)
+ * Special handling of ALTER TABLE SET TABLESPACE for partitioned tables and
+ * indexes. These have no storage (so not handled in Phase 3 like other
+ * relation types)
  */
 static void
-ATExecPartedIdxSetTableSpace(Relation rel, Oid newTableSpace)
+ATExecSetTableSpaceNoStorage(Relation rel, Oid newTableSpace)
 {
 	HeapTuple	tuple;
 	Oid			oldTableSpace;
@@ -10937,7 +10962,8 @@ ATExecPartedIdxSetTableSpace(Relation rel, Oid newTableSpace)
 	Form_pg_class rd_rel;
 	Oid			indexOid = RelationGetRelid(rel);
 
-	Assert(rel->rd_rel->relkind == RELKIND_PARTITIONED_INDEX);
+	Assert(rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE ||
+		   rel->rd_rel->relkind == RELKIND_PARTITIONED_INDEX);
 
 	/* Can't allow a non-shared relation in pg_global */
 	if (newTableSpace == GLOBALTABLESPACE_OID)
diff --git a/src/test/regress/input/tablespace.source b/src/test/regress/input/tablespace.source
index 2ac757cfab..96ee3917f6 100644
--- a/src/test/regress/input/tablespace.source
+++ b/src/test/regress/input/tablespace.source
@@ -44,6 +44,17 @@ CREATE INDEX foo_idx on testschema.foo(i) TABLESPACE regress_tblspace;
 SELECT relname, spcname FROM pg_catalog.pg_tablespace t, pg_catalog.pg_class c
     where c.reltablespace = t.oid AND c.relname = 'foo_idx';
 
+-- partitioned table
+CREATE TABLE testschema.part (a int) PARTITION BY LIST (a) TABLESPACE regress_tblspace;
+CREATE TABLE testschema.part1 PARTITION OF testschema.part FOR VALUES IN (1);
+ALTER TABLE testschema.part SET TABLESPACE pg_default;
+CREATE TABLE testschema.part2 PARTITION OF testschema.part FOR VALUES IN (2);
+-- Ensure part1 defaulted to regress_tblspace and part2 defaulted to pg_default.
+SELECT relname, spcname FROM pg_catalog.pg_class c
+    LEFT JOIN pg_catalog.pg_tablespace t ON c.reltablespace = t.oid
+    where c.relname LIKE 'part%' order by relname;
+DROP TABLE testschema.part;
+
 -- partitioned index
 CREATE TABLE testschema.part (a int) PARTITION BY LIST (a);
 CREATE TABLE testschema.part1 PARTITION OF testschema.part FOR VALUES IN (1);
diff --git a/src/test/regress/output/tablespace.source b/src/test/regress/output/tablespace.source
index 2e78e5ece6..436b42c8f6 100644
--- a/src/test/regress/output/tablespace.source
+++ b/src/test/regress/output/tablespace.source
@@ -61,6 +61,23 @@ SELECT relname, spcname FROM pg_catalog.pg_tablespace t, pg_catalog.pg_class c
  foo_idx | regress_tblspace
 (1 row)
 
+-- partitioned table
+CREATE TABLE testschema.part (a int) PARTITION BY LIST (a) TABLESPACE regress_tblspace;
+CREATE TABLE testschema.part1 PARTITION OF testschema.part FOR VALUES IN (1);
+ALTER TABLE testschema.part SET TABLESPACE pg_default;
+CREATE TABLE testschema.part2 PARTITION OF testschema.part FOR VALUES IN (2);
+-- Ensure part1 defaulted to regress_tblspace and part2 defaulted to pg_default.
+SELECT relname, spcname FROM pg_catalog.pg_class c
+    LEFT JOIN pg_catalog.pg_tablespace t ON c.reltablespace = t.oid
+    where c.relname LIKE 'part%' order by relname;
+ relname |     spcname      
+---------+------------------
+ part    | 
+ part1   | regress_tblspace
+ part2   | 
+(3 rows)
+
+DROP TABLE testschema.part;
 -- partitioned index
 CREATE TABLE testschema.part (a int) PARTITION BY LIST (a);
 CREATE TABLE testschema.part1 PARTITION OF testschema.part FOR VALUES IN (1);
-- 
2.16.2.windows.1

