From 1fe7305114c5e13fece12c468210b91f9079da34 Mon Sep 17 00:00:00 2001
From: Alvaro Herrera <alvherre@alvh.no-ip.org>
Date: Mon, 11 Jan 2021 20:02:06 -0300
Subject: [PATCH] Prevent drop of tablespace used by partitioned relation

Bug #16577
---
 doc/src/sgml/catalogs.sgml                | 13 ++++-
 src/backend/catalog/heap.c                |  9 +++
 src/backend/catalog/pg_shdepend.c         | 68 ++++++++++++++++++++---
 src/backend/commands/tablecmds.c          |  4 ++
 src/backend/commands/tablespace.c         | 12 ++++
 src/include/catalog/dependency.h          | 13 +++++
 src/test/regress/input/tablespace.source  |  3 +
 src/test/regress/output/tablespace.source |  5 ++
 8 files changed, 117 insertions(+), 10 deletions(-)

diff --git a/doc/src/sgml/catalogs.sgml b/doc/src/sgml/catalogs.sgml
index 3a2266526c..43d7a1ad90 100644
--- a/doc/src/sgml/catalogs.sgml
+++ b/doc/src/sgml/catalogs.sgml
@@ -6849,10 +6849,21 @@ SCRAM-SHA-256$<replaceable>&lt;iteration count&gt;</replaceable>:<replaceable>&l
       </para>
      </listitem>
     </varlistentry>
+
+    <varlistentry>
+     <term><symbol>SHARED_DEPENDENCY_TABLESPACE</symbol> (<literal>t</literal>)</term>
+     <listitem>
+      <para>
+       The referenced object (which must be a tablespace) is mentioned as
+       the tablespace for a relation that doesn't have storage.
+      </para>
+     </listitem>
+    </varlistentry>
    </variablelist>
 
    Other dependency flavors might be needed in future.  Note in particular
-   that the current definition only supports roles as referenced objects.
+   that the current definition only supports roles and tablespaces as referenced
+   objects.
   </para>
 
  </sect1>
diff --git a/src/backend/catalog/heap.c b/src/backend/catalog/heap.c
index 21f2240ade..9abc4a1f55 100644
--- a/src/backend/catalog/heap.c
+++ b/src/backend/catalog/heap.c
@@ -440,6 +440,15 @@ heap_create(const char *relname,
 		}
 	}
 
+	/*
+	 * If a tablespace is specified, removal of that tablespace is normally
+	 * protected by the existence of a physical file; but for relations with
+	 * no files, add a pg_shdepend entry to account for that.
+	 */
+	if (!create_storage && reltablespace != InvalidOid)
+		recordDependencyOnTablespace(RelationRelationId, relid,
+									 reltablespace);
+
 	return rel;
 }
 
diff --git a/src/backend/catalog/pg_shdepend.c b/src/backend/catalog/pg_shdepend.c
index 86c3be107f..90b7a5de29 100644
--- a/src/backend/catalog/pg_shdepend.c
+++ b/src/backend/catalog/pg_shdepend.c
@@ -59,6 +59,7 @@
 #include "commands/schemacmds.h"
 #include "commands/subscriptioncmds.h"
 #include "commands/tablecmds.h"
+#include "commands/tablespace.h"
 #include "commands/typecmds.h"
 #include "miscadmin.h"
 #include "storage/lmgr.h"
@@ -186,11 +187,14 @@ recordDependencyOnOwner(Oid classId, Oid objectId, Oid owner)
  *
  * There must be no more than one existing entry for the given dependent
  * object and dependency type!	So in practice this can only be used for
- * updating SHARED_DEPENDENCY_OWNER entries, which should have that property.
+ * updating SHARED_DEPENDENCY_OWNER and SHARED_DEPENDENCY_TABLESPACE
+ * entries, which should have that property.
  *
  * If there is no previous entry, we assume it was referencing a PINned
  * object, so we create a new entry.  If the new referenced object is
  * PINned, we don't create an entry (and drop the old one, if any).
+ * (For tablespaces, we don't record dependencies in certain cases, so
+ * there are other possible reasons for entries to be missing.)
  *
  * sdepRel must be the pg_shdepend relation, already opened and suitably
  * locked.
@@ -344,6 +348,58 @@ changeDependencyOnOwner(Oid classId, Oid objectId, Oid newOwnerId)
 	table_close(sdepRel, RowExclusiveLock);
 }
 
+/*
+ * recordDependencyOnTablespace
+ *
+ * A convenient wrapper of recordSharedDependencyOn -- register the specified
+ * tablespace as default for the given object.
+ *
+ * Note: it's the caller's responsibility to ensure that there isn't a
+ * tablespace entry for the object already.
+ */
+void
+recordDependencyOnTablespace(Oid classId, Oid objectId, Oid tablespace)
+{
+	ObjectAddress myself,
+				  referenced;
+
+	ObjectAddressSet(myself, classId, objectId);
+	ObjectAddressSet(referenced, TableSpaceRelationId, tablespace);
+
+	recordSharedDependencyOn(&myself, &referenced,
+							 SHARED_DEPENDENCY_TABLESPACE);
+}
+
+/*
+ * changeDependencyOnTablespace
+ *
+ * Update the shared dependencies to account for the new tablespace.
+ *
+ * Note: we don't need an objsubid argument because only whole objects
+ * have tablespaces.
+ */
+void
+changeDependencyOnTablespace(Oid classId, Oid objectId, Oid newTablespaceId)
+{
+	Relation	sdepRel;
+
+	sdepRel = table_open(SharedDependRelationId, RowExclusiveLock);
+
+	if (newTablespaceId != DEFAULTTABLESPACE_OID &&
+		newTablespaceId != InvalidOid)
+		shdepChangeDep(sdepRel,
+					   classId, objectId, 0,
+					   TableSpaceRelationId, newTablespaceId,
+					   SHARED_DEPENDENCY_TABLESPACE);
+	else
+		shdepDropDependency(sdepRel,
+							classId, objectId, 0, true,
+							InvalidOid, InvalidOid,
+							SHARED_DEPENDENCY_INVALID);
+
+	table_close(sdepRel, RowExclusiveLock);
+}
+
 /*
  * getOidListDiff
  *		Helper for updateAclDependencies.
@@ -1121,13 +1177,6 @@ shdepLockAndCheckObject(Oid classId, Oid objectId)
 								objectId)));
 			break;
 
-			/*
-			 * Currently, this routine need not support any other shared
-			 * object types besides roles.  If we wanted to record explicit
-			 * dependencies on databases or tablespaces, we'd need code along
-			 * these lines:
-			 */
-#ifdef NOT_USED
 		case TableSpaceRelationId:
 			{
 				/* For lack of a syscache on pg_tablespace, do this: */
@@ -1141,7 +1190,6 @@ shdepLockAndCheckObject(Oid classId, Oid objectId)
 				pfree(tablespace);
 				break;
 			}
-#endif
 
 		case DatabaseRelationId:
 			{
@@ -1201,6 +1249,8 @@ storeObjectDescription(StringInfo descs,
 				appendStringInfo(descs, _("privileges for %s"), objdesc);
 			else if (deptype == SHARED_DEPENDENCY_POLICY)
 				appendStringInfo(descs, _("target of %s"), objdesc);
+			else if (deptype == SHARED_DEPENDENCY_TABLESPACE)
+				appendStringInfo(descs, _("tablespace for %s"), objdesc);
 			else
 				elog(ERROR, "unrecognized dependency type: %d",
 					 (int) deptype);
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 993da56d43..fd55bf4ac7 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -13340,6 +13340,10 @@ ATExecSetTableSpaceNoStorage(Relation rel, Oid newTableSpace)
 	rd_rel->reltablespace = (newTableSpace == MyDatabaseTableSpace) ? InvalidOid : newTableSpace;
 	CatalogTupleUpdate(pg_class, &tuple->t_self, tuple);
 
+	/* Record dependency on tablespace */
+	changeDependencyOnTablespace(RelationRelationId,
+								 reloid, rd_rel->reltablespace);
+
 	InvokeObjectPostAlterHook(RelationRelationId, reloid, 0);
 
 	heap_freetuple(tuple);
diff --git a/src/backend/commands/tablespace.c b/src/backend/commands/tablespace.c
index a9a2db2834..69ea155d50 100644
--- a/src/backend/commands/tablespace.c
+++ b/src/backend/commands/tablespace.c
@@ -420,6 +420,8 @@ DropTableSpace(DropTableSpaceStmt *stmt)
 	Form_pg_tablespace spcform;
 	ScanKeyData entry[1];
 	Oid			tablespaceoid;
+	char	   *detail;
+	char	   *detail_log;
 
 	/*
 	 * Find the target tuple
@@ -468,6 +470,16 @@ DropTableSpace(DropTableSpaceStmt *stmt)
 		aclcheck_error(ACLCHECK_NO_PRIV, OBJECT_TABLESPACE,
 					   tablespacename);
 
+	/* Check for pg_shdepend entries depending on this tablespace */
+	if (checkSharedDependencies(TableSpaceRelationId, tablespaceoid,
+								&detail, &detail_log))
+		ereport(ERROR,
+				(errcode(ERRCODE_DEPENDENT_OBJECTS_STILL_EXIST),
+				 errmsg("tablespace \"%s\" cannot be dropped because some objects depend on it",
+						tablespacename),
+				 errdetail_internal("%s", detail),
+				 errdetail_log("%s", detail_log)));
+
 	/* DROP hook for the tablespace being removed */
 	InvokeObjectDropHook(TableSpaceRelationId, tablespaceoid, 0);
 
diff --git a/src/include/catalog/dependency.h b/src/include/catalog/dependency.h
index 6d6c23d846..6065b0312f 100644
--- a/src/include/catalog/dependency.h
+++ b/src/include/catalog/dependency.h
@@ -67,6 +67,12 @@ typedef enum DependencyType
  * a role mentioned in a policy object.  The referenced object must be a
  * pg_authid entry.
  *
+ * (e) a SHARED_DEPENDENCY_TABLESPACE entry means that the referenced
+ * object is a tablespace mentioned in a relation without storage.  The
+ * referenced object must be a pg_tablespace entry.  (Relations that have
+ * storage don't need this: they are protected by the existence of a physical
+ * file in the tablespace.)
+ *
  * SHARED_DEPENDENCY_INVALID is a value used as a parameter in internal
  * routines, and is not valid in the catalog itself.
  */
@@ -76,6 +82,7 @@ typedef enum SharedDependencyType
 	SHARED_DEPENDENCY_OWNER = 'o',
 	SHARED_DEPENDENCY_ACL = 'a',
 	SHARED_DEPENDENCY_POLICY = 'r',
+	SHARED_DEPENDENCY_TABLESPACE = 't',
 	SHARED_DEPENDENCY_INVALID = 0
 } SharedDependencyType;
 
@@ -253,6 +260,12 @@ extern void recordDependencyOnOwner(Oid classId, Oid objectId, Oid owner);
 extern void changeDependencyOnOwner(Oid classId, Oid objectId,
 									Oid newOwnerId);
 
+extern void recordDependencyOnTablespace(Oid classId, Oid objectId,
+										 Oid tablespace);
+
+extern void changeDependencyOnTablespace(Oid classId, Oid objectId,
+										 Oid newTablespaceId);
+
 extern void updateAclDependencies(Oid classId, Oid objectId, int32 objectSubId,
 								  Oid ownerId,
 								  int noldmembers, Oid *oldmembers,
diff --git a/src/test/regress/input/tablespace.source b/src/test/regress/input/tablespace.source
index a5f61a35dc..1a181016d7 100644
--- a/src/test/regress/input/tablespace.source
+++ b/src/test/regress/input/tablespace.source
@@ -249,6 +249,9 @@ CREATE TABLESPACE regress_badspace LOCATION '/no/such/location';
 -- No such tablespace
 CREATE TABLE bar (i int) TABLESPACE regress_nosuchspace;
 
+-- Fail, in use for some partitioned object
+DROP TABLESPACE regress_tblspace;
+ALTER INDEX testschema.part_a_idx SET TABLESPACE pg_default;
 -- Fail, not empty
 DROP TABLESPACE regress_tblspace;
 
diff --git a/src/test/regress/output/tablespace.source b/src/test/regress/output/tablespace.source
index 162b591b31..94c5f023c6 100644
--- a/src/test/regress/output/tablespace.source
+++ b/src/test/regress/output/tablespace.source
@@ -712,6 +712,11 @@ ERROR:  directory "/no/such/location" does not exist
 -- No such tablespace
 CREATE TABLE bar (i int) TABLESPACE regress_nosuchspace;
 ERROR:  tablespace "regress_nosuchspace" does not exist
+-- Fail, in use for some partitioned object
+DROP TABLESPACE regress_tblspace;
+ERROR:  tablespace "regress_tblspace" cannot be dropped because some objects depend on it
+DETAIL:  tablespace for index testschema.part_a_idx
+ALTER INDEX testschema.part_a_idx SET TABLESPACE pg_default;
 -- Fail, not empty
 DROP TABLESPACE regress_tblspace;
 ERROR:  tablespace "regress_tblspace" is not empty
-- 
2.20.1

