Added environment variable control

jazzband · May 16, 2016 · e30755d · e30755d
1 parent 8f47ad1
commit e30755d
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 9 deletions.
diff --git a/dbbackup/db/base.py b/dbbackup/db/base.py
@@ -105,23 +105,31 @@ 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.
 
         :param command: Command line to launch
         :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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
 ------