Pass proper old rlocator to smgr when creating relations and indexes (#156)

contrib/pg_tde/Makefile

@@ -25,7 +25,8 @@ alter_index_basic \
 merge_join_basic \
 tablespace_basic \
 relocate \
-version
+version \
+recreate_storage
 TAP_TESTS = 1
 
 OBJS = src/encryption/enc_tde.o \

contrib/pg_tde/expected/recreate_storage.out

@@ -0,0 +1,134 @@
+CREATE EXTENSION IF NOT EXISTS pg_tde;
+SELECT pg_tde_add_key_provider_file('file-vault','/tmp/pg_tde_test_keyring.per');
+ pg_tde_add_key_provider_file 
+------------------------------
+                            1
+(1 row)
+
+SELECT pg_tde_set_principal_key('test-db-principal-key','file-vault');
+ pg_tde_set_principal_key 
+--------------------------
+ t
+(1 row)
+
+SET default_table_access_method = "tde_heap";
+CREATE TABLE t1(n integer);
+SELECT pg_tde_is_encrypted('t1');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+VACUUM FULL t1;
+SELECT pg_tde_is_encrypted('t1');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+CREATE TABLE test_tab1 AS SELECT generate_series(1,10) a;
+CREATE INDEX test_idx1 ON test_tab1(a);
+SELECT pg_tde_is_encrypted('test_tab1');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+SELECT pg_tde_is_encrypted('test_idx1');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+REINDEX index CONCURRENTLY test_idx1;
+SELECT pg_tde_is_encrypted('test_tab1');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+SELECT pg_tde_is_encrypted('test_idx1');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+CREATE TABLE mvtest_t (id int NOT NULL PRIMARY KEY, type text NOT NULL, amt numeric NOT NULL);
+INSERT INTO mvtest_t VALUES
+  (1, 'x', 2),
+  (2, 'x', 3),
+  (3, 'y', 5),
+  (4, 'y', 7),
+  (5, 'z', 11);
+CREATE MATERIALIZED VIEW mvtest_tm AS SELECT type, sum(amt) AS totamt FROM mvtest_t GROUP BY type WITH NO DATA;
+SELECT pg_tde_is_encrypted('mvtest_tm');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+REFRESH MATERIALIZED VIEW mvtest_tm;
+SELECT pg_tde_is_encrypted('mvtest_tm');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+CREATE TYPE rewritetype AS (a int);
+CREATE TABLE rewritemetoo1 OF rewritetype;
+CREATE TABLE rewritemetoo2 OF rewritetype;
+SELECT pg_tde_is_encrypted('rewritemetoo1');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+SELECT pg_tde_is_encrypted('rewritemetoo2');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+ALTER TYPE rewritetype ALTER ATTRIBUTE a TYPE text cascade;
+SELECT pg_tde_is_encrypted('rewritemetoo1');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+SELECT pg_tde_is_encrypted('rewritemetoo2');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+CREATE TABLE encrypted_table (
+    id SERIAL,
+    data TEXT,
+    created_at DATE NOT NULL,
+    PRIMARY KEY (id, created_at)
+) USING tde_heap;
+CREATE INDEX idx_date ON encrypted_table (created_at);
+SELECT pg_tde_is_encrypted('encrypted_table');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+CLUSTER encrypted_table USING idx_date;
+SELECT pg_tde_is_encrypted('encrypted_table');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
+DROP EXTENSION pg_tde CASCADE;
+NOTICE:  drop cascades to 7 other objects
+DETAIL:  drop cascades to table t1
+drop cascades to table test_tab1
+drop cascades to table mvtest_t
+drop cascades to materialized view mvtest_tm
+drop cascades to table rewritemetoo1
+drop cascades to table rewritemetoo2
+drop cascades to table encrypted_table
+RESET default_table_access_method;

contrib/pg_tde/meson.build

@@ -108,6 +108,7 @@ sql_tests = [
       'update_basic',
       'key_provider',
       'relocate',
+      'recreate_storage',
 ]
 
 tap_tests = [

contrib/pg_tde/sql/recreate_storage.sql

@@ -0,0 +1,54 @@
+CREATE EXTENSION IF NOT EXISTS pg_tde;
+
+SELECT pg_tde_add_key_provider_file('file-vault','/tmp/pg_tde_test_keyring.per');
+SELECT pg_tde_set_principal_key('test-db-principal-key','file-vault');
+
+SET default_table_access_method = "tde_heap";
+
+CREATE TABLE t1(n integer);
+SELECT pg_tde_is_encrypted('t1');
+VACUUM FULL t1;
+SELECT pg_tde_is_encrypted('t1');
+
+CREATE TABLE test_tab1 AS SELECT generate_series(1,10) a;
+CREATE INDEX test_idx1 ON test_tab1(a);
+SELECT pg_tde_is_encrypted('test_tab1');
+SELECT pg_tde_is_encrypted('test_idx1');
+REINDEX index CONCURRENTLY test_idx1;
+SELECT pg_tde_is_encrypted('test_tab1');
+SELECT pg_tde_is_encrypted('test_idx1');
+
+CREATE TABLE mvtest_t (id int NOT NULL PRIMARY KEY, type text NOT NULL, amt numeric NOT NULL);
+INSERT INTO mvtest_t VALUES
+  (1, 'x', 2),
+  (2, 'x', 3),
+  (3, 'y', 5),
+  (4, 'y', 7),
+  (5, 'z', 11);
+CREATE MATERIALIZED VIEW mvtest_tm AS SELECT type, sum(amt) AS totamt FROM mvtest_t GROUP BY type WITH NO DATA;
+SELECT pg_tde_is_encrypted('mvtest_tm');
+REFRESH MATERIALIZED VIEW mvtest_tm;
+SELECT pg_tde_is_encrypted('mvtest_tm');
+
+CREATE TYPE rewritetype AS (a int);
+CREATE TABLE rewritemetoo1 OF rewritetype;
+CREATE TABLE rewritemetoo2 OF rewritetype;
+SELECT pg_tde_is_encrypted('rewritemetoo1');
+SELECT pg_tde_is_encrypted('rewritemetoo2');
+ALTER TYPE rewritetype ALTER ATTRIBUTE a TYPE text cascade;
+SELECT pg_tde_is_encrypted('rewritemetoo1');
+SELECT pg_tde_is_encrypted('rewritemetoo2');
+
+CREATE TABLE encrypted_table (
+    id SERIAL,
+    data TEXT,
+    created_at DATE NOT NULL,
+    PRIMARY KEY (id, created_at)
+) USING tde_heap;
+CREATE INDEX idx_date ON encrypted_table (created_at);
+SELECT pg_tde_is_encrypted('encrypted_table');
+CLUSTER encrypted_table USING idx_date;
+SELECT pg_tde_is_encrypted('encrypted_table');
+
+DROP EXTENSION pg_tde CASCADE;
+RESET default_table_access_method;

src/backend/bootstrap/bootparse.y

@@ -210,7 +210,8 @@ Boot_CreateStmt:
 												   true,
 												   &relfrozenxid,
 												   &relminmxid,
-												   true);
+												   true,
+												   NULL);
 						elog(DEBUG4, "bootstrap relation created");
 					}
 					else

src/backend/catalog/heap.c

@@ -300,7 +300,8 @@ heap_create(const char *relname,
 			bool allow_system_table_mods,
 			TransactionId *relfrozenxid,
 			MultiXactId *relminmxid,
-			bool create_storage)
+			bool create_storage,
+			RelFileLocator *old_rlocator)
 {
 	Relation	rel;
 
@@ -385,14 +386,31 @@ heap_create(const char *relname,
 	 */
 	if (create_storage)
 	{
+		RelFileLocator prev_rlocator = rel->rd_locator;
+		RelFileLocator new_rlocator = rel->rd_locator;
+
+		if (old_rlocator != NULL)
+		{
+			prev_rlocator = *old_rlocator;
+
+			/*
+			 * table_relation_set_new_filelocator() takes old_rlocator from
+			 * rel->rd_locator
+			 */
+			rel->rd_locator = prev_rlocator;
+		}
+
 		if (RELKIND_HAS_TABLE_AM(rel->rd_rel->relkind))
-			table_relation_set_new_filelocator(rel, &rel->rd_locator,
+			table_relation_set_new_filelocator(rel, &new_rlocator,
 											   relpersistence,
 											   relfrozenxid, relminmxid);
 		else if (RELKIND_HAS_STORAGE(rel->rd_rel->relkind))
-			RelationCreateStorage(rel->rd_locator, rel->rd_locator, relpersistence, true);
+			RelationCreateStorage(prev_rlocator, new_rlocator, relpersistence, true);
 		else
 			Assert(false);
+
+		/* restore the orginal rel's locator */
+		rel->rd_locator = new_rlocator;
 	}
 
 	/*
@@ -1129,6 +1147,8 @@ heap_create_with_catalog(const char *relname,
 	Oid			existing_relid;
 	Oid			old_type_oid;
 	Oid			new_type_oid;
+	RelFileLocator *old_rlocator = NULL;
+	Relation	old_rel;
 
 	/* By default set to InvalidOid unless overridden by binary-upgrade */
 	RelFileNumber relfilenumber = InvalidRelFileNumber;
@@ -1283,6 +1303,12 @@ heap_create_with_catalog(const char *relname,
 	else
 		relacl = NULL;
 
+	if (relrewrite != InvalidOid)
+	{
+		old_rel = table_open(relrewrite, AccessShareLock);
+		old_rlocator = &old_rel->rd_locator;
+	}
+
 	/*
 	 * Create the relcache entry (mostly dummy at this point) and the physical
 	 * disk file.  (If we fail further down, it's the smgr's responsibility to
@@ -1306,7 +1332,11 @@ heap_create_with_catalog(const char *relname,
 							   allow_system_table_mods,
 							   &relfrozenxid,
 							   &relminmxid,
-							   true);
+							   true,
+							   old_rlocator);
+
+	if (relrewrite != InvalidOid)
+		table_close(old_rel, AccessShareLock);
 
 	Assert(relid == RelationGetRelid(new_rel_desc));
 

src/backend/catalog/index.c

@@ -741,7 +741,8 @@ index_create(Relation heapRelation,
 			 bits16 constr_flags,
 			 bool allow_system_table_mods,
 			 bool is_internal,
-			 Oid *constraintId)
+			 Oid *constraintId,
+			 RelFileLocator *old_rlocator)
 {
 	Oid			heapRelationId = RelationGetRelid(heapRelation);
 	Relation	pg_class;
@@ -985,7 +986,8 @@ index_create(Relation heapRelation,
 								allow_system_table_mods,
 								&relfrozenxid,
 								&relminmxid,
-								create_storage);
+								create_storage,
+								old_rlocator);
 
 	Assert(relfrozenxid == InvalidTransactionId);
 	Assert(relminmxid == InvalidMultiXactId);
@@ -1459,7 +1461,8 @@ index_concurrently_create_copy(Relation heapRelation, Oid oldIndexId,
 							  0,
 							  true, /* allow table to be a system catalog? */
 							  false,	/* is_internal? */
-							  NULL);
+							  NULL,
+							  &indexRelation->rd_locator);
 
 	/* Close the relations used and clean up */
 	index_close(indexRelation, NoLock);

src/backend/catalog/toasting.c

@@ -325,7 +325,7 @@ create_toast_table(Relation rel, Oid toastOid, Oid toastIndexOid,
 				 BTREE_AM_OID,
 				 rel->rd_rel->reltablespace,
 				 collationIds, opclassIds, NULL, coloptions, NULL, (Datum) 0,
-				 INDEX_CREATE_IS_PRIMARY, 0, true, true, NULL);
+				 INDEX_CREATE_IS_PRIMARY, 0, true, true, NULL, NULL);
 
 	table_close(toast_rel, NoLock);
 

src/backend/commands/indexcmds.c

@@ -1208,7 +1208,7 @@ DefineIndex(Oid tableId,
 					 coloptions, NULL, reloptions,
 					 flags, constr_flags,
 					 allowSystemTableMods, !check_rights,
-					 &createdConstraintId);
+					 &createdConstraintId, NULL);
 
 	ObjectAddressSet(address, RelationRelationId, indexRelationId);
 

src/include/catalog/heap.h

@@ -17,6 +17,7 @@
 #include "catalog/indexing.h"
 #include "catalog/objectaddress.h"
 #include "parser/parse_node.h"
+#include "storage/relfilelocator.h"
 
 
 /* flag bits for CheckAttributeType/CheckAttributeNamesTypes */
@@ -60,7 +61,8 @@ extern Relation heap_create(const char *relname,
 							bool allow_system_table_mods,
 							TransactionId *relfrozenxid,
 							MultiXactId *relminmxid,
-							bool create_storage);
+							bool create_storage,
+							RelFileLocator *old_rlocator);
 
 extern Oid	heap_create_with_catalog(const char *relname,
 									 Oid relnamespace,

src/include/catalog/index.h

@@ -16,6 +16,7 @@
 
 #include "catalog/objectaddress.h"
 #include "nodes/execnodes.h"
+#include "storage/relfilelocator.h"
 
 
 #define DEFAULT_INDEX_TYPE	"btree"
@@ -86,7 +87,8 @@ extern Oid	index_create(Relation heapRelation,
 						 bits16 constr_flags,
 						 bool allow_system_table_mods,
 						 bool is_internal,
-						 Oid *constraintId);
+						 Oid *constraintId,
+						 RelFileLocator *old_rlocator);
 
 #define	INDEX_CONSTR_CREATE_MARK_AS_PRIMARY	(1 << 0)
 #define	INDEX_CONSTR_CREATE_DEFERRABLE		(1 << 1)