PG-1213 Add function that returns intrnal WAL keys ranges

Add user function that shows internal WAL keys ranges (start timeline,
start lsn, end timeline, end lsn)

Author: artemgavrilov

contrib/pg_tde/meson.build

@@ -127,6 +127,7 @@ tap_tests = [
   't/unlogged_tables.pl',
   't/wal_archiving.pl',
   't/wal_encrypt.pl',
+  't/wal_key_ranges.pl',
   't/wal_key_tli.pl',
 ]
 

contrib/pg_tde/pg_tde--1.0--2.0.sql

@@ -1,5 +1,17 @@
+-- Function to check if a WAL record is encrypted
 CREATE FUNCTION pg_tde_is_wal_record_encrypted(lsn pg_lsn, tli integer DEFAULT 0)
 RETURNS BOOLEAN
 LANGUAGE C
 AS 'MODULE_PATHNAME';
 REVOKE ALL ON FUNCTION pg_tde_is_wal_record_encrypted(pg_lsn, integer) FROM PUBLIC;
+
+-- Function to get internal WAL key ranges
+CREATE FUNCTION pg_tde_get_wal_key_ranges
+    (OUT start_tli integer,
+    OUT start_lsn pg_lsn,
+    OUT end_tli integer,
+    OUT end_lsn pg_lsn)
+RETURNS SETOF RECORD
+LANGUAGE C
+AS 'MODULE_PATHNAME';
+REVOKE ALL ON FUNCTION pg_tde_get_wal_key_ranges() FROM PUBLIC;

contrib/pg_tde/src/pg_tde.c

@@ -35,6 +35,8 @@
 
 PG_MODULE_MAGIC;
 
+#define PG_TDE_LIST_WAL_KEYS_RANGES_COLS 4
+
 static void pg_tde_init_data_dir(void);
 
 static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
@@ -44,6 +46,7 @@ PG_FUNCTION_INFO_V1(pg_tde_extension_initialize);
 PG_FUNCTION_INFO_V1(pg_tde_version);
 PG_FUNCTION_INFO_V1(pg_tdeam_handler);
 PG_FUNCTION_INFO_V1(pg_tde_is_wal_record_encrypted);
+PG_FUNCTION_INFO_V1(pg_tde_get_wal_key_ranges);
 
 static void
 tde_shmem_request(void)
@@ -201,3 +204,64 @@ pg_tde_is_wal_record_encrypted(PG_FUNCTION_ARGS)
 
 	PG_RETURN_BOOL(false);
 }
+
+/*
+ * Returns internal WAL key ranges.
+ */
+Datum
+pg_tde_get_wal_key_ranges(PG_FUNCTION_ARGS)
+{
+	Tuplestorestate *tupstore;
+	TupleDesc	tupdesc;
+	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	MemoryContext per_query_ctx;
+	MemoryContext oldcontext;
+	WALKeyCacheRec *keys;
+	WalLocation loc = {.tli = 0,.lsn = 0};
+
+	/* check to see if caller supports us returning a tuplestore */
+	if (rsinfo == NULL || !IsA(rsinfo, ReturnSetInfo))
+		ereport(ERROR,
+				errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				errmsg("set-valued function called in context that cannot accept a set"));
+	if (!(rsinfo->allowedModes & SFRM_Materialize))
+		ereport(ERROR,
+				errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				errmsg("materialize mode required, but it is not allowed in this context"));
+
+	/* Switch into long-lived context to construct returned data structures */
+	per_query_ctx = rsinfo->econtext->ecxt_per_query_memory;
+	oldcontext = MemoryContextSwitchTo(per_query_ctx);
+
+	/* Build a tuple descriptor for our result type */
+	if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
+		elog(ERROR, "return type must be a row type");
+
+	tupstore = tuplestore_begin_heap(true, false, work_mem);
+	rsinfo->returnMode = SFRM_Materialize;
+	rsinfo->setResult = tupstore;
+	rsinfo->setDesc = tupdesc;
+
+	MemoryContextSwitchTo(oldcontext);
+
+	keys = pg_tde_fetch_wal_keys(loc);
+
+	for (WALKeyCacheRec *curr_key = keys; curr_key != NULL; curr_key = curr_key->next)
+	{
+		Datum		values[PG_TDE_LIST_WAL_KEYS_RANGES_COLS] = {0};
+		bool		nulls[PG_TDE_LIST_WAL_KEYS_RANGES_COLS] = {0};
+		int			i = 0;
+
+		if (curr_key->key.type != WAL_KEY_TYPE_ENCRYPTED)
+			continue;
+
+		values[i++] = Int64GetDatum(curr_key->start.tli);
+		values[i++] = Int64GetDatum(curr_key->start.lsn);
+		values[i++] = Int64GetDatum(curr_key->end.tli);
+		values[i++] = Int64GetDatum(curr_key->end.lsn);
+
+		tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+	}
+
+	return (Datum) 0;
+}

contrib/pg_tde/t/expected/wal_encrypt.out

@@ -103,7 +103,7 @@ SHOW pg_tde.wal_encrypt;
 (1 row)
 
 INSERT INTO test_wal (k) VALUES (7), (8);
-SELECT pg_tde_is_wal_record_encrypted('0/15883E8'::pg_lsn);
+SELECT pg_tde_is_wal_record_encrypted(pg_current_wal_lsn());
  pg_tde_is_wal_record_encrypted 
 --------------------------------
  t

contrib/pg_tde/t/wal_encrypt.pl

@@ -97,8 +97,7 @@
 
 PGTDE::psql($node, 'postgres', 'INSERT INTO test_wal (k) VALUES (7), (8);');
 
-PGTDE::psql($node, 'postgres',
-	"SELECT pg_tde_is_wal_record_encrypted(pg_current_wal_lsn());");
+PGTDE::psql($node, 'postgres', "SELECT pg_tde_is_wal_record_encrypted(pg_current_wal_lsn());");
 
 PGTDE::psql($node, 'postgres',
 	"SELECT data FROM pg_logical_slot_get_changes('tde_slot', NULL, NULL);");

contrib/pg_tde/t/wal_key_ranges.pl

@@ -0,0 +1,99 @@
+#!/usr/bin/perl
+
+use strict;
+use warnings;
+use File::Basename;
+use Test::More;
+use lib 't';
+use pgtde;
+
+PGTDE::setup_files_dir(basename($0));
+
+unlink('/tmp/wal_key_ranges.per');
+
+my $psql_out = '';
+
+my $node = PostgreSQL::Test::Cluster->new('main');
+$node->init;
+$node->append_conf('postgresql.conf', "shared_preload_libraries = 'pg_tde'");
+$node->append_conf('postgresql.conf', "wal_level = 'logical'");
+# We don't test that it can't start: the test framework doesn't have an easy way to do this
+#$node->append_conf('postgresql.conf', "pg_tde.wal_encrypt = 1");
+$node->start;
+
+# Create and configure pg_tde extension
+$node->psql('postgres', "CREATE EXTENSION pg_tde;");
+
+$node->psql('postgres',
+	"SELECT pg_tde_add_global_key_provider_file('file-keyring', '/tmp/wal_key_ranges.per');"
+);
+
+$node->psql('postgres',
+	"SELECT pg_tde_create_key_using_global_key_provider('server-key', 'file-keyring');"
+);
+$node->psql('postgres',
+	"SELECT pg_tde_set_server_key_using_global_key_provider('server-key', 'file-keyring');"
+);
+
+$node->psql('postgres', 'SELECT pg_tde_verify_server_key();');
+
+# Create test table and enable WAL encryption
+$node->psql('postgres',
+	'CREATE TABLE test_wal (id SERIAL, k INTEGER, PRIMARY KEY (id));');
+
+$node->psql('postgres', 'ALTER SYSTEM SET pg_tde.wal_encrypt = on;');
+
+$node->restart;
+
+# Insert some data to generate WAL records and check that WAL records are encrypted
+$node->psql('postgres', 'INSERT INTO test_wal (k) VALUES (1), (2);');
+
+my $enc_lsn = $node->safe_psql('postgres', "SELECT pg_current_wal_lsn();");
+$node->psql(
+	'postgres',
+	"SELECT pg_tde_is_wal_record_encrypted('$enc_lsn'::pg_lsn);",
+	stdout => \$psql_out);
+is($psql_out, 't', "Check that WAL record is encrypted");
+
+# Force PG to switch to a new WAL segment to avoid situation when we decrypt
+# non-full WAL page on WAL encryption off.
+$node->psql('postgres', 'SELECT pg_switch_wal();');
+$node->psql('postgres', 'ALTER SYSTEM SET pg_tde.wal_encrypt = off;');
+
+$node->restart;
+
+# Insert more data to generate WAL records and check that they are not encrypted
+$node->psql('postgres', 'INSERT INTO test_wal (k) VALUES (3), (4);');
+
+my $dec_lsn = $node->safe_psql('postgres', "SELECT pg_current_wal_lsn();");
+$node->psql(
+	'postgres',
+	"SELECT pg_tde_is_wal_record_encrypted('$dec_lsn'::pg_lsn);",
+	stdout => \$psql_out);
+is($psql_out, 'f', "Check that WAL record is not encrypted");
+
+# Check that previously encrypted record is still encrypted
+$node->psql(
+	'postgres',
+	"SELECT pg_tde_is_wal_record_encrypted('$enc_lsn'::pg_lsn);",
+	stdout => \$psql_out);
+is($psql_out, 't', "Check that WAL record is still encrypted");
+
+# Check that WAL records that we recoreded before match the key ranges. We doint this
+# that way because LSN numbers are not stable across different runs of the test.
+my $key_count = $node->safe_psql(
+	'postgres', "SELECT count(*) FROM pg_tde_get_wal_key_ranges() 
+            WHERE start_lsn <= '$enc_lsn'::pg_lsn
+            AND end_lsn > '$enc_lsn'::pg_lsn
+            AND end_lsn <= '$dec_lsn'::pg_lsn
+            AND start_tli = 1
+            AND end_tli = 1;");
+is($key_count, 1, "Check that there is one key range for encrypted record");
+
+$node->psql('postgres', 'DROP EXTENSION pg_tde;');
+
+$node->stop;
+
+
+done_testing();
+