Add optional transaction parameter

aanari · Oct 18, 2017 · 3f88d94 · 3f88d94
1 parent 4b7421e
commit 3f88d94
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 11 deletions.
diff --git a/pg_materialize/pg_materialize.py b/pg_materialize/pg_materialize.py
@@ -79,12 +79,12 @@ def format_content(entity):
         .replace('COMMIT;', '')
     return "  -- %s%s" % (file_name, content_no_txn)
 
-def generate_script(views, spacer=''):
-    return "\n\n".join([
-        'BEGIN;',
-        spacer.join(views),
-        'COMMIT;'
-    ])
+def generate_script(views, transaction, spacer=''):
+    lines = [ spacer.join(views) ]
+    if transaction:
+        lines.insert(0, 'BEGIN;')
+        lines.append('COMMIT;')
+    return "\n\n".join(lines)
 
 def serialize_script(name, suffix, content, output_dir, verbose):
     script_name = '%s-%s.sql' % (name, suffix)
@@ -102,7 +102,8 @@ def serialize_script(name, suffix, content, output_dir, verbose):
 @click.option('-o', '--output-dir', default='.', help='Output directory for the generated creation and refresh scripts')
 @click.option('-p', '--pattern', default='_m?v', help='View regex pattern to match (i.e. "_m?v" for both "users_mv_hist" and "users_v")')
 @click.option('-v', '--verbose', is_flag=True, help='Enables verbose logging')
-def cli(dry_run, input_dir, ignore_refresh, output_dir, pattern, verbose):
+@click.option('-t', '--transaction', is_flag=True, help='Enables transaction wrapping around refresh')
+def cli(dry_run, input_dir, ignore_refresh, output_dir, pattern, verbose, transaction):
 
     dag = {}
 
@@ -153,18 +154,21 @@ def cli(dry_run, input_dir, ignore_refresh, output_dir, pattern, verbose):
         list
     )
 
-    create_script = generate_script(create_views)
+    create_script = generate_script(create_views, transaction)
 
+    refresh_prefix = 'REFRESH MATERIALIZED VIEW CONCURRENTLY '
+    if transaction:
+        refresh_prefix = '  ' + refresh_prefix
     refresh_views = pipe(sorted_views,
         filter(lambda view: re.search(pattern, view) and not (ignore_refresh and re.search(ignore_refresh , view))),
-        map(lambda view: '  REFRESH MATERIALIZED VIEW CONCURRENTLY ' + view + ';'),
+        map(lambda view: refresh_prefix + view + ';'),
         list
     )
 
     if verbose:
         print('Selecting %d Materialized Views for Refresh' % len(refresh_views))
 
-    refresh_script = generate_script(refresh_views, "\n\n")
+    refresh_script = generate_script(refresh_views, transaction, "\n\n")
 
     if dry_run:
         print('Dry Run Option Enabled - Skipping Script Generation')

diff --git a/tests/test_pg_materialize.py b/tests/test_pg_materialize.py
@@ -95,7 +95,7 @@ def test_format_content(self):
     def test_generate_script(self):
 
         views = ['  DEFINE_VIEW_1;', '  DEFINE_VIEW_2;', '  DEFINE_VIEW_3;']
-        actual = pgm.generate_script(views, "\n")
+        actual = pgm.generate_script(views, True, "\n")
         expected = """
 BEGIN;
 
@@ -108,6 +108,17 @@ def test_generate_script(self):
 
         self.assertEqual(actual.strip(), expected.strip())
 
+        actual = pgm.generate_script(views, False, "\n")
+        expected = """
+
+  DEFINE_VIEW_1;
+  DEFINE_VIEW_2;
+  DEFINE_VIEW_3;
+
+        """
+
+        self.assertEqual(actual.strip(), expected.strip())
+
 
     def test_process_file(self):