Skip to content

Commit 7ca7be7

Browse files
artemgavrilovartemgavrilov
committed
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)
1 parent 9ec9225 commit 7ca7be7

File tree

5 files changed

+177
-1
lines changed

5 files changed

+177
-1
lines changed

‎contrib/pg_tde/meson.build

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,7 @@ tap_tests = [
128128
't/unlogged_tables.pl',
129129
't/wal_archiving.pl',
130130
't/wal_encrypt.pl',
131+
't/wal_key_ranges.pl',
131132
't/wal_key_tli.pl',
132133
]
133134

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,17 @@
1+
-- Function to check if a WAL record is encrypted
12
CREATE FUNCTION pg_tde_is_wal_record_encrypted(lsn pg_lsn, tli integer DEFAULT 0)
23
RETURNS BOOLEAN
34
LANGUAGE C
45
AS 'MODULE_PATHNAME';
56
REVOKE ALL ON FUNCTION pg_tde_is_wal_record_encrypted(pg_lsn, integer) FROM PUBLIC;
7+
8+
-- Function to get internal WAL key ranges
9+
CREATE FUNCTION pg_tde_get_wal_key_ranges
10+
(OUT start_tli integer,
11+
OUT start_lsn pg_lsn,
12+
OUT end_tli integer,
13+
OUT end_lsn pg_lsn)
14+
RETURNS SETOF RECORD
15+
LANGUAGE C
16+
AS 'MODULE_PATHNAME';
17+
REVOKE ALL ON FUNCTION pg_tde_get_wal_key_ranges() FROM PUBLIC;

‎contrib/pg_tde/src/pg_tde.c

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@
3535

3636
PG_MODULE_MAGIC;
3737

38+
#define PG_TDE_LIST_WAL_KEYS_RANGES_COLS 4
39+
3840
static void pg_tde_init_data_dir(void);
3941

4042
static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
@@ -44,6 +46,7 @@ PG_FUNCTION_INFO_V1(pg_tde_extension_initialize);
4446
PG_FUNCTION_INFO_V1(pg_tde_version);
4547
PG_FUNCTION_INFO_V1(pg_tdeam_handler);
4648
PG_FUNCTION_INFO_V1(pg_tde_is_wal_record_encrypted);
49+
PG_FUNCTION_INFO_V1(pg_tde_get_wal_key_ranges);
4750

4851
static void
4952
tde_shmem_request(void)
@@ -201,3 +204,64 @@ pg_tde_is_wal_record_encrypted(PG_FUNCTION_ARGS)
201204

202205
PG_RETURN_BOOL(false);
203206
}
207+
208+
/*
209+
* Returns internal WAL key ranges.
210+
*/
211+
Datum
212+
pg_tde_get_wal_key_ranges(PG_FUNCTION_ARGS)
213+
{
214+
Tuplestorestate *tupstore;
215+
TupleDesc tupdesc;
216+
ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
217+
MemoryContext per_query_ctx;
218+
MemoryContext oldcontext;
219+
WALKeyCacheRec *keys;
220+
WalLocation loc = {.tli = 0,.lsn = 0};
221+
222+
/* check to see if caller supports us returning a tuplestore */
223+
if (rsinfo == NULL || !IsA(rsinfo, ReturnSetInfo))
224+
ereport(ERROR,
225+
errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
226+
errmsg("set-valued function called in context that cannot accept a set"));
227+
if (!(rsinfo->allowedModes & SFRM_Materialize))
228+
ereport(ERROR,
229+
errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
230+
errmsg("materialize mode required, but it is not allowed in this context"));
231+
232+
/* Switch into long-lived context to construct returned data structures */
233+
per_query_ctx = rsinfo->econtext->ecxt_per_query_memory;
234+
oldcontext = MemoryContextSwitchTo(per_query_ctx);
235+
236+
/* Build a tuple descriptor for our result type */
237+
if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
238+
elog(ERROR, "return type must be a row type");
239+
240+
tupstore = tuplestore_begin_heap(true, false, work_mem);
241+
rsinfo->returnMode = SFRM_Materialize;
242+
rsinfo->setResult = tupstore;
243+
rsinfo->setDesc = tupdesc;
244+
245+
MemoryContextSwitchTo(oldcontext);
246+
247+
keys = pg_tde_fetch_wal_keys(loc);
248+
249+
for (WALKeyCacheRec *curr_key = keys; curr_key != NULL; curr_key = curr_key->next)
250+
{
251+
Datum values[PG_TDE_LIST_WAL_KEYS_RANGES_COLS] = {0};
252+
bool nulls[PG_TDE_LIST_WAL_KEYS_RANGES_COLS] = {0};
253+
int i = 0;
254+
255+
if (curr_key->key.type != WAL_KEY_TYPE_ENCRYPTED)
256+
continue;
257+
258+
values[i++] = Int64GetDatum(curr_key->start.tli);
259+
values[i++] = Int64GetDatum(curr_key->start.lsn);
260+
values[i++] = Int64GetDatum(curr_key->end.tli);
261+
values[i++] = Int64GetDatum(curr_key->end.lsn);
262+
263+
tuplestore_putvalues(tupstore, tupdesc, values, nulls);
264+
}
265+
266+
return (Datum) 0;
267+
}

‎contrib/pg_tde/t/expected/wal_encrypt.out

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ SHOW pg_tde.wal_encrypt;
103103
(1 row)
104104

105105
INSERT INTO test_wal (k) VALUES (7), (8);
106-
SELECT pg_tde_is_wal_record_encrypted('0/15883E8'::pg_lsn);
106+
SELECT pg_tde_is_wal_record_encrypted(pg_current_wal_lsn());
107107
pg_tde_is_wal_record_encrypted
108108
--------------------------------
109109
t

‎contrib/pg_tde/t/wal_key_ranges.pl

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
#!/usr/bin/perl
2+
3+
use strict;
4+
use warnings;
5+
use File::Basename;
6+
use Test::More;
7+
use lib 't';
8+
use pgtde;
9+
10+
PGTDE::setup_files_dir(basename($0));
11+
12+
unlink('/tmp/wal_key_ranges.per');
13+
14+
my $psql_out = '';
15+
16+
my $node = PostgreSQL::Test::Cluster->new('main');
17+
$node->init;
18+
$node->append_conf('postgresql.conf', "shared_preload_libraries = 'pg_tde'");
19+
$node->append_conf('postgresql.conf', "wal_level = 'logical'");
20+
# We don't test that it can't start: the test framework doesn't have an easy way to do this
21+
#$node->append_conf('postgresql.conf', "pg_tde.wal_encrypt = 1");
22+
$node->start;
23+
24+
# Create and configure pg_tde extension
25+
$node->psql('postgres', "CREATE EXTENSION pg_tde;");
26+
27+
$node->psql('postgres',
28+
"SELECT pg_tde_add_global_key_provider_file('file-keyring', '/tmp/wal_key_ranges.per');"
29+
);
30+
31+
$node->psql('postgres',
32+
"SELECT pg_tde_create_key_using_global_key_provider('server-key', 'file-keyring');"
33+
);
34+
$node->psql('postgres',
35+
"SELECT pg_tde_set_server_key_using_global_key_provider('server-key', 'file-keyring');"
36+
);
37+
38+
$node->psql('postgres', 'SELECT pg_tde_verify_server_key();');
39+
40+
# Create test table and enable WAL encryption
41+
$node->psql('postgres',
42+
'CREATE TABLE test_wal (id SERIAL, k INTEGER, PRIMARY KEY (id));');
43+
44+
$node->psql('postgres', 'ALTER SYSTEM SET pg_tde.wal_encrypt = on;');
45+
46+
$node->restart;
47+
48+
# Insert some data to generate WAL records and check that WAL records are encrypted
49+
$node->psql('postgres', 'INSERT INTO test_wal (k) VALUES (1), (2);');
50+
51+
my $enc_lsn = $node->safe_psql('postgres', "SELECT pg_current_wal_lsn();");
52+
$node->psql(
53+
'postgres',
54+
"SELECT pg_tde_is_wal_record_encrypted('$enc_lsn'::pg_lsn);",
55+
stdout => \$psql_out);
56+
is($psql_out, 't', "Check that WAL record is encrypted");
57+
58+
# Force PG to switch to a new WAL segment to avoid situation when we decrypt
59+
# non-full WAL page on WAL encryption off.
60+
$node->psql('postgres', 'SELECT pg_switch_wal();');
61+
$node->psql('postgres', 'ALTER SYSTEM SET pg_tde.wal_encrypt = off;');
62+
63+
$node->restart;
64+
65+
# Insert more data to generate WAL records and check that they are not encrypted
66+
$node->psql('postgres', 'INSERT INTO test_wal (k) VALUES (3), (4);');
67+
68+
my $dec_lsn = $node->safe_psql('postgres', "SELECT pg_current_wal_lsn();");
69+
$node->psql(
70+
'postgres',
71+
"SELECT pg_tde_is_wal_record_encrypted('$dec_lsn'::pg_lsn);",
72+
stdout => \$psql_out);
73+
is($psql_out, 'f', "Check that WAL record is not encrypted");
74+
75+
# Check that previously encrypted record is still encrypted
76+
$node->psql(
77+
'postgres',
78+
"SELECT pg_tde_is_wal_record_encrypted('$enc_lsn'::pg_lsn);",
79+
stdout => \$psql_out);
80+
is($psql_out, 't', "Check that WAL record is still encrypted");
81+
82+
# Check that WAL records that we recoreded before match the key ranges. We doint this
83+
# that way because LSN numbers are not stable across different runs of the test.
84+
my $key_count = $node->safe_psql(
85+
'postgres', "SELECT count(*) FROM pg_tde_get_wal_key_ranges()
86+
WHERE start_lsn <= '$enc_lsn'::pg_lsn
87+
AND end_lsn > '$enc_lsn'::pg_lsn
88+
AND end_lsn <= '$dec_lsn'::pg_lsn
89+
AND start_tli = 1
90+
AND end_tli = 1;");
91+
is($key_count, 1, "Check that there is one key range for encrypted record");
92+
93+
$node->psql('postgres', 'DROP EXTENSION pg_tde;');
94+
95+
$node->stop;
96+
97+
98+
done_testing();
99+

0 commit comments

Comments
 (0)