Move relations key

percona · Oct 8, 2024 · 7eb8f74 · 7eb8f74
1 parent 6ab94e6
commit 7eb8f74
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 5 deletions.
diff --git a/sql/tablespace.sql b/sql/tablespace.sql
@@ -1,12 +1,9 @@
 CREATE EXTENSION pg_tde;
 
-SELECT  * FROM pg_tde_principal_key_info();
-
 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');
 
-
-CREATE TABLE test(num1 bigint, num2 double precision, t text);
+CREATE TABLE test(num1 bigint, num2 double precision, t text) using tde_heap_basic;;
 INSERT INTO test(num1, num2, t)
   SELECT round(random()*100), random(), 'text'
   FROM generate_series(1, 10) s(i);
@@ -16,10 +13,13 @@ SET allow_in_place_tablespaces = true;
 CREATE TABLESPACE test_tblspace LOCATION '';
 
 ALTER TABLE test SET TABLESPACE test_tblspace;
+SELECT count(*) FROM test;
 ALTER TABLE test SET TABLESPACE pg_default;
 
 REINDEX (TABLESPACE test_tblspace, CONCURRENTLY) TABLE test;
-INSERT INTO test VALUES (10, 2);
+INSERT INTO test VALUES (110, 2);
+
+SELECT * FROM test WHERE num1=110;
 
 DROP TABLE test;
 DROP TABLESPACE test_tblspace;

diff --git a/src/access/pg_tde_tdemap.c b/src/access/pg_tde_tdemap.c
@@ -832,6 +832,38 @@ pg_tde_write_map_keydata_files(off_t map_size, char *m_file_data, off_t keydata_
 	return !is_err;
 }
 
+/* 
+ * Move relation's key to the new physical location and cache it with the new
+ * relfilenode.
+ * Due to ALTER TABLE SET TABLESPACE for example.
+ */
+bool
+pg_tde_move_rel_key(const RelFileLocator *newrlocator, const RelFileLocator *oldrlocator)
+{
+	RelKeyData 	*rel_key;
+	RelKeyData 	*enc_key;
+	TDEPrincipalKey *principal_key;
+	XLogRelKey xlrec;
+	LWLock *lock_pk = tde_lwlock_enc_keys();
+
+	LWLockAcquire(lock_pk, LW_EXCLUSIVE);
+	principal_key = GetPrincipalKey(oldrlocator->dbOid, oldrlocator->spcOid, LW_EXCLUSIVE);
+	rel_key = GetRelationKey(*oldrlocator);
+	Assert(rel_key);
+	enc_key = tde_encrypt_rel_key(principal_key, rel_key, newrlocator);
+	pg_tde_write_key_map_entry(newrlocator, enc_key, &principal_key->keyInfo);
+	pg_tde_put_key_into_cache(newrlocator->relNumber, rel_key);
+
+	xlrec.rlocator = *newrlocator;
+	xlrec.relKey = *enc_key;
+	XLogBeginInsert();
+	XLogRegisterData((char *) &xlrec, sizeof(xlrec));
+	XLogInsert(RM_TDERMGR_ID, XLOG_TDE_ADD_RELATION_KEY);
+
+	LWLockRelease(lock_pk);
+	pfree(enc_key);
+}
+
 #endif		/* !FRONTEND */
 
 /*

diff --git a/src/include/access/pg_tde_tdemap.h b/src/include/access/pg_tde_tdemap.h
@@ -50,6 +50,7 @@ extern RelKeyData* tde_create_rel_key(Oid rel_id, InternalKey *key, TDEPrincipal
 extern RelKeyData *tde_encrypt_rel_key(TDEPrincipalKey *principal_key, RelKeyData *rel_key_data, const RelFileLocator *rlocator);
 extern RelKeyData *tde_decrypt_rel_key(TDEPrincipalKey *principal_key, RelKeyData *enc_rel_key_data, const RelFileLocator *rlocator);
 extern RelKeyData *pg_tde_get_key_from_file(const RelFileLocator *rlocator);
+extern bool pg_tde_move_rel_key(const RelFileLocator *newrlocator, const RelFileLocator *oldrlocator);
 
 extern void pg_tde_set_db_file_paths(const RelFileLocator *rlocator, char *map_path, char *keydata_path);
 

diff --git a/src16/access/pg_tdeam_handler.c b/src16/access/pg_tdeam_handler.c
@@ -707,6 +707,7 @@ pg_tdeam_relation_copy_data(Relation rel, const RelFileLocator *newrlocator)
 		}
 	}
 
+	pg_tde_move_rel_key(newrlocator, rel->rd_locator);
 
 	/* drop old relation, and close new one */
 	RelationDropStorage(rel);

diff --git a/src17/access/pg_tdeam_handler.c b/src17/access/pg_tdeam_handler.c
@@ -702,6 +702,7 @@ pg_tdeam_relation_copy_data(Relation rel, const RelFileLocator *newrlocator)
 		}
 	}
 
+	pg_tde_move_rel_key(newrlocator, rel->rd_locator);
 
 	/* drop old relation, and close new one */
 	RelationDropStorage(rel);