diff --git a/contrib/pg_tde/src/access/pg_tde_xlog_smgr.c b/contrib/pg_tde/src/access/pg_tde_xlog_smgr.c
index 78c830e0022..6e075dadf14 100644
--- a/contrib/pg_tde/src/access/pg_tde_xlog_smgr.c
+++ b/contrib/pg_tde/src/access/pg_tde_xlog_smgr.c
@@ -266,16 +266,33 @@ TDEXLogSmgrInitWrite(bool encrypt_xlog)
 		pfree(key);
 }
 
+/*
+ * Used by pg_tde_restore_encrypt to simulate being constantly in recovery
+ * since the command does not have access to any information about if we are in
+ * recovery or not.
+ *
+ * Creates a dummy key which points at the very end of the WAL stream.
+ */
 void
-TDEXLogSmgrInitWriteReuseKey()
+TDEXLogSmgrInitWriteOldKeys()
 {
-	WalEncryptionKey *key = pg_tde_read_last_wal_key();
+	WALKeyCacheRec *keys;
+	WalEncryptionKey dummy = {
+		.type = WAL_KEY_TYPE_UNENCRYPTED,
+		.wal_start = {.tli = -1,.lsn = -1}
+	};
 
-	if (key)
+	EncryptionKey = dummy;
+	TDEXLogSetEncKeyLocation(dummy.wal_start);
+
+	keys = pg_tde_get_wal_cache_keys();
+
+	if (keys == NULL)
 	{
-		EncryptionKey = *key;
-		TDEXLogSetEncKeyLocation(EncryptionKey.wal_start);
-		pfree(key);
+		WalLocation start = {.tli = 1,.lsn = 0};
+
+		/* cache is empty, prefetch keys from disk */
+		pg_tde_fetch_wal_keys(start);
 	}
 }
 
diff --git a/contrib/pg_tde/src/bin/pg_tde_restore_encrypt.c b/contrib/pg_tde/src/bin/pg_tde_restore_encrypt.c
index a02519c787b..03f73cc1e6a 100644
--- a/contrib/pg_tde/src/bin/pg_tde_restore_encrypt.c
+++ b/contrib/pg_tde/src/bin/pg_tde_restore_encrypt.c
@@ -68,7 +68,7 @@ write_encrypted_segment(const char *segpath, const char *segname, const char *tm
 
 	XLogFromFileName(segname, &tli, &segno, walsegsz);
 
-	TDEXLogSmgrInitWriteReuseKey();
+	TDEXLogSmgrInitWriteOldKeys();
 
 	w = xlog_smgr->seg_write(segfd, buf.data, r, pos, tli, segno, walsegsz);
 
diff --git a/contrib/pg_tde/src/include/access/pg_tde_xlog_smgr.h b/contrib/pg_tde/src/include/access/pg_tde_xlog_smgr.h
index 272880299f0..16a0ba3cd7b 100644
--- a/contrib/pg_tde/src/include/access/pg_tde_xlog_smgr.h
+++ b/contrib/pg_tde/src/include/access/pg_tde_xlog_smgr.h
@@ -11,7 +11,7 @@ extern Size TDEXLogEncryptStateSize(void);
 extern void TDEXLogShmemInit(void);
 extern void TDEXLogSmgrInit(void);
 extern void TDEXLogSmgrInitWrite(bool encrypt_xlog);
-extern void TDEXLogSmgrInitWriteReuseKey(void);
+extern void TDEXLogSmgrInitWriteOldKeys(void);
 
 extern void TDEXLogCryptBuffer(const void *buf, void *out_buf, size_t count, off_t offset,
 							   TimeLineID tli, XLogSegNo segno, int segSize);
diff --git a/contrib/pg_tde/t/wal_archiving.pl b/contrib/pg_tde/t/wal_archiving.pl
index e3ac1ce4239..f35f5d6aa07 100644
--- a/contrib/pg_tde/t/wal_archiving.pl
+++ b/contrib/pg_tde/t/wal_archiving.pl
@@ -58,7 +58,7 @@ like(
 unlike(
 	`strings $data_dir/pg_wal/0000000100000000000000??`,
 	qr/foobar_enc/,
-	'should not find foobar_enc in WAL');
+	'should not find foobar_enc in WAL since it is encrypted');
 
 $primary->stop;
 
@@ -84,10 +84,10 @@ $replica->start;
 
 $data_dir = $replica->data_dir;
 
-unlike(
+like(
 	`strings $data_dir/pg_wal/0000000100000000000000??`,
 	qr/foobar_plain/,
-	'should not find foobar_plain in WAL since it is encrypted');
+	'should find foobar_plain in WAL since we use the same key file');
 unlike(
 	`strings $data_dir/pg_wal/0000000100000000000000??`,
 	qr/foobar_enc/,