Merge pull request #23 from andersinno/add-extra-params-support

JoaRiski · web-flow · commit 861976f41154 · 2019-06-07T13:52:16.000+03:00
Add support for dump CLI extra parameters
diff --git a/README.md b/README.md
@@ -77,6 +77,11 @@ config:
   addons:
     - some.other.package
     - yet.another.package
+  extra_parameters: # These parameters will be passed to the dump tool CLI
+    mysqldump:
+      - "--single-transaction" # Included by default
+    pg_dump:
+      - "--exclude-table=something"
 strategy:
   user:
     first_name: name.first_name
@@ -90,6 +95,11 @@ be looking for sanitizer functions. They are completely optional and can
 be omitted, in which case only sanitizer functions defined in package
 called `sanitizers` and built-in sanitizers will be used instead.
 
+It's also possible to define extra parameters to pass to the dump tool (
+`mysqldump` or `pg_dump`). By default, `mysqldump` will include the
+`--single-transaction` extra parameter. You can disable this by defining the
+extra parameters in the config file explicitly, e.g. with an empty array `[]`.
+
 The `strategy` portion of the configuration contains the actual
 sanitation rules. First you define name of the database table (in the
 example that would be `user`) followed by column names in that table
diff --git a/database_sanitizer/config.py b/database_sanitizer/config.py
@@ -9,6 +9,9 @@
 
 __all__ = ("Configuration", "ConfigurationError")
 
+MYSQLDUMP_DEFAULT_PARAMETERS = ["--single-transaction"]
+PG_DUMP_DEFAULT_PARAMETERS = []
+
 
 class ConfigurationError(ValueError):
     """
@@ -24,6 +27,8 @@ class Configuration(object):
     def __init__(self):
         self.sanitizers = {}
         self.addon_packages = []
+        self.mysqldump_params = []
+        self.pg_dump_params = []
 
     @classmethod
     def from_file(cls, filename):
@@ -64,6 +69,51 @@ def load(self, config_data):
 
         self.load_addon_packages(config_data)
         self.load_sanitizers(config_data)
+        self.load_dump_extra_parameters(config_data)
+
+    def load_dump_extra_parameters(self, config_data):
+        """
+        Loads extra parameters for mysqldump and/or pg_dump CLI usage. These
+        parameters should be added to the mysqldump and/or pg_dump command call
+        when taking a dump.
+
+        :param config_data: Already parsed configuration data, as dictionary.
+        :type config_data: dict[str,any]
+        """
+        section_config = config_data.get("config", {})
+        if not isinstance(section_config, dict):
+            raise ConfigurationError(
+                "'config' is %s instead of dict" % (
+                    type(section_config),
+                ),
+            )
+
+        section_extra_parameters = section_config.get("extra_parameters", {})
+        if not isinstance(section_extra_parameters, dict):
+            raise ConfigurationError(
+                "'config.extra_parameters' is %s instead of dict" % (
+                    type(section_extra_parameters),
+                ),
+            )
+
+        mysqldump_params = section_extra_parameters.get("mysqldump", MYSQLDUMP_DEFAULT_PARAMETERS)
+        if not isinstance(mysqldump_params, list):
+            raise ConfigurationError(
+                "'config.extra_parameters.mysqldump' is %s instead of list" % (
+                    type(mysqldump_params),
+                ),
+            )
+
+        pg_dump_params = section_extra_parameters.get("pg_dump", PG_DUMP_DEFAULT_PARAMETERS)
+        if not isinstance(pg_dump_params, list):
+            raise ConfigurationError(
+                "'config.extra_parameters.pg_dump' is %s instead of list" % (
+                    type(pg_dump_params),
+                ),
+            )
+
+        self.mysqldump_params = mysqldump_params
+        self.pg_dump_params = pg_dump_params
 
     def load_addon_packages(self, config_data):
         """
diff --git a/database_sanitizer/dump/mysql.py b/database_sanitizer/dump/mysql.py
@@ -11,6 +11,7 @@
     encode_mysql_literal,
     get_mysqldump_args_and_env_from_url,
 )
+from ..config import MYSQLDUMP_DEFAULT_PARAMETERS
 
 #: Regular expression which matches `INSERT INTO` statements produced by the
 #: `mysqldump` utility, even when extended inserts have been enabled.
@@ -57,8 +58,12 @@ def sanitize(url, config):
 
     args, env = get_mysqldump_args_and_env_from_url(url=url)
 
+    extra_params = MYSQLDUMP_DEFAULT_PARAMETERS
+    if config:
+        extra_params = config.mysqldump_params
+
     process = subprocess.Popen(
-        args=["mysqldump"] + args,
+        args=["mysqldump"] + args + extra_params,
         env=env,
         stdout=subprocess.PIPE,
     )
diff --git a/database_sanitizer/dump/postgres.py b/database_sanitizer/dump/postgres.py
@@ -7,6 +7,7 @@
 import subprocess
 
 from ..utils.postgres import decode_copy_value, encode_copy_value
+from ..config import PG_DUMP_DEFAULT_PARAMETERS
 
 COPY_LINE_PATTERN = re.compile(
     r"^COPY \"(?P<schema>[^\"]*)\".\"(?P<table>[^\"]*)\" "
@@ -31,6 +32,10 @@ def sanitize(url, config):
     if url.scheme not in ("postgres", "postgresql", "postgis"):
         raise ValueError("Unsupported database type: '%s'" % (url.scheme,))
 
+    extra_params = PG_DUMP_DEFAULT_PARAMETERS
+    if config:
+        extra_params = config.pg_dump_params
+
     process = subprocess.Popen(
         (
             "pg_dump",
@@ -42,7 +47,7 @@ def sanitize(url, config):
             # URL as argument to the command.
             "--dbname",
             url.geturl().replace('postgis://', 'postgresql://'),
-        ),
+         ) + tuple(extra_params),
         stdout=subprocess.PIPE,
     )
 
diff --git a/database_sanitizer/tests/test_config.py b/database_sanitizer/tests/test_config.py
@@ -28,6 +28,54 @@ def test_load_config_data_must_be_dict():
         config.load(config_data="test")
 
 
+def test_load_dump_extra_parameters():
+    config = Configuration()
+
+    config.load_dump_extra_parameters({})
+    assert config.mysqldump_params == ["--single-transaction"]
+    assert config.pg_dump_params == []
+
+    with pytest.raises(ConfigurationError):
+        config.load_dump_extra_parameters({"config": "test"})
+
+    config.load_dump_extra_parameters({"config": {}})
+    assert config.mysqldump_params == ["--single-transaction"]
+    assert config.pg_dump_params == []
+
+    with pytest.raises(ConfigurationError):
+        config.load_dump_extra_parameters({"config": {
+            "extra_parameters": "test"
+        }})
+
+    with pytest.raises(ConfigurationError):
+        config.load_dump_extra_parameters({"config": {
+            "extra_parameters": [True]
+        }})
+
+    with pytest.raises(ConfigurationError):
+        config.load_dump_extra_parameters({"config": {
+            "extra_parameters": {
+                "mysqldump": "hernekeitto",
+            },
+        }})
+
+    with pytest.raises(ConfigurationError):
+        config.load_dump_extra_parameters({"config": {
+            "extra_parameters": {
+                "pg_dump": "viina",
+            },
+        }})
+
+    config.load_dump_extra_parameters({"config": {
+        "extra_parameters": {
+            "mysqldump": ["--double-transaction"],
+            "pg_dump": ["--exclude-table=something"],
+        },
+    }})
+    assert config.mysqldump_params == ["--double-transaction"]
+    assert config.pg_dump_params == ["--exclude-table=something"]
+
+
 def test_load_addon_packages():
     config = Configuration()
 
diff --git a/database_sanitizer/tests/test_dump.py b/database_sanitizer/tests/test_dump.py
@@ -5,12 +5,15 @@
 import pytest
 
 from database_sanitizer import dump
+from database_sanitizer.config import Configuration
 
 EXPECTED_POPEN_KWARGS = {
     'mysql://User:Pass@HostName/Db': {
         'args': (
             'mysqldump --complete-insert --extended-insert'
-            ' --net_buffer_length=10240 -h hostname -u User Db').split(),
+            ' --net_buffer_length=10240 -h hostname -u User Db'
+            ' --single-transaction'
+        ).split(),
         'env': {'MYSQL_PWD': 'Pass'},
         'stdout': subprocess.PIPE,
     },
@@ -45,6 +48,73 @@ def test_run(mocked_popen, url):
     assert popen_kwargs == expected_popen_kwargs
 
 
+@mock.patch('subprocess.Popen')
+def test_run_with_mysql_extra_params(mocked_popen):
+    mocked_popen.return_value.stdout = BytesIO(b'INPUT DUMP')
+    output = StringIO()
+
+    url = "mysql://User:Pass@HostName/Db"
+    config = Configuration()
+    config.load({
+        "config": {
+            "extra_parameters": {
+                "mysqldump": ["--double-transaction"]
+            }
+        }
+    })
+
+    dump.run(url, output, config)
+
+    expected = {
+        'args': (
+            'mysqldump --complete-insert --extended-insert'
+            ' --net_buffer_length=10240 -h hostname -u User Db'
+            ' --double-transaction'
+        ).split(),
+        'env': {'MYSQL_PWD': 'Pass'},
+        'stdout': subprocess.PIPE,
+    }
+
+    (popen_args, popen_kwargs) = mocked_popen.call_args
+    expected_popen_args = (
+        (expected.pop('args'),) if popen_args else ())
+    assert popen_args == expected_popen_args
+    assert popen_kwargs == expected
+
+
+@mock.patch('subprocess.Popen')
+def test_run_with_pg_dump_extra_params(mocked_popen):
+    mocked_popen.return_value.stdout = BytesIO(b'INPUT DUMP')
+    output = StringIO()
+
+    url = "postgres:///Db"
+    config = Configuration()
+    config.load({
+        "config": {
+            "extra_parameters": {
+                "pg_dump": ["--exclude-table=something"]
+            }
+        }
+    })
+
+    dump.run(url, output, config)
+
+    expected = {
+        'args': tuple((
+            'pg_dump --encoding=utf-8 --quote-all-identifiers'
+            ' --dbname postgres:///Db'
+            ' --exclude-table=something'
+        ).split()),
+        'stdout': subprocess.PIPE,
+    }
+
+    (popen_args, popen_kwargs) = mocked_popen.call_args
+    expected_popen_args = (
+        (expected.pop('args'),) if popen_args else ())
+    assert popen_args == expected_popen_args
+    assert popen_kwargs == expected
+
+
 @mock.patch('subprocess.Popen')
 def test_run_unknown_scheme(mocked_popen):
     with pytest.raises(ValueError) as excinfo:
