diff --git a/contrib/pg_tde/src/access/pg_tde_tdemap.c b/contrib/pg_tde/src/access/pg_tde_tdemap.c
index 4fd8a7f5cb9..0e6aecdb709 100644
--- a/contrib/pg_tde/src/access/pg_tde_tdemap.c
+++ b/contrib/pg_tde/src/access/pg_tde_tdemap.c
@@ -20,7 +20,6 @@
 #include "access/xlog_internal.h"
 #include "access/xloginsert.h"
 #include "utils/builtins.h"
-#include "utils/hsearch.h"
 #include "miscadmin.h"
 
 #include "access/pg_tde_tdemap.h"
@@ -28,6 +27,7 @@
 #include "catalog/tde_global_space.h"
 #include "catalog/tde_principal_key.h"
 #include "encryption/enc_aes.h"
+#include "encryption/enc_tde.h"
 #include "keyring/keyring_api.h"
 
 #include <openssl/rand.h>
@@ -66,24 +66,6 @@ typedef struct TDEFileHeader
 	TDESignedPrincipalKeyInfo signed_key_info;
 } TDEFileHeader;
 
-typedef struct
-{
-	RelFileLocator rel;
-	InternalKey key;
-} TempRelKeyEntry;
-
-#ifndef FRONTEND
-
-/* Arbitrarily picked small number of temporary relations */
-#define INIT_TEMP_RELS 16
-
-/*
- * Each backend has a hashtable that stores the keys for all temporary tables.
- */
-static HTAB *TempRelKeys = NULL;
-
-#endif
-
 static WALKeyCacheRec *tde_wal_key_cache = NULL;
 static WALKeyCacheRec *tde_wal_key_last_rec = NULL;
 
@@ -98,72 +80,25 @@ static int	pg_tde_open_file_read(const char *tde_filename, bool ignore_missing,
 static WALKeyCacheRec *pg_tde_add_wal_key_to_cache(InternalKey *cached_key, XLogRecPtr start_lsn);
 
 #ifndef FRONTEND
-static InternalKey *pg_tde_create_smgr_key_temp(const RelFileLocator *newrlocator);
-static InternalKey *pg_tde_create_smgr_key_perm(const RelFileLocator *newrlocator);
-static void pg_tde_generate_internal_key(InternalKey *int_key, TDEMapEntryType entry_type);
 static int	pg_tde_file_header_write(const char *tde_filename, int fd, const TDESignedPrincipalKeyInfo *signed_key_info, off_t *bytes_written);
 static void pg_tde_sign_principal_key_info(TDESignedPrincipalKeyInfo *signed_key_info, const TDEPrincipalKey *principal_key);
 static void pg_tde_write_one_map_entry(int fd, const TDEMapEntry *map_entry, off_t *offset, const char *db_map_path);
-static void pg_tde_write_key_map_entry(const RelFileLocator *rlocator, InternalKey *rel_key_data, TDEPrincipalKey *principal_key);
-static void pg_tde_free_key_map_entry(const RelFileLocator *rlocator);
+static void pg_tde_write_key_map_entry(const RelFileLocator *rlocator, const InternalKey *rel_key_data, TDEPrincipalKey *principal_key);
 static int	keyrotation_init_file(const TDESignedPrincipalKeyInfo *signed_key_info, char *rotated_filename, const char *filename, off_t *curr_pos);
 static void finalize_key_rotation(const char *path_old, const char *path_new);
 static int	pg_tde_open_file_write(const char *tde_filename, const TDESignedPrincipalKeyInfo *signed_key_info, bool truncate, off_t *curr_pos);
 
-InternalKey *
-pg_tde_create_smgr_key(const RelFileLocatorBackend *newrlocator)
-{
-	if (RelFileLocatorBackendIsTemp(*newrlocator))
-		return pg_tde_create_smgr_key_temp(&newrlocator->locator);
-	else
-		return pg_tde_create_smgr_key_perm(&newrlocator->locator);
-}
-
-static InternalKey *
-pg_tde_create_smgr_key_temp(const RelFileLocator *newrlocator)
-{
-	InternalKey *rel_key_data = palloc_object(InternalKey);
-	TempRelKeyEntry *entry;
-	bool		found;
-
-	pg_tde_generate_internal_key(rel_key_data, TDE_KEY_TYPE_SMGR);
-
-	if (TempRelKeys == NULL)
-	{
-		HASHCTL		ctl;
-
-		ctl.keysize = sizeof(RelFileLocator);
-		ctl.entrysize = sizeof(TempRelKeyEntry);
-		TempRelKeys = hash_create("pg_tde temporary relation keys",
-								  INIT_TEMP_RELS,
-								  &ctl,
-								  HASH_ELEM | HASH_BLOBS);
-	}
-
-	entry = (TempRelKeyEntry *) hash_search(TempRelKeys,
-											newrlocator,
-											HASH_ENTER, &found);
-	Assert(!found);
-
-	entry->key = *rel_key_data;
-
-	return rel_key_data;
-}
-
-static InternalKey *
-pg_tde_create_smgr_key_perm(const RelFileLocator *newrlocator)
+void
+pg_tde_save_smgr_key(RelFileLocator rel, const InternalKey *rel_key_data, bool write_xlog)
 {
-	InternalKey *rel_key_data = palloc_object(InternalKey);
 	TDEPrincipalKey *principal_key;
 	LWLock	   *lock_pk = tde_lwlock_enc_keys();
 	XLogRelKey	xlrec = {
-		.rlocator = *newrlocator,
+		.rlocator = rel,
 	};
 
-	pg_tde_generate_internal_key(rel_key_data, TDE_KEY_TYPE_SMGR);
-
 	LWLockAcquire(lock_pk, LW_EXCLUSIVE);
-	principal_key = GetPrincipalKey(newrlocator->dbOid, LW_EXCLUSIVE);
+	principal_key = GetPrincipalKey(rel.dbOid, LW_EXCLUSIVE);
 	if (principal_key == NULL)
 	{
 		ereport(ERROR,
@@ -171,65 +106,19 @@ pg_tde_create_smgr_key_perm(const RelFileLocator *newrlocator)
 				errhint("create one using pg_tde_set_key before using encrypted tables"));
 	}
 
-	pg_tde_write_key_map_entry(newrlocator, rel_key_data, principal_key);
+	pg_tde_write_key_map_entry(&rel, rel_key_data, principal_key);
 	LWLockRelease(lock_pk);
 
-	/*
-	 * It is fine to write the to WAL after writing to the file since we have
-	 * not WAL logged the SMGR CREATE event either.
-	 */
-	XLogBeginInsert();
-	XLogRegisterData((char *) &xlrec, sizeof(xlrec));
-	XLogInsert(RM_TDERMGR_ID, XLOG_TDE_ADD_RELATION_KEY);
-
-	return rel_key_data;
-}
-
-void
-pg_tde_create_smgr_key_perm_redo(const RelFileLocator *newrlocator)
-{
-	InternalKey rel_key_data;
-	InternalKey *old_key;
-	TDEPrincipalKey *principal_key;
-	LWLock	   *lock_pk = tde_lwlock_enc_keys();
-
-	if ((old_key = pg_tde_get_key_from_file(newrlocator, TDE_KEY_TYPE_SMGR)))
-	{
-		pfree(old_key);
-		return;
-	}
-
-	pg_tde_generate_internal_key(&rel_key_data, TDE_KEY_TYPE_SMGR);
-
-	LWLockAcquire(lock_pk, LW_EXCLUSIVE);
-	principal_key = GetPrincipalKey(newrlocator->dbOid, LW_EXCLUSIVE);
-	if (principal_key == NULL)
+	if (write_xlog)
 	{
-		ereport(ERROR,
-				errmsg("principal key not configured"),
-				errhint("create one using pg_tde_set_key before using encrypted tables"));
+		/*
+		 * It is fine to write the to WAL after writing to the file since we
+		 * have not WAL logged the SMGR CREATE event either.
+		 */
+		XLogBeginInsert();
+		XLogRegisterData((char *) &xlrec, sizeof(xlrec));
+		XLogInsert(RM_TDERMGR_ID, XLOG_TDE_ADD_RELATION_KEY);
 	}
-
-	pg_tde_write_key_map_entry(newrlocator, &rel_key_data, principal_key);
-	LWLockRelease(lock_pk);
-}
-
-static void
-pg_tde_generate_internal_key(InternalKey *int_key, TDEMapEntryType entry_type)
-{
-	int_key->type = entry_type;
-	int_key->start_lsn = InvalidXLogRecPtr;
-
-	if (!RAND_bytes(int_key->key, INTERNAL_KEY_LEN))
-		ereport(ERROR,
-				errcode(ERRCODE_INTERNAL_ERROR),
-				errmsg("could not generate internal key: %s",
-					   ERR_error_string(ERR_get_error(), NULL)));
-	if (!RAND_bytes(int_key->base_iv, INTERNAL_KEY_IV_LEN))
-		ereport(ERROR,
-				errcode(ERRCODE_INTERNAL_ERROR),
-				errmsg("could not generate IV: %s",
-					   ERR_error_string(ERR_get_error(), NULL)));
 }
 
 const char *
@@ -275,18 +164,6 @@ pg_tde_create_wal_key(InternalKey *rel_key_data, const RelFileLocator *newrlocat
 	LWLockRelease(tde_lwlock_enc_keys());
 }
 
-void
-DeleteSMGRRelationKey(RelFileLocatorBackend rel)
-{
-	if (RelFileLocatorBackendIsTemp(rel))
-	{
-		Assert(TempRelKeys);
-		hash_search(TempRelKeys, &rel.locator, HASH_REMOVE, NULL);
-	}
-	else
-		pg_tde_free_key_map_entry(&rel.locator);
-}
-
 /*
  * Deletes the key map file for a given database.
  */
@@ -463,7 +340,7 @@ pg_tde_write_one_map_entry(int fd, const TDEMapEntry *map_entry, off_t *offset,
  * concurrent in place updates leading to data conflicts.
  */
 void
-pg_tde_write_key_map_entry(const RelFileLocator *rlocator, InternalKey *rel_key_data, TDEPrincipalKey *principal_key)
+pg_tde_write_key_map_entry(const RelFileLocator *rlocator, const InternalKey *rel_key_data, TDEPrincipalKey *principal_key)
 {
 	char		db_map_path[MAXPGPATH];
 	int			map_fd;
@@ -518,16 +395,14 @@ pg_tde_write_key_map_entry(const RelFileLocator *rlocator, InternalKey *rel_key_
  * This fucntion is called by the pg_tde SMGR when storage is unlinked on
  * transaction commit/abort.
  */
-static void
-pg_tde_free_key_map_entry(const RelFileLocator *rlocator)
+void
+pg_tde_free_key_map_entry(const RelFileLocator rlocator)
 {
 	char		db_map_path[MAXPGPATH];
 	File		map_fd;
 	off_t		curr_pos = 0;
 
-	Assert(rlocator);
-
-	pg_tde_set_db_file_path(rlocator->dbOid, db_map_path);
+	pg_tde_set_db_file_path(rlocator.dbOid, db_map_path);
 
 	LWLockAcquire(tde_lwlock_enc_keys(), LW_EXCLUSIVE);
 
@@ -542,7 +417,7 @@ pg_tde_free_key_map_entry(const RelFileLocator *rlocator)
 		if (!pg_tde_read_one_map_entry(map_fd, &map_entry, &curr_pos))
 			break;
 
-		if (map_entry.type != MAP_ENTRY_EMPTY && map_entry.spcOid == rlocator->spcOid && map_entry.relNumber == rlocator->relNumber)
+		if (map_entry.type != MAP_ENTRY_EMPTY && map_entry.spcOid == rlocator.spcOid && map_entry.relNumber == rlocator.relNumber)
 		{
 			TDEMapEntry empty_map_entry = {
 				.type = MAP_ENTRY_EMPTY,
@@ -1084,57 +959,27 @@ pg_tde_get_principal_key_info(Oid dbOid)
 	return signed_key_info;
 }
 
-static InternalKey *
-pg_tde_get_temporary_rel_key(const RelFileLocator *rel)
-{
-#ifndef FRONTEND
-	TempRelKeyEntry *entry;
-
-	if (TempRelKeys == NULL)
-		return NULL;
-
-	entry = hash_search(TempRelKeys, rel, HASH_FIND, NULL);
-
-	if (entry)
-	{
-		InternalKey *key = palloc_object(InternalKey);
-
-		*key = entry->key;
-		return key;
-	}
-#endif
-
-	return NULL;
-}
-
 /*
  * Figures out whether a relation is encrypted or not, but without trying to
  * decrypt the key if it is.
  */
 bool
-IsSMGRRelationEncrypted(RelFileLocatorBackend rel)
+pg_tde_has_smgr_key(RelFileLocator rel)
 {
 	bool		result;
 	TDEMapEntry map_entry;
 	char		db_map_path[MAXPGPATH];
 
-	Assert(rel.locator.relNumber != InvalidRelFileNumber);
-
-	if (RelFileLocatorBackendIsTemp(rel))
-#ifndef FRONTEND
-		return TempRelKeys && hash_search(TempRelKeys, &rel.locator, HASH_FIND, NULL);
-#else
-		return false;
-#endif
+	Assert(rel.relNumber != InvalidRelFileNumber);
 
-	pg_tde_set_db_file_path(rel.locator.dbOid, db_map_path);
+	pg_tde_set_db_file_path(rel.dbOid, db_map_path);
 
 	if (access(db_map_path, F_OK) == -1)
 		return false;
 
 	LWLockAcquire(tde_lwlock_enc_keys(), LW_SHARED);
 
-	result = pg_tde_find_map_entry(&rel.locator, TDE_KEY_TYPE_SMGR, db_map_path, &map_entry);
+	result = pg_tde_find_map_entry(&rel, TDE_KEY_TYPE_SMGR, db_map_path, &map_entry);
 
 	LWLockRelease(tde_lwlock_enc_keys());
 	return result;
@@ -1144,14 +989,11 @@ IsSMGRRelationEncrypted(RelFileLocatorBackend rel)
  * Returns TDE key for a given relation.
  */
 InternalKey *
-GetSMGRRelationKey(RelFileLocatorBackend rel)
+pg_tde_get_smgr_key(RelFileLocator rel)
 {
-	Assert(rel.locator.relNumber != InvalidRelFileNumber);
+	Assert(rel.relNumber != InvalidRelFileNumber);
 
-	if (RelFileLocatorBackendIsTemp(rel))
-		return pg_tde_get_temporary_rel_key(&rel.locator);
-	else
-		return pg_tde_get_key_from_file(&rel.locator, TDE_KEY_TYPE_SMGR);
+	return pg_tde_get_key_from_file(&rel, TDE_KEY_TYPE_SMGR);
 }
 
 /*
diff --git a/contrib/pg_tde/src/access/pg_tde_xlog.c b/contrib/pg_tde/src/access/pg_tde_xlog.c
index 9d817036bd8..0d0581d0604 100644
--- a/contrib/pg_tde/src/access/pg_tde_xlog.c
+++ b/contrib/pg_tde/src/access/pg_tde_xlog.c
@@ -25,6 +25,7 @@
 
 #include "access/pg_tde_xlog.h"
 #include "encryption/enc_tde.h"
+#include "smgr/pg_tde_smgr.h"
 
 static void tdeheap_rmgr_redo(XLogReaderState *record);
 static void tdeheap_rmgr_desc(StringInfo buf, XLogReaderState *record);
@@ -52,7 +53,7 @@ tdeheap_rmgr_redo(XLogReaderState *record)
 	{
 		XLogRelKey *xlrec = (XLogRelKey *) XLogRecGetData(record);
 
-		pg_tde_create_smgr_key_perm_redo(&xlrec->rlocator);
+		tde_smgr_create_key_redo(&xlrec->rlocator);
 	}
 	else if (info == XLOG_TDE_ADD_PRINCIPAL_KEY)
 	{
diff --git a/contrib/pg_tde/src/common/pg_tde_utils.c b/contrib/pg_tde/src/common/pg_tde_utils.c
index 109657f2b9d..e092557f90c 100644
--- a/contrib/pg_tde/src/common/pg_tde_utils.c
+++ b/contrib/pg_tde/src/common/pg_tde_utils.c
@@ -16,8 +16,7 @@
 
 #ifndef FRONTEND
 #include "fmgr.h"
-#include "catalog/pg_class.h"
-#include "access/pg_tde_tdemap.h"
+#include "smgr/pg_tde_smgr.h"
 #include "access/relation.h"
 #include "utils/rel.h"
 
@@ -28,7 +27,6 @@ pg_tde_is_encrypted(PG_FUNCTION_ARGS)
 	Oid			relationOid = PG_GETARG_OID(0);
 	LOCKMODE	lockmode = AccessShareLock;
 	Relation	rel = relation_open(relationOid, lockmode);
-	RelFileLocatorBackend rlocator = {.locator = rel->rd_locator,.backend = rel->rd_backend};
 	bool		result;
 
 	if (!RELKIND_HAS_STORAGE(rel->rd_rel->relkind))
@@ -42,7 +40,7 @@ pg_tde_is_encrypted(PG_FUNCTION_ARGS)
 				errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 				errmsg("we cannot check if temporary relations from other backends are encrypted"));
 
-	result = IsSMGRRelationEncrypted(rlocator);
+	result = tde_smgr_rel_is_encrypted(RelationGetSmgr(rel));
 
 	relation_close(rel, lockmode);
 
diff --git a/contrib/pg_tde/src/encryption/enc_tde.c b/contrib/pg_tde/src/encryption/enc_tde.c
index 59e08dc1e9c..2676efb92aa 100644
--- a/contrib/pg_tde/src/encryption/enc_tde.c
+++ b/contrib/pg_tde/src/encryption/enc_tde.c
@@ -7,6 +7,13 @@
 #include "encryption/enc_aes.h"
 #include "storage/bufmgr.h"
 
+#ifdef FRONTEND
+#include "pg_tde_fe.h"
+#endif
+
+#include <openssl/rand.h>
+#include <openssl/err.h>
+
 #define AES_BLOCK_SIZE 		        16
 #define NUM_AES_BLOCKS_IN_BATCH     200
 #define DATA_BYTES_PER_AES_BATCH    (NUM_AES_BLOCKS_IN_BATCH * AES_BLOCK_SIZE)
@@ -23,6 +30,24 @@ iv_prefix_debug(const char *iv_prefix, char *out_hex)
 }
 #endif
 
+void
+pg_tde_generate_internal_key(InternalKey *int_key, TDEMapEntryType entry_type)
+{
+	int_key->type = entry_type;
+	int_key->start_lsn = InvalidXLogRecPtr;
+
+	if (!RAND_bytes(int_key->key, INTERNAL_KEY_LEN))
+		ereport(ERROR,
+				errcode(ERRCODE_INTERNAL_ERROR),
+				errmsg("could not generate internal key: %s",
+					   ERR_error_string(ERR_get_error(), NULL)));
+	if (!RAND_bytes(int_key->base_iv, INTERNAL_KEY_IV_LEN))
+		ereport(ERROR,
+				errcode(ERRCODE_INTERNAL_ERROR),
+				errmsg("could not generate IV: %s",
+					   ERR_error_string(ERR_get_error(), NULL)));
+}
+
 /*
  * Encrypts/decrypts `data` with a given `key`. The result is written to `out`.
  *
diff --git a/contrib/pg_tde/src/include/access/pg_tde_tdemap.h b/contrib/pg_tde/src/include/access/pg_tde_tdemap.h
index b65b063be17..e19f787af2b 100644
--- a/contrib/pg_tde/src/include/access/pg_tde_tdemap.h
+++ b/contrib/pg_tde/src/include/access/pg_tde_tdemap.h
@@ -83,8 +83,6 @@ extern WALKeyCacheRec *pg_tde_fetch_wal_keys(XLogRecPtr start_lsn);
 extern WALKeyCacheRec *pg_tde_get_wal_cache_keys(void);
 extern void pg_tde_wal_last_key_set_lsn(XLogRecPtr lsn, const char *keyfile_path);
 
-extern InternalKey *pg_tde_create_smgr_key(const RelFileLocatorBackend *newrlocator);
-extern void pg_tde_create_smgr_key_perm_redo(const RelFileLocator *newrlocator);
 extern void pg_tde_create_wal_key(InternalKey *rel_key_data, const RelFileLocator *newrlocator, TDEMapEntryType flags);
 
 #define PG_TDE_MAP_FILENAME			"%d_keys"
@@ -95,9 +93,10 @@ pg_tde_set_db_file_path(Oid dbOid, char *path)
 	join_path_components(path, pg_tde_get_data_dir(), psprintf(PG_TDE_MAP_FILENAME, dbOid));
 }
 
-extern bool IsSMGRRelationEncrypted(RelFileLocatorBackend rel);
-extern InternalKey *GetSMGRRelationKey(RelFileLocatorBackend rel);
-extern void DeleteSMGRRelationKey(RelFileLocatorBackend rel);
+extern void pg_tde_save_smgr_key(RelFileLocator rel, const InternalKey *key, bool write_xlog);
+extern bool pg_tde_has_smgr_key(RelFileLocator rel);
+extern InternalKey *pg_tde_get_smgr_key(RelFileLocator rel);
+extern void pg_tde_free_key_map_entry(RelFileLocator rel);
 
 extern int	pg_tde_count_relations(Oid dbOid);
 
diff --git a/contrib/pg_tde/src/include/encryption/enc_tde.h b/contrib/pg_tde/src/include/encryption/enc_tde.h
index ac417c5254f..879733cc0a1 100644
--- a/contrib/pg_tde/src/include/encryption/enc_tde.h
+++ b/contrib/pg_tde/src/include/encryption/enc_tde.h
@@ -12,6 +12,7 @@
 
 #include "access/pg_tde_tdemap.h"
 
+extern void pg_tde_generate_internal_key(InternalKey *int_key, TDEMapEntryType entry_type);
 extern void pg_tde_stream_crypt(const char *iv_prefix, uint32 start_offset, const char *data, uint32 data_len, char *out, InternalKey *key, void **ctxPtr);
 
 #endif							/* ENC_TDE_H */
diff --git a/contrib/pg_tde/src/include/smgr/pg_tde_smgr.h b/contrib/pg_tde/src/include/smgr/pg_tde_smgr.h
index 65b54e5a5e9..e0f09efc4c3 100644
--- a/contrib/pg_tde/src/include/smgr/pg_tde_smgr.h
+++ b/contrib/pg_tde/src/include/smgr/pg_tde_smgr.h
@@ -9,6 +9,11 @@
 #ifndef PG_TDE_SMGR_H
 #define PG_TDE_SMGR_H
 
+#include "storage/relfilelocator.h"
+#include "storage/smgr.h"
+
 extern void RegisterStorageMgr(void);
+extern void tde_smgr_create_key_redo(const RelFileLocator *rlocator);
+extern bool tde_smgr_rel_is_encrypted(SMgrRelation reln);
 
 #endif							/* PG_TDE_SMGR_H */
diff --git a/contrib/pg_tde/src/smgr/pg_tde_smgr.c b/contrib/pg_tde/src/smgr/pg_tde_smgr.c
index 9a6537b728c..cf3de6cd5bc 100644
--- a/contrib/pg_tde/src/smgr/pg_tde_smgr.c
+++ b/contrib/pg_tde/src/smgr/pg_tde_smgr.c
@@ -1,10 +1,13 @@
-#include "smgr/pg_tde_smgr.h"
 #include "postgres.h"
+
+#include "smgr/pg_tde_smgr.h"
 #include "storage/smgr.h"
 #include "storage/md.h"
 #include "catalog/catalog.h"
 #include "encryption/enc_aes.h"
+#include "encryption/enc_tde.h"
 #include "access/pg_tde_tdemap.h"
+#include "utils/hsearch.h"
 #include "pg_tde_event_capture.h"
 
 typedef enum TDEMgrRelationEncryptionStatus
@@ -43,8 +46,82 @@ typedef struct TDESMgrRelation
 	InternalKey relKey;
 } TDESMgrRelation;
 
+typedef struct
+{
+	RelFileLocator rel;
+	InternalKey key;
+} TempRelKeyEntry;
+
+#define INIT_TEMP_RELS 16
+
+/*
+ * Each backend has a hashtable that stores the keys for all temproary tables.
+ */
+static HTAB *TempRelKeys = NULL;
+
+static SMgrId OurSMgrId = MaxSMgrId;
+
+static void tde_smgr_save_temp_key(const RelFileLocator *newrlocator, const InternalKey *key);
+static InternalKey *tde_smgr_get_temp_key(const RelFileLocator *rel);
+static bool tde_smgr_has_temp_key(const RelFileLocator *rel);
+static void tde_smgr_remove_temp_key(const RelFileLocator *rel);
 static void CalcBlockIv(ForkNumber forknum, BlockNumber bn, const unsigned char *base_iv, unsigned char *iv);
 
+static InternalKey *
+tde_smgr_create_key(const RelFileLocatorBackend *smgr_rlocator)
+{
+	InternalKey *key = palloc_object(InternalKey);
+
+	pg_tde_generate_internal_key(key, TDE_KEY_TYPE_SMGR);
+
+	if (RelFileLocatorBackendIsTemp(*smgr_rlocator))
+		tde_smgr_save_temp_key(&smgr_rlocator->locator, key);
+	else
+		pg_tde_save_smgr_key(smgr_rlocator->locator, key, true);
+
+	return key;
+}
+
+void
+tde_smgr_create_key_redo(const RelFileLocator *rlocator)
+{
+	InternalKey key;
+
+	if (pg_tde_has_smgr_key(*rlocator))
+		return;
+
+	pg_tde_generate_internal_key(&key, TDE_KEY_TYPE_SMGR);
+
+	pg_tde_save_smgr_key(*rlocator, &key, false);
+}
+
+static bool
+tde_smgr_is_encrypted(const RelFileLocatorBackend *smgr_rlocator)
+{
+	if (RelFileLocatorBackendIsTemp(*smgr_rlocator))
+		return tde_smgr_has_temp_key(&smgr_rlocator->locator);
+	else
+		return pg_tde_has_smgr_key(smgr_rlocator->locator);
+}
+
+static InternalKey *
+tde_smgr_get_key(const RelFileLocatorBackend *smgr_rlocator)
+{
+	if (RelFileLocatorBackendIsTemp(*smgr_rlocator))
+		return tde_smgr_get_temp_key(&smgr_rlocator->locator);
+	else
+		return pg_tde_get_smgr_key(smgr_rlocator->locator);
+}
+
+static void
+tde_smgr_remove_key(const RelFileLocatorBackend *smgr_rlocator)
+{
+	if (RelFileLocatorBackendIsTemp(*smgr_rlocator))
+		tde_smgr_remove_temp_key(&smgr_rlocator->locator);
+	else
+		pg_tde_free_key_map_entry(smgr_rlocator->locator);
+}
+
 static bool
 tde_smgr_should_encrypt(const RelFileLocatorBackend *smgr_rlocator, RelFileLocator *old_locator)
 {
@@ -66,13 +143,25 @@ tde_smgr_should_encrypt(const RelFileLocatorBackend *smgr_rlocator, RelFileLocat
 					.backend = smgr_rlocator->backend,
 				};
 
-				return IsSMGRRelationEncrypted(old_smgr_locator);
+				return tde_smgr_is_encrypted(&old_smgr_locator);
 			}
 	}
 
 	return false;
 }
 
+bool
+tde_smgr_rel_is_encrypted(SMgrRelation reln)
+{
+	TDESMgrRelation *tdereln = (TDESMgrRelation *) reln;
+
+	if (reln->smgr_which != OurSMgrId)
+		return false;
+
+	return tdereln->encryption_status == RELATION_KEY_AVAILABLE ||
+		tdereln->encryption_status == RELATION_KEY_NOT_AVAILABLE;
+}
+
 static void
 tde_mdwritev(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 			 const void **buffers, BlockNumber nblocks, bool skipFsync)
@@ -90,7 +179,7 @@ tde_mdwritev(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 
 		if (tdereln->encryption_status == RELATION_KEY_NOT_AVAILABLE)
 		{
-			InternalKey *int_key = GetSMGRRelationKey(reln->smgr_rlocator);
+			InternalKey *int_key = tde_smgr_get_key(&reln->smgr_rlocator);
 
 			tdereln->relKey = *int_key;
 			tdereln->encryption_status = RELATION_KEY_AVAILABLE;
@@ -139,8 +228,8 @@ tde_mdunlink(RelFileLocatorBackend rlocator, ForkNumber forknum, bool isRedo)
 	 */
 	if (forknum == MAIN_FORKNUM || forknum == InvalidForkNumber)
 	{
-		if (IsSMGRRelationEncrypted(rlocator))
-			DeleteSMGRRelationKey(rlocator);
+		if (tde_smgr_is_encrypted(&rlocator))
+			tde_smgr_remove_key(&rlocator);
 	}
 }
 
@@ -161,7 +250,7 @@ tde_mdextend(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 
 		if (tdereln->encryption_status == RELATION_KEY_NOT_AVAILABLE)
 		{
-			InternalKey *int_key = GetSMGRRelationKey(reln->smgr_rlocator);
+			InternalKey *int_key = tde_smgr_get_key(&reln->smgr_rlocator);
 
 			tdereln->relKey = *int_key;
 			tdereln->encryption_status = RELATION_KEY_AVAILABLE;
@@ -190,7 +279,7 @@ tde_mdreadv(SMgrRelation reln, ForkNumber forknum, BlockNumber blocknum,
 		return;
 	else if (tdereln->encryption_status == RELATION_KEY_NOT_AVAILABLE)
 	{
-		InternalKey *int_key = GetSMGRRelationKey(reln->smgr_rlocator);
+		InternalKey *int_key = tde_smgr_get_key(&reln->smgr_rlocator);
 
 		tdereln->relKey = *int_key;
 		tdereln->encryption_status = RELATION_KEY_AVAILABLE;
@@ -259,10 +348,10 @@ tde_mdcreate(RelFileLocator relold, SMgrRelation reln, ForkNumber forknum, bool
 		 * Since event triggers do not fire on the standby or in recovery we
 		 * do not try to generate any new keys and instead trust the xlog.
 		 */
-		InternalKey *key = GetSMGRRelationKey(reln->smgr_rlocator);
+		InternalKey *key = tde_smgr_get_key(&reln->smgr_rlocator);
 
 		if (!isRedo && !key && tde_smgr_should_encrypt(&reln->smgr_rlocator, &relold))
-			key = pg_tde_create_smgr_key(&reln->smgr_rlocator);
+			key = tde_smgr_create_key(&reln->smgr_rlocator);
 
 		if (key)
 		{
@@ -291,7 +380,7 @@ tde_mdopen(SMgrRelation reln)
 
 	mdopen(reln);
 
-	if (IsSMGRRelationEncrypted(reln->smgr_rlocator))
+	if (tde_smgr_is_encrypted(&reln->smgr_rlocator))
 	{
 		tdereln->encryption_status = RELATION_KEY_NOT_AVAILABLE;
 	}
@@ -327,7 +416,68 @@ RegisterStorageMgr(void)
 {
 	if (storage_manager_id != MdSMgrId)
 		elog(FATAL, "Another storage manager was loaded before pg_tde. Multiple storage managers is unsupported.");
-	storage_manager_id = smgr_register(&tde_smgr, sizeof(TDESMgrRelation));
+	OurSMgrId = smgr_register(&tde_smgr, sizeof(TDESMgrRelation));
+	storage_manager_id = OurSMgrId;
+}
+
+static void
+tde_smgr_save_temp_key(const RelFileLocator *newrlocator, const InternalKey *key)
+{
+	TempRelKeyEntry *entry;
+	bool		found;
+
+	if (TempRelKeys == NULL)
+	{
+		HASHCTL		ctl;
+
+		ctl.keysize = sizeof(RelFileLocator);
+		ctl.entrysize = sizeof(TempRelKeyEntry);
+		TempRelKeys = hash_create("pg_tde temporary relation keys",
+								  INIT_TEMP_RELS,
+								  &ctl,
+								  HASH_ELEM | HASH_BLOBS);
+	}
+
+	entry = (TempRelKeyEntry *) hash_search(TempRelKeys,
+											newrlocator,
+											HASH_ENTER, &found);
+	Assert(!found);
+
+	entry->key = *key;
+}
+
+static InternalKey *
+tde_smgr_get_temp_key(const RelFileLocator *rel)
+{
+	TempRelKeyEntry *entry;
+
+	if (TempRelKeys == NULL)
+		return NULL;
+
+	entry = hash_search(TempRelKeys, rel, HASH_FIND, NULL);
+
+	if (entry)
+	{
+		InternalKey *key = palloc_object(InternalKey);
+
+		*key = entry->key;
+		return key;
+	}
+
+	return NULL;
+}
+
+static bool
+tde_smgr_has_temp_key(const RelFileLocator *rel)
+{
+	return TempRelKeys && hash_search(TempRelKeys, rel, HASH_FIND, NULL);
+}
+
+static void
+tde_smgr_remove_temp_key(const RelFileLocator *rel)
+{
+	Assert(TempRelKeys);
+	hash_search(TempRelKeys, rel, HASH_REMOVE, NULL);
 }
 
 /*
diff --git a/src/bin/pg_checksums/pg_checksums.c b/src/bin/pg_checksums/pg_checksums.c
index 73c8f320a60..66919c4afa2 100644
--- a/src/bin/pg_checksums/pg_checksums.c
+++ b/src/bin/pg_checksums/pg_checksums.c
@@ -138,10 +138,9 @@ pg_tde_init(const char *datadir)
 static bool
 is_pg_tde_encypted(Oid spcOid, Oid dbOid, RelFileNumber relNumber)
 {
-	RelFileLocator locator = {.spcOid = spcOid, .dbOid = dbOid,.relNumber = relNumber};
-	RelFileLocatorBackend rlocator = {.locator = locator,.backend = INVALID_PROC_NUMBER};
+	RelFileLocator locator = {.spcOid = spcOid,.dbOid = dbOid,.relNumber = relNumber};
 
-	return IsSMGRRelationEncrypted(rlocator);
+	return pg_tde_has_smgr_key(locator);
 }
 #endif