pyrox.dev
1import ./make-test-python.nix ({ pkgs, lib, ... }: {
2 name = "pg_anonymizer";
3 meta.maintainers = lib.teams.flyingcircus.members;
4
5 nodes.machine = { pkgs, ... }: {
6 environment.systemPackages = [ pkgs.pg-dump-anon ];
7 services.postgresql = {
8 enable = true;
9 extraPlugins = ps: [ ps.anonymizer ];
10 settings.shared_preload_libraries = [ "anon" ];
11 };
12 };
13
14 testScript = ''
15 start_all()
16 machine.wait_for_unit("multi-user.target")
17 machine.wait_for_unit("postgresql.service")
18
19 with subtest("Setup"):
20 machine.succeed("sudo -u postgres psql --command 'create database demo'")
21 machine.succeed(
22 "sudo -u postgres psql -d demo -f ${pkgs.writeText "init.sql" ''
23 create extension anon cascade;
24 select anon.init();
25 create table player(id serial, name text, points int);
26 insert into player(id,name,points) values (1,'Foo', 23);
27 insert into player(id,name,points) values (2,'Bar',42);
28 security label for anon on column player.name is 'MASKED WITH FUNCTION anon.fake_last_name();';
29 security label for anon on column player.points is 'MASKED WITH VALUE NULL';
30 ''}"
31 )
32
33 def get_player_table_contents():
34 return [
35 x.split(',') for x in machine.succeed("sudo -u postgres psql -d demo --csv --command 'select * from player'").splitlines()[1:]
36 ]
37
38 def check_anonymized_row(row, id, original_name):
39 assert row[0] == id, f"Expected first row to have ID {id}, but got {row[0]}"
40 assert row[1] != original_name, f"Expected first row to have a name other than {original_name}"
41 assert not bool(row[2]), "Expected points to be NULL in first row"
42
43 def find_xsv_in_dump(dump, sep=','):
44 """
45 Expecting to find a CSV (for pg_dump_anon) or TSV (for pg_dump) structure, looking like
46
47 COPY public.player ...
48 1,Shields,
49 2,Salazar,
50 \.
51
52 in the given dump (the commas are tabs in case of pg_dump).
53 Extract the CSV lines and split by `sep`.
54 """
55
56 try:
57 from itertools import dropwhile, takewhile
58 return [x.split(sep) for x in list(takewhile(
59 lambda x: x != "\\.",
60 dropwhile(
61 lambda x: not x.startswith("COPY public.player"),
62 dump.splitlines()
63 )
64 ))[1:]]
65 except:
66 print(f"Dump to process: {dump}")
67 raise
68
69 def check_original_data(output):
70 assert output[0] == ['1','Foo','23'], f"Expected first row from player table to be 1,Foo,23; got {output[0]}"
71 assert output[1] == ['2','Bar','42'], f"Expected first row from player table to be 2,Bar,42; got {output[1]}"
72
73 def check_anonymized_rows(output):
74 check_anonymized_row(output[0], '1', 'Foo')
75 check_anonymized_row(output[1], '2', 'Bar')
76
77 with subtest("Check initial state"):
78 check_original_data(get_player_table_contents())
79
80 with subtest("Anonymous dumps"):
81 check_original_data(find_xsv_in_dump(
82 machine.succeed("sudo -u postgres pg_dump demo"),
83 sep='\t'
84 ))
85 check_anonymized_rows(find_xsv_in_dump(
86 machine.succeed("sudo -u postgres pg_dump_anon -U postgres -h /run/postgresql -d demo"),
87 sep=','
88 ))
89
90 with subtest("Anonymize"):
91 machine.succeed("sudo -u postgres psql -d demo --command 'select anon.anonymize_database();'")
92 check_anonymized_rows(get_player_table_contents())
93 '';
94})