From e30755d0f9b9e969ef20109bd080f59c829b51b2 Mon Sep 17 00:00:00 2001 From: ZuluPro Date: Mon, 16 May 2016 05:20:39 -0400 Subject: [PATCH] Added environment variable control --- dbbackup/db/base.py | 14 +++++++++++--- dbbackup/db/mongodb.py | 4 ++-- dbbackup/db/mysql.py | 4 ++-- dbbackup/db/postgresql.py | 4 ++-- dbbackup/tests/test_connectors.py | 13 +++++++++++++ docs/databases.rst | 7 +++++++ 6 files changed, 37 insertions(+), 9 deletions(-) diff --git a/dbbackup/db/base.py b/dbbackup/db/base.py index cecb4efb..8eb21495 100644 --- a/dbbackup/db/base.py +++ b/dbbackup/db/base.py @@ -105,8 +105,11 @@ class BaseCommandDBConnector(BaseDBConnector): dump_suffix = '' restore_prefix = '' restore_suffix = '' + env = {} + dump_env = {} + restore_env = {} - def run_command(self, command, stdin=None): + def run_command(self, command, stdin=None, env=None): """ Launch a shell command line. @@ -114,14 +117,19 @@ def run_command(self, command, stdin=None): :type command: str :param stdin: Standard input of command :type stdin: file + :param env: Environment variable used in command + :type env: dict :return: Standard output of command :rtype: file """ + cmd = shlex.split(command) stdout = SpooledTemporaryFile(max_size=10 * 1024 * 1024) stderr = SpooledTemporaryFile(max_size=10 * 1024 * 1024) - cmd = shlex.split(command) + full_env = self.env.copy() + full_env.update(env or {}) try: - process = Popen(cmd, stdin=stdin, stdout=stdout, stderr=stderr) + process = Popen(cmd, stdin=stdin, stdout=stdout, stderr=stderr, + env=full_env) process.wait() if process.poll(): stderr.seek(0) diff --git a/dbbackup/db/mongodb.py b/dbbackup/db/mongodb.py index 270f51ed..0c636585 100644 --- a/dbbackup/db/mongodb.py +++ b/dbbackup/db/mongodb.py @@ -22,7 +22,7 @@ def _create_dump(self): cmd += ' --excludeCollection {}'.format(collection) cmd += ' --archive' cmd = '{} {} {}'.format(self.dump_prefix, cmd, self.dump_suffix) - stdout, stderr = self.run_command(cmd) + stdout, stderr = self.run_command(cmd, env=self.dump_env) return stdout def _restore_dump(self, dump): @@ -38,4 +38,4 @@ def _restore_dump(self, dump): cmd += ' --drop' cmd += ' --archive' cmd = '{} {} {}'.format(self.restore_prefix, cmd, self.restore_suffix) - return self.run_command(cmd, stdin=dump) + return self.run_command(cmd, stdin=dump, env=self.restore_env) diff --git a/dbbackup/db/mysql.py b/dbbackup/db/mysql.py index 428560bb..6e4000cc 100644 --- a/dbbackup/db/mysql.py +++ b/dbbackup/db/mysql.py @@ -22,7 +22,7 @@ def _create_dump(self): for table in self.exclude: cmd += ' --ignore-table={}.{}'.format(self.settings['NAME'], table) cmd = '{} {} {}'.format(self.dump_prefix, cmd, self.dump_suffix) - stdout, stderr = self.run_command(cmd) + stdout, stderr = self.run_command(cmd, env=self.dump_env) return stdout def _restore_dump(self, dump): @@ -36,5 +36,5 @@ def _restore_dump(self, dump): if self.settings.get('PASSWORD'): cmd += ' --password={}'.format(self.settings['PASSWORD']) cmd = '{} {} {}'.format(self.restore_prefix, cmd, self.restore_suffix) - stdout, stderr = self.run_command(cmd, stdin=dump) + stdout, stderr = self.run_command(cmd, stdin=dump, env=self.restore_env) return stdout, stderr diff --git a/dbbackup/db/postgresql.py b/dbbackup/db/postgresql.py index 9a2e8b25..85a03867 100644 --- a/dbbackup/db/postgresql.py +++ b/dbbackup/db/postgresql.py @@ -29,7 +29,7 @@ def _create_dump(self): if self.drop: cmd += ' --clean' cmd = '{} {} {}'.format(self.dump_prefix, cmd, self.dump_suffix) - stdout, stderr = self.run_command(cmd) + stdout, stderr = self.run_command(cmd, env=self.dump_env) return stdout def _restore_dump(self, dump): @@ -47,7 +47,7 @@ def _restore_dump(self, dump): if self.single_transaction: cmd += ' --single-transaction' cmd = '{} {} {}'.format(self.restore_prefix, cmd, self.restore_suffix) - stdout, stderr = self.run_command(cmd, stdin=dump) + stdout, stderr = self.run_command(cmd, stdin=dump, env=self.restore_env) return stdout, stderr diff --git a/dbbackup/tests/test_connectors.py b/dbbackup/tests/test_connectors.py index c1c7ce2c..54275867 100644 --- a/dbbackup/tests/test_connectors.py +++ b/dbbackup/tests/test_connectors.py @@ -53,6 +53,19 @@ def test_run_command_stdin(self): self.assertEqual(stdout.read(), b'foo') self.assertFalse(stderr.read()) + def test_run_command_with_env(self): + connector = BaseCommandDBConnector() + # Empty env + stdout, stderr = connector.run_command('env') + self.assertFalse(stdout.read()) + # env from self.env + connector.env = {'foo': 'bar'} + stdout, stderr = connector.run_command('env') + self.assertEqual(stdout.read(), b'foo=bar\n') + # method overide gloabal env + stdout, stderr = connector.run_command('env', env={'foo': 'ham'}) + self.assertEqual(stdout.read(), b'foo=ham\n') + class SqliteConnectorTest(TestCase): def test_write_dump(self): diff --git a/docs/databases.rst b/docs/databases.rst index 3eb0b8d2..20d9dac1 100644 --- a/docs/databases.rst +++ b/docs/databases.rst @@ -89,6 +89,13 @@ DUMP_SUFFIX and RESTORE_PREFIX String to include as suffix of dump or restore command. It will be add with a space between launched command and its suffix. +ENV, DUMP_ENV and RESTORE_ENV +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Environment variables used during command running, default are ``{}``. ``ENV`` +is used for every command, ``DUMP_ENV`` and ``RESTORE_ENV`` override the +values defined in ``ENV`` during the dedicated commands. + SQLite ------