PG-1073, PG-1215: Improve pg_tde_is_encrypted (#363)

* Lookup is now based on Oid with the proper Oid getter, so explicit
  schemas are also supported (testcase added for this)
* There's an additional function written in C, which verifies that
  we do have an encryption key for the specified table. This can
  help in finding bugs with tde_heap, as now the function only
  returns true if we are actually writing encrypted data to the disk
  for the table.

expected/pg_tde_is_encrypted.out

@@ -51,6 +51,12 @@ SELECT pg_tde_is_encrypted('test_norm');
  f
 (1 row)
 
+SELECT pg_tde_is_encrypted('public.test_enc');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
 SELECT  key_provider_id, key_provider_name, principal_key_name
 		FROM pg_tde_principal_key_info();
  key_provider_id | key_provider_name |  principal_key_name   

expected/pg_tde_is_encrypted_basic.out

@@ -51,6 +51,12 @@ SELECT pg_tde_is_encrypted('test_norm');
  f
 (1 row)
 
+SELECT pg_tde_is_encrypted('public.test_enc');
+ pg_tde_is_encrypted 
+---------------------
+ t
+(1 row)
+
 SELECT  key_provider_id, key_provider_name, principal_key_name
 		FROM pg_tde_principal_key_info();
  key_provider_id | key_provider_name |  principal_key_name   

pg_tde--1.0.sql

@@ -230,15 +230,21 @@ RETURNS table_am_handler
 AS 'MODULE_PATHNAME'
 LANGUAGE C;
 
+CREATE FUNCTION pg_tde_internal_has_key(oid OID)
+RETURNS boolean
+AS 'MODULE_PATHNAME'
+LANGUAGE C;
+
 CREATE FUNCTION pg_tde_is_encrypted(table_name VARCHAR)
 RETURNS boolean
 AS $$
 SELECT EXISTS (
     SELECT 1
     FROM   pg_catalog.pg_class
-    WHERE  relname = table_name
+    WHERE  oid = table_name::regclass::oid
     AND    (relam = (SELECT oid FROM pg_catalog.pg_am WHERE amname = 'tde_heap_basic')
-        OR relam = (SELECT oid FROM pg_catalog.pg_am WHERE amname = 'tde_heap'))
+        OR (relam = (SELECT oid FROM pg_catalog.pg_am WHERE amname = 'tde_heap'))
+            AND pg_tde_internal_has_key(table_name::regclass::oid))
     )$$
 LANGUAGE SQL;
 

sql/pg_tde_is_encrypted.inc

@@ -23,6 +23,8 @@ SELECT amname FROM pg_class INNER JOIN pg_am ON pg_am.oid = pg_class.relam WHERE
 SELECT pg_tde_is_encrypted('test_enc');
 SELECT pg_tde_is_encrypted('test_norm');
 
+SELECT pg_tde_is_encrypted('public.test_enc');
+
 SELECT  key_provider_id, key_provider_name, principal_key_name
 		FROM pg_tde_principal_key_info();
 

src/common/pg_tde_utils.c

@@ -14,6 +14,9 @@
 #include "utils/snapmgr.h"
 #include "commands/defrem.h"
 #include "common/pg_tde_utils.h"
+#include "miscadmin.h"
+#include "catalog/tde_principal_key.h"
+#include "access/pg_tde_tdemap.h"
 
 #ifndef FRONTEND
 #include "access/genam.h"
@@ -31,6 +34,56 @@ get_tde_table_am_oid(void)
 	return get_table_am_oid("tde_heap", false);
 }
 
+PG_FUNCTION_INFO_V1(pg_tde_internal_has_key);
+Datum
+pg_tde_internal_has_key(PG_FUNCTION_ARGS)
+{
+	Oid tableOid = InvalidOid;
+	Oid	dbOid = MyDatabaseId;
+	TDEPrincipalKey* principalKey = NULL;
+	
+	if (!PG_ARGISNULL(0))
+	{
+		tableOid = PG_GETARG_OID(0);
+	}
+
+	if(tableOid == InvalidOid)
+	{
+		PG_RETURN_BOOL(false);
+	}
+
+	LWLockAcquire(tde_lwlock_enc_keys(), LW_SHARED);
+	principalKey = GetPrincipalKey(dbOid, LW_SHARED);
+	LWLockRelease(tde_lwlock_enc_keys());
+
+	if(principalKey == NULL)
+	{
+		PG_RETURN_BOOL(false);
+	}
+
+	{
+		LOCKMODE	lockmode = AccessShareLock;
+		Relation	rel = table_open(tableOid, lockmode);
+		RelKeyData *rkd;
+
+		if (
+			#ifdef PERCONA_EXT
+			rel->rd_rel->relam != get_tde_table_am_oid() && 
+			#endif
+			rel->rd_rel->relam != get_tde_basic_table_am_oid())
+		{
+			table_close(rel, lockmode);
+			PG_RETURN_BOOL(false);
+		}
+
+		rkd = GetSMGRRelationKey(rel->rd_locator);
+
+		table_close(rel, lockmode);
+		
+		PG_RETURN_BOOL(rkd != NULL);
+	}
+}
+
 /*
  * Returns the list of OIDs for all TDE tables in a database
  */