Use proper tablespace ID in key rotation

meson.build

@@ -110,6 +110,7 @@ tap_tests = [
       't/005_multiple_extensions.pl',
       't/006_remote_vault_config.pl',
       't/007_access_control.pl',
+      't/009_key_rotate_tablespace.pl',
     ]
 
 if get_variable('percona_ext', false)

src/access/pg_tde_tdemap.c

@@ -727,13 +727,9 @@ pg_tde_perform_rotate_key(TDEPrincipalKey *principal_key, TDEPrincipalKey *new_p
 		if (found == false)
 			continue;
 
-		/*
-		 * Set the relNumber of rlocator. Ignore the tablespace Oid since we
-		 * only place our files under the default.
-		 */
 		rloc.relNumber = read_map_entry.relNumber;
 		rloc.dbOid = principal_key->keyInfo.databaseId;
-		rloc.spcOid = DEFAULTTABLESPACE_OID;	/* TODO. Revisit */
+		rloc.spcOid = principal_key->keyInfo.tablespaceId;
 
 		/* Let's get the decrypted key and re-encrypt it with the new key. */
 		enc_rel_key_data[OLD_PRINCIPAL_KEY] = pg_tde_read_one_keydata(k_fd[OLD_PRINCIPAL_KEY], key_index[OLD_PRINCIPAL_KEY], principal_key);

t/009_key_rotate_tablespace.pl

@@ -0,0 +1,96 @@
+#!/usr/bin/perl
+
+use strict;
+use warnings;
+use File::Basename;
+use File::Compare;
+use File::Copy;
+use Test::More;
+use lib 't';
+use pgtde;
+
+# Get file name and CREATE out file name and dirs WHERE requried
+PGTDE::setup_files_dir(basename($0));
+
+my ($cmdret, $stdout);
+
+# CREATE new PostgreSQL node and do initdb
+my $node = PGTDE->pgtde_init_pg();
+my $pgdata = $node->data_dir;
+
+# UPDATE postgresql.conf to include/load pg_tde library
+open my $conf, '>>', "$pgdata/postgresql.conf";
+print $conf "shared_preload_libraries = 'pg_tde'\n";
+close $conf;
+
+# Start server
+my $rt_value = $node->start;
+ok($rt_value == 1, "Start Server");
+
+$node->safe_psql('postgres',
+	q{
+SET allow_in_place_tablespaces = true;
+CREATE TABLESPACE test_tblspace LOCATION '';
+CREATE DATABASE tbc TABLESPACE = test_tblspace;
+});
+
+$stdout = $node->safe_psql('tbc',
+	q{
+CREATE EXTENSION pg_tde;
+SELECT pg_tde_add_key_provider_file('file-vault','/tmp/pg_tde_test_keyring.per');
+SELECT pg_tde_set_principal_key('test-db-principal-key','file-vault');
+
+CREATE TABLE country_table (
+     country_id        serial primary key,
+     country_name    text unique not null,
+     continent        text not null
+) USING tde_heap_basic;
+
+INSERT INTO country_table (country_name, continent)
+     VALUES ('Japan', 'Asia'),
+            ('UK', 'Europe'),
+            ('USA', 'North America');
+
+SELECT * FROM country_table;
+
+}, extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+
+$cmdret = $node->psql('tbc', "SELECT pg_tde_rotate_principal_key('new-k', 'file-vault');", extra_params => ['-a']);
+ok($cmdret == 0, "ROTATE KEY");
+PGTDE::append_to_file($stdout);
+
+# Restart the server
+PGTDE::append_to_file("-- server restart");
+$node->stop();
+
+$rt_value = $node->start();
+ok($rt_value == 1, "Restart Server");
+
+$stdout = $node->safe_psql('tbc', 'SELECT * FROM country_table;', extra_params => ['-a']);
+PGTDE::append_to_file($stdout);
+
+
+# DROP EXTENSION
+$stdout = $node->safe_psql('tbc', 'DROP EXTENSION pg_tde CASCADE;', extra_params => ['-a']);
+ok($cmdret == 0, "DROP PGTDE EXTENSION");
+PGTDE::append_to_file($stdout);
+
+$stdout = $node->safe_psql('postgres', q{
+DROP DATABASE tbc;
+DROP TABLESPACE test_tblspace;
+}, extra_params => ['-a']);
+ok($cmdret == 0, "DROP DATABSE");
+PGTDE::append_to_file($stdout);
+# Stop the server
+$node->stop();
+
+# compare the expected and out file
+my $compare = PGTDE->compare_results();
+
+# Test/check if expected and result/out file match. If Yes, test passes.
+is($compare,0,"Compare Files: $PGTDE::expected_filename_with_path and $PGTDE::out_filename_with_path files.");
+
+# Done testing for this testcase file.
+done_testing();

t/expected/009_key_rotate_tablespace.out

@@ -0,0 +1,44 @@
+CREATE EXTENSION pg_tde;
+SELECT pg_tde_add_key_provider_file('file-vault','/tmp/pg_tde_test_keyring.per');
+1
+SELECT pg_tde_set_principal_key('test-db-principal-key','file-vault');
+t
+CREATE TABLE country_table (
+     country_id        serial primary key,
+     country_name    text unique not null,
+     continent        text not null
+) USING tde_heap_basic;
+INSERT INTO country_table (country_name, continent)
+     VALUES ('Japan', 'Asia'),
+            ('UK', 'Europe'),
+            ('USA', 'North America');
+SELECT * FROM country_table;
+1|Japan|Asia
+2|UK|Europe
+3|USA|North America
+CREATE EXTENSION pg_tde;
+SELECT pg_tde_add_key_provider_file('file-vault','/tmp/pg_tde_test_keyring.per');
+1
+SELECT pg_tde_set_principal_key('test-db-principal-key','file-vault');
+t
+CREATE TABLE country_table (
+     country_id        serial primary key,
+     country_name    text unique not null,
+     continent        text not null
+) USING tde_heap_basic;
+INSERT INTO country_table (country_name, continent)
+     VALUES ('Japan', 'Asia'),
+            ('UK', 'Europe'),
+            ('USA', 'North America');
+SELECT * FROM country_table;
+1|Japan|Asia
+2|UK|Europe
+3|USA|North America
+-- server restart
+SELECT * FROM country_table;
+1|Japan|Asia
+2|UK|Europe
+3|USA|North America
+DROP EXTENSION pg_tde CASCADE;
+DROP DATABASE tbc;
+DROP TABLESPACE test_tblspace;