Merge pull request #24 from andersinno/skip_rows

Add option to leave out table contents completely

Author: suutari-ai

README.md

@@ -87,6 +87,7 @@ strategy:
     first_name: name.first_name
     last_name: name.last_name
     secret_key: string.empty
+  access_log: skip_rows
 ```
 
 In the example configuration above, there are first listed two "addon
@@ -108,3 +109,9 @@ sanitation function consists from two parts separated from each other by
 a dot: Python module name and name of the actual function, which will
 be prefixed with `sanitize_`, so `name.first_name` would be a function
 called `sanitize_first_name` in a file called `name.py`.
+
+Table content can be left out completely from the sanitized dump by
+setting table strategy to `skip_rows` (check `access_log` table in the
+example config). This will leave out all `INSERT INTO` (MySQL) or `COPY`
+(PostgreSQL) statements from the sanitized dump file. `CREATE TABLE`
+statements will not be removed.

database_sanitizer/config.py

@@ -9,6 +9,7 @@
 
 __all__ = ("Configuration", "ConfigurationError")
 
+SKIP_ROWS_CONFIG_VALUE = "skip_rows"
 MYSQLDUMP_DEFAULT_PARAMETERS = ["--single-transaction"]
 PG_DUMP_DEFAULT_PARAMETERS = []
 
@@ -26,6 +27,7 @@ class Configuration(object):
     """
     def __init__(self):
         self.sanitizers = {}
+        self.skip_rows_for_tables = []
         self.addon_packages = []
         self.mysqldump_params = []
         self.pg_dump_params = []
@@ -166,13 +168,18 @@ def load_sanitizers(self, config_data):
         if not isinstance(section_strategy, dict):
             if section_strategy is None:
                 return
-            raise ConfigurationError(
-                "'strategy' is %s instead of dict" % (
-                    type(section_strategy),
-                ),
-            )
+            if section_strategy != SKIP_ROWS_CONFIG_VALUE:
+                raise ConfigurationError(
+                    "'strategy' is %s instead of dict" % (
+                        type(section_strategy),
+                    ),
+                )
 
         for table_name, column_data in six.iteritems(section_strategy):
+            if column_data == SKIP_ROWS_CONFIG_VALUE:
+                self.skip_rows_for_tables.append(table_name)
+                continue
+
             if not isinstance(column_data, dict):
                 if column_data is None:
                     continue

database_sanitizer/dump/mysql.py

@@ -103,6 +103,11 @@ def sanitize_from_stream(stream, config):
         table_name = insert_into_match.group("table")
         column_names = parse_column_names(insert_into_match.group("columns"))
 
+        # Skip `INSERT INTO` statement if table rows are configured
+        # to be skipped.
+        if table_name in config.skip_rows_for_tables:
+            continue
+
         # Collect sanitizers possibly used for this table and place them into
         # a dictionary from which we can look them up by index later.
         sanitizers = {}

database_sanitizer/dump/postgres.py

@@ -54,6 +54,7 @@ def sanitize(url, config):
     sanitize_value_line = None
     current_table = None
     current_table_columns = None
+    skip_table = False
 
     for line in codecs.getreader("utf-8")(process.stdout):
         # Eat the trailing new line.
@@ -65,7 +66,12 @@ def sanitize(url, config):
             if line == "\\.":
                 current_table = None
                 current_table_columns = None
-                yield "\\."
+                if not skip_table:
+                    yield "\\."
+                skip_table = False
+                continue
+
+            if skip_table:
                 continue
 
             if not sanitize_value_line:
@@ -84,6 +90,12 @@ def sanitize(url, config):
         current_table = copy_line_match.group("table")
         current_table_columns = parse_column_names(copy_line_match.group("columns"))
 
+        # Skip `COPY` statement if table rows are configured
+        # to be skipped.
+        if config and current_table in config.skip_rows_for_tables:
+            skip_table = True
+            continue
+
         sanitize_value_line = get_value_line_sanitizer(
             config, current_table, current_table_columns)
 

database_sanitizer/tests/test_config.py

@@ -137,6 +137,29 @@ def mock_find_sanitizer(*args):
     assert "table2.column1" in config.sanitizers
 
 
+def test_table_skip_rows_configuration():
+    config = Configuration()
+
+    with pytest.raises(ConfigurationError):
+        config.load_sanitizers({"strategy": "test"})
+
+    def mock_find_sanitizer(*args):
+        return lambda value: value
+
+    with mock.patch("database_sanitizer.config.Configuration.find_sanitizer",
+                    side_effect=mock_find_sanitizer):
+
+        config.load_sanitizers({"strategy": {
+            "table1": "skip_rows",
+            "table2": {
+                "column1": "test",
+            }
+        }})
+
+    assert "table2.column1" in config.sanitizers
+    assert "table1" in config.skip_rows_for_tables
+
+
 def test_find_sanitizer():
     config = Configuration()
 

database_sanitizer/tests/test_dump_mysql.py

@@ -20,6 +20,13 @@
 
 DROP TABLE IF EXISTS `test`;
 
+CREATE TABLE `test` (
+`id` int(11) NOT NULL AUTO_INCREMENT,
+`created_at` date NOT NULL,
+`notes` varchar(255) NOT NULL,
+PRIMARY KEY (`id`)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
+
 INSERT INTO `test` (`id`, `created_at`, `notes`) VALUES \
 (1,'2018-01-01','Test data 1'),\
 (2,'2018-01-02','Test data 2'),\
@@ -61,6 +68,31 @@ def test_sanitize_from_stream():
 """ in dump_output_lines
 
 
+def test_skip_table_rows():
+    stream = io.BytesIO(MOCK_MYSQLDUMP_OUTPUT)
+    config = Configuration()
+    config.skip_rows_for_tables.append('test')
+
+    output = list(sanitize_from_stream(stream, config))
+
+    assert output == [
+        '',
+        '--- Fake MySQL database dump',
+        '',
+        'DROP TABLE IF EXISTS `test`;',
+        '',
+        'CREATE TABLE `test` (',
+        '`id` int(11) NOT NULL AUTO_INCREMENT,',
+        '`created_at` date NOT NULL,',
+        '`notes` varchar(255) NOT NULL,',
+        'PRIMARY KEY (`id`)',
+        ') ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;',
+        '',
+        '',
+        '--- Final line after `INSERT INTO` statement.',
+    ]
+
+
 def test_sanitizer_invalid_input():
     stream = io.BytesIO(INVALID_MOCK_MYSQLDUMP_OUTPUT)
     config = Configuration()

database_sanitizer/tests/test_dump_postgres.py

@@ -19,6 +19,12 @@
 
 COMMENT ON SCHEMA "public" IS 'standard public schema';
 
+CREATE TABLE "public"."test" (
+"id" integer NOT NULL,
+"created_at" timestamp with time zone NOT NULL,
+"notes" character varying(255) NOT NULL
+);
+
 COPY "public"."test" ("id", "created_at", "notes") FROM stdin;
 1\t2018-01-01 00:00:00\tTest data 1
 2\t2018-01-02 00:00:00\tTest data 2
@@ -65,6 +71,31 @@ def test_sanitize():
     assert "2\t2018-01-02 00:00:00\tSanitized" in dump_output_lines
 
 
+def test_skip_table_rows():
+    url = urlparse.urlparse("postgres://localhost/test")
+    config = Configuration()
+    config.skip_rows_for_tables.append('test')
+
+    with mock.patch("subprocess.Popen",
+                    side_effect=create_mock_popen(MOCK_PG_DUMP_OUTPUT)):
+        output = list(sanitize(url, config))
+
+    assert output == [
+        '--- Fake PostgreSQL database dump',
+        '',
+        'COMMENT ON SCHEMA "public" IS \'standard public schema\';',
+        '',
+        'CREATE TABLE "public"."test" (',
+        '"id" integer NOT NULL,',
+        '"created_at" timestamp with time zone NOT NULL,',
+        '"notes" character varying(255) NOT NULL',
+        ');',
+        '',
+        '',
+        '--- Final line after `COPY` statement'
+    ]
+
+
 def test_sanitizer_invalid_input():
     url = urlparse.urlparse("postgres://localhost/test")