Merge pull request #287 from dutow/pg952

PG-952: Handle non existent principal key in tde_heap
percona · Sep 30, 2024 · d715cb1 · d715cb1
2 parents cb7557f + 6d02b5e
commit d715cb1
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 3 deletions.
diff --git a/src/access/pg_tde_tdemap.c b/src/access/pg_tde_tdemap.c
@@ -152,7 +152,7 @@ pg_tde_create_key_map_entry(const RelFileLocator *newrlocator)
 	{
 		LWLockRelease(lock_pk);
 		ereport(ERROR,
-				(errmsg("failed to retrieve principal key")));
+				(errmsg("failed to retrieve principal key. Create one using pg_tde_set_principal_key before using encrypted tables.")));
 
 		return NULL;
 	}
@@ -870,7 +870,7 @@ pg_tde_get_key_from_file(const RelFileLocator *rlocator)
 	{
 		LWLockRelease(lock_pk);
 		ereport(ERROR,
-				(errmsg("failed to retrieve principal key")));
+				(errmsg("failed to retrieve principal key. Create one using pg_tde_set_principal_key before using encrypted tables.")));
 	}
 
 	/* Get the file paths */

diff --git a/src/catalog/tde_principal_key.c b/src/catalog/tde_principal_key.c
@@ -257,7 +257,7 @@ set_principal_key_with_keyring(const char *key_name, GenericKeyring *keyring,
             LWLockRelease(lock_files);
 
             ereport(ERROR,
-                    (errmsg("failed to retrieve principal key")));
+                    (errmsg("failed to retrieve principal key. Create one using pg_tde_set_principal_key before using encrypted tables.")));
         }
 
         principalKey->keyLength = keyInfo->data.len;

diff --git a/src/pg_tde_event_capture.c b/src/pg_tde_event_capture.c
@@ -21,6 +21,9 @@
 #include "commands/event_trigger.h"
 #include "common/pg_tde_utils.h"
 #include "pg_tde_event_capture.h"
+#include "commands/tablespace.h"
+#include "catalog/tde_principal_key.h"
+#include "miscadmin.h"
 
 /* Global variable that gets set at ddl start and cleard out at ddl end*/
 TdeCreateEvent tdeCurrentCreateEvent = {.relation = NULL};
@@ -97,6 +100,8 @@ pg_tde_ddl_command_start_capture(PG_FUNCTION_ARGS)
 	else if (IsA(parsetree, CreateStmt))
 	{
 		CreateStmt *stmt = (CreateStmt *) parsetree;
+		TDEPrincipalKey * principal_key;
+		Oid tablespace_oid;
 
 		tdeCurrentCreateEvent.eventType = TDE_TABLE_CREATE_EVENT;
 		tdeCurrentCreateEvent.relation = stmt->relation;
@@ -105,6 +110,17 @@ pg_tde_ddl_command_start_capture(PG_FUNCTION_ARGS)
 		{
 			tdeCurrentCreateEvent.encryptMode = true;
 		}
+
+		tablespace_oid = stmt->tablespacename != NULL ? get_tablespace_oid(stmt->tablespacename, false) 
+						 : MyDatabaseTableSpace;  
+		principal_key = GetPrincipalKey(MyDatabaseId, tablespace_oid);
+		if (principal_key == NULL)
+		{
+			ereport(ERROR,
+					(errmsg("failed to retrieve principal key. Create one using pg_tde_set_principal_key before using encrypted tables.")));
+
+		}
+
 	}
 #endif
 	PG_RETURN_NULL();