Omit Postgres dump COPY statements if table is configured to skip rows

Leave out any `COPY` statements for table if that table is listed in
Configuration.skip_rows_for_table.

Author: ehaivala

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_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")