Skip to content

Commit bc94e8d

Browse files
fblackburn1fblackburn1
committed
add dump-options to modify pg_dump command
1 parent f30f926 commit bc94e8d

File tree

5 files changed

+37
-21
lines changed

5 files changed

+37
-21
lines changed

‎README.rst‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,9 @@ Usage
8585
--port PORT Port of the database
8686
--dry-run Don't commit changes made on the database
8787
--dump-file DUMP_FILE
88-
Create a database dump file with the given name
88+
Create a database dump file with the given name
89+
--dump-options DUMP_OPTIONS
90+
Options to pass to the pg_dump command
8991
--init-sql INIT_SQL SQL to run before starting anonymization
9092
--parallel Data anonymization is done in parallel
9193

‎pganonymize/cli.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,8 @@ def get_arg_parser():
4646
parser.add_argument('--dry-run', action='store_true', help='Don\'t commit changes made on the database',
4747
default=False)
4848
parser.add_argument('--dump-file', help='Create a database dump file with the given name')
49+
parser.add_argument('--dump-options', help='Options to pass to the pg_dump command',
50+
default='--format custom --compress 9')
4951
parser.add_argument('--init-sql', help='SQL to run before starting anonymization', default=False)
5052
parser.add_argument(
5153
'--parallel',
@@ -101,4 +103,4 @@ def main(args):
101103
logging.info('Anonymization took {:.2f}s'.format(end_time - start_time))
102104

103105
if args.dump_file:
104-
create_database_dump(args.dump_file, pg_args)
106+
create_database_dump(args.dump_file, pg_args, args.dump_options)

‎pganonymize/utils.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -288,7 +288,7 @@ def truncate_tables(connection):
288288
cursor.close()
289289

290290

291-
def create_database_dump(filename, db_args):
291+
def create_database_dump(filename, db_args, dump_args):
292292
"""
293293
Create a dump file from the current database.
294294
@@ -299,10 +299,11 @@ def create_database_dump(filename, db_args):
299299
if db_args.get('password'):
300300
env_vars += 'PGPASSWORD={password}'.format(password=db_args['password'])
301301
arguments = '--dbname {dbname} --username {user} --host {host} --port {port}'.format(**db_args)
302-
cmd = '{env_vars}pg_dump --format custom --compress 9 {args} --file {filename}'.format(
302+
cmd = '{env_vars}pg_dump {dump_args} {db_args} --file {filename}'.format(
303303
env_vars='{} '.format(env_vars) if env_vars else '',
304-
args=arguments,
305-
filename=filename
304+
dump_args=dump_args,
305+
db_args=arguments,
306+
filename=filename,
306307
)
307308
logging.info('Creating database dump file "%s"', filename)
308309
subprocess.call(cmd, shell=True)

‎tests/test_cli.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ class TestCli(object):
1717
@pytest.mark.parametrize('cli_args, expected, expected_executes, commit_calls, call_dump', [
1818
['--host localhost --port 5432 --user root --password my-cool-password --dbname db --schema ./tests/schemes/valid_schema.yml -v --init-sql "set work_mem=\'1GB\'"', # noqa
1919
Namespace(verbose=1, list_providers=False, schema='./tests/schemes/valid_schema.yml', dbname='db', user='root',
20-
password='my-cool-password', host='localhost', port='5432', dry_run=False, dump_file=None, init_sql="set work_mem='1GB'", parallel=False), # noqa
20+
password='my-cool-password', host='localhost', port='5432', dry_run=False, dump_file=None, dump_options='--format custom --compress 9', init_sql="set work_mem='1GB'", parallel=False), # noqa
2121
[call("set work_mem='1GB'"),
2222
call('TRUNCATE TABLE "django_session"'),
2323
call('SELECT COUNT(*) FROM "auth_user"'),
@@ -32,7 +32,7 @@ class TestCli(object):
3232
],
3333
['--dry-run --host localhost --port 5432 --user root --password my-cool-password --dbname db --schema ./tests/schemes/valid_schema.yml -v --init-sql "set work_mem=\'1GB\'"', # noqa
3434
Namespace(verbose=1, list_providers=False, schema='./tests/schemes/valid_schema.yml', dbname='db', user='root',
35-
password='my-cool-password', host='localhost', port='5432', dry_run=True, dump_file=None, init_sql="set work_mem='1GB'", parallel=False), # noqa
35+
password='my-cool-password', host='localhost', port='5432', dry_run=True, dump_file=None, dump_options='--format custom --compress 9', init_sql="set work_mem='1GB'", parallel=False), # noqa
3636
[call("set work_mem='1GB'"),
3737
call('TRUNCATE TABLE "django_session"'),
3838
call('SELECT "id", "first_name", "last_name", "email" FROM "auth_user" LIMIT 100'),
@@ -42,9 +42,9 @@ class TestCli(object):
4242
],
4343
0, []
4444
],
45-
['--dump-file ./dump.sql --host localhost --port 5432 --user root --password my-cool-password --dbname db --schema ./tests/schemes/valid_schema.yml -v --init-sql "set work_mem=\'1GB\'"', # noqa
45+
['--dump-file ./dump.sql --dump-options "--format plain" --host localhost --port 5432 --user root --password my-cool-password --dbname db --schema ./tests/schemes/valid_schema.yml -v --init-sql "set work_mem=\'1GB\'"', # noqa
4646
Namespace(verbose=1, list_providers=False, schema='./tests/schemes/valid_schema.yml', dbname='db', user='root',
47-
password='my-cool-password', host='localhost', port='5432', dry_run=False, dump_file='./dump.sql', init_sql="set work_mem='1GB'", parallel=False), # noqa
47+
password='my-cool-password', host='localhost', port='5432', dry_run=False, dump_file='./dump.sql', dump_options='--format plain', init_sql="set work_mem='1GB'", parallel=False), # noqa
4848
[
4949
call("set work_mem='1GB'"),
5050
call('TRUNCATE TABLE "django_session"'),
@@ -56,14 +56,14 @@ class TestCli(object):
5656
call('UPDATE "auth_user" t SET "first_name" = s."first_name", "last_name" = s."last_name", "email" = s."email" FROM "tmp_auth_user" s WHERE t."id" = s."id"') # noqa
5757
],
5858
1,
59-
[call('PGPASSWORD=my-cool-password pg_dump --format custom --compress 9 --dbname db --username root --host localhost --port 5432 --file ./dump.sql', shell=True)] # noqa
59+
[call('PGPASSWORD=my-cool-password pg_dump --format plain --dbname db --username root --host localhost --port 5432 --file ./dump.sql', shell=True)] # noqa
6060
],
6161
6262
['--list-providers --parallel',
6363
Namespace(verbose=None, list_providers=True, schema='schema.yml', dbname=None, user=None,
64-
password='', host='localhost', port='5432', dry_run=False, dump_file=None, init_sql=False, parallel=True), # noqa
64+
password='', host='localhost', port='5432', dry_run=False, dump_file=None, dump_options='--format custom --compress 9', init_sql=False, parallel=True), # noqa
6565
[], 0, []
66-
]
66+
],
6767
])
6868
def test_cli_args(self, subprocess, patched_connect, quote_ident, cli_args, expected, expected_executes, commit_calls, call_dump): # noqa
6969
arg_parser = get_arg_parser()

‎tests/test_utils.py‎

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -236,22 +236,33 @@ class TestCreateDatabaseDump(object):
236236

237237
@patch('pganonymize.utils.subprocess.call')
238238
def test(self, mock_call):
239-
create_database_dump(
240-
'/tmp/dump.gz',
241-
{'dbname': 'database', 'user': 'foo', 'host': 'localhost', 'port': 5432},
242-
)
239+
filename = '/tmp/dump.gz'
240+
db_args = {'dbname': 'database', 'user': 'foo', 'host': 'localhost', 'port': 5432}
241+
dump_args = '--format custom --compress 9'
242+
create_database_dump(filename, db_args, dump_args)
243243
mock_call.assert_called_once_with(
244244
'pg_dump --format custom --compress 9 --dbname database --username foo --host localhost --port 5432 --file /tmp/dump.gz',
245245
shell=True,
246246
)
247247

248248
@patch('pganonymize.utils.subprocess.call')
249249
def test_with_password(self, mock_call):
250-
create_database_dump(
251-
'/tmp/dump.gz',
252-
{'dbname': 'database', 'user': 'foo', 'host': 'localhost', 'port': 5432, 'password': 'pass'},
253-
)
250+
filename = '/tmp/dump.gz'
251+
db_args = {'dbname': 'database', 'user': 'foo', 'host': 'localhost', 'port': 5432, 'password': 'pass'}
252+
dump_args = '--format custom --compress 9'
253+
create_database_dump(filename, db_args, dump_args)
254254
mock_call.assert_called_once_with(
255255
'PGPASSWORD=pass pg_dump --format custom --compress 9 --dbname database --username foo --host localhost --port 5432 --file /tmp/dump.gz',
256256
shell=True,
257257
)
258+
259+
@patch('pganonymize.utils.subprocess.call')
260+
def test_with_custom_dump_args(self, mock_call):
261+
filename = '/tmp/dump.gz'
262+
db_args = {'dbname': 'database', 'user': 'foo', 'host': 'localhost', 'port': 5432}
263+
dump_args = '--format plain'
264+
create_database_dump(filename, db_args, dump_args)
265+
mock_call.assert_called_once_with(
266+
'pg_dump --format plain --dbname database --username foo --host localhost --port 5432 --file /tmp/dump.gz',
267+
shell=True,
268+
)

0 commit comments

Comments
 (0)